import networkx as nx
import pygraphviz as pgv
from loguru import logger
from .asset import Asset, AssetStatus
#import matplotlib.pyplot as plt
from collections import Counter, defaultdict
from collections.abc import Iterable
import pendulum as plm
from notifiers.logging import NotificationHandler
from kombu import Connection, Exchange, Queue
from .schedule import Schedule
import asyncio
from enum import Enum
from queue import Queue as ThreadSafeQueue, Empty
from threading import Thread
from coolname import generate_slug
from .message import publish_msg, listen
from .registry import register_graph_process, unregister_graph_process, list_running_graphs
import signal
import inspect
import pickle
import base64
class ShutdownSignal(Exception):
pass
class BuildResult(Enum):
Rebuilt = 0
Skipped = 1
Failed = 2
__RABBIT_URL__ = "amqp://guest:guest@localhost//"
__MESSAGE_TTL__ = 1
_NODE_COLORS = {AssetStatus.Current : "green",
AssetStatus.Stale : "khaki",
AssetStatus.Building : "blue",
AssetStatus.Paused : "slategray",
AssetStatus.UpstreamStopped : "lightslategray",
AssetStatus.Unavailable : "gray",
AssetStatus.Failed : "red"
}
# TODO: use HEX to match above colors <fg #00005f>, <fg #EE1>
_LOGURU_COLORS = {AssetStatus.Current : "32",
AssetStatus.Stale : "33",
AssetStatus.Building : "34",
AssetStatus.Paused : "35",
AssetStatus.UpstreamStopped : "35",
AssetStatus.Unavailable : "37",
AssetStatus.Failed : "91"
}
# TODO: use better formatting than ANSI...
def fmt(status):
#return f"<{_LOGURU_COLORS[status]}>{status}</>"
return f"\033[{_LOGURU_COLORS[status]}m{status}\033[00m\u001b[1m"
async def wait_for_dependencies(task, parents):
for parent in parents:
await parent
result = await task
return result
[docs]
class AssetGraph:
"""A directed acyclic graph of :class:`~fwirl.Asset` objects with a
built-in scheduler and messaging interface.
:class:`AssetGraph` is the central object in fwirl. It holds a DAG of
assets, tracks their statuses, and orchestrates builds and refreshes both
on demand and on a cron-based schedule. It can also receive commands at
runtime through RabbitMQ by calling :meth:`run`.
Args:
key: A unique string identifier for this graph. The same key is
used as the RabbitMQ routing key when running in server mode.
notifiers: Optional mapping of `notifiers
<https://notifiers.readthedocs.io/>`_ service configurations used
to send build notifications. Each entry has the form::
{
"<service_name>": {
"params": {...}, # notifiers service params
"level": "ERROR", # loguru level string
}
}
Example::
import fwirl
g = fwirl.AssetGraph("my_graph")
g.add_assets([asset_a, asset_b])
g.schedule("hourly_build", "build", "0 * * * *")
g.run() # blocks; listens for RabbitMQ commands
"""
def __init__(self, key, notifiers=None):
self.graph = nx.DiGraph()
self.key = key
self.schedules = {}
if notifiers is not None:
for service in notifiers:
handler = NotificationHandler(service, defaults=notifiers[service]["params"])
logger.add(handler, level=notifiers[service]["level"])
self.message_queue = ThreadSafeQueue()
self.job_queue = []
self.job_running = None
self.workers = []
[docs]
def add_assets(self, assets):
"""Add a list of assets (and all their transitive dependencies) to the graph.
The method walks the dependency tree rooted at each asset in *assets*
and inserts every discovered node and edge into the internal DiGraph.
Duplicate assets (same key) are silently skipped.
Args:
assets: An iterable of :class:`~fwirl.Asset` objects to add.
Raises:
ValueError: If a self-dependency loop is detected, or if any
asset's :meth:`~fwirl.Asset.build` or
:meth:`~fwirl.Asset.timestamp` method is not a coroutine.
"""
logger.info(f"Gathering edges, assets, and upstream assets to add to the graph")
edges = []
unq_assets = set()
asset_queue = []
asset_queue.extend(assets)
while len(asset_queue) != 0:
a = asset_queue.pop()
if a in a.dependencies:
# TODO test this check...
raise ValueError("Self-dependency loop detected!")
if a in unq_assets or a in self.graph:
continue
unq_assets.add(a)
edges.extend([(d, a) for d in a.dependencies])
asset_queue.extend(a.dependencies)
logger.info(f"Adding {len(edges)} edges, {len(assets)} assets, and upstream assets to the graph")
old_edges = self.graph.size()
old_nodes = self.graph.number_of_nodes()
self.graph.add_nodes_from(unq_assets)
self.graph.add_edges_from(edges)
new_edges = self.graph.size()
new_nodes = self.graph.number_of_nodes()
#self._stale_topological_sort = True
logger.info(f"Added {new_edges-old_edges} unique edges and {new_nodes-old_nodes} unique assets to the graph")
logger.info(f"Verifying that all assets have async timestamp and build functions")
for asset in self.graph:
if not inspect.iscoroutinefunction(asset.build):
logger.error(f"Asset {asset}'s build routine is not async. Please ensure that .build() and .timestamp() are asynchronous.")
raise ValueError
if not inspect.iscoroutinefunction(asset.timestamp):
logger.error(f"Asset {asset}'s timestamp routine is not async. Please ensure that .build() and .timestamp() are asynchronous.")
raise ValueError
[docs]
def remove_assets(self, assets):
"""Remove a list of assets from the graph.
Downstream assets that depended on any removed node are also removed
automatically by NetworkX.
Args:
assets: An iterable of :class:`~fwirl.Asset` objects to remove.
"""
logger.info(f"Removing {len(assets)} assets and downstream assets from graph")
old_edges = self.graph.size()
old_nodes = self.graph.number_of_nodes()
self.graph.remove_nodes_from(assets)
new_edges = self.graph.size()
new_nodes = self.graph.number_of_nodes()
#self._stale_topological_sort = True
logger.info(f"Removed {old_nodes - new_nodes} assets and {old_edges - new_edges} edges from the graph")
[docs]
def list_assets(self, display=True):
"""Return (and optionally log) a formatted list of all assets.
Args:
display: When ``True`` (the default) the string is also emitted
via the loguru logger.
Returns:
str: A multi-line string listing each asset key and its current
:class:`~fwirl.AssetStatus`.
"""
s = ''
if self.graph.number_of_nodes() == 0:
s += "No assets to list."
else:
s += "Assets:\n"
for asset in self.graph:
s += f"- {asset.key}: {fmt(asset.status)}\n"
if display:
logger.info(s)
return s
[docs]
def list_schedules(self, display=True):
"""Return (and optionally log) a formatted list of all schedules.
Args:
display: When ``True`` (the default) the string is also emitted
via the loguru logger.
Returns:
str: A multi-line string listing each schedule name and its
:class:`~fwirl.schedule.Schedule` representation.
"""
s = ''
if len(self.schedules) == 0:
s += "No schedules to list."
else:
s += "Schedules:\n"
for sch in self.schedules:
s += f"- {sch}: {self.schedules[sch]}\n"
if display:
logger.info(s)
return s
[docs]
def list_jobs(self, display=True):
"""Return (and optionally log) the current job queue.
Args:
display: When ``True`` (the default) the string is also emitted
via the loguru logger.
Returns:
str: A multi-line string describing the currently running job
(if any) and the pending job queue.
"""
s = ''
if self.job_running is None:
s += 'No running job.\n'
else:
s += f'Running job: {self.job_running}\n'
if len(self.job_queue) == 0:
s += "No jobs to list."
else:
s += "Jobs:\n"
for j in self.job_queue:
s += f"- {j}\n"
if display:
logger.info(s)
return s
def _asset_by_key(self, key):
for asset in self.graph:
if asset.key == key:
return asset
return None
def _serialize_timestamp(self, ts):
if ts == AssetStatus.Unavailable or ts is None:
return None
if hasattr(ts, "to_iso8601_string"):
return ts.to_iso8601_string()
return str(ts)
def _serialize_property_value(self, value):
if value is None or isinstance(value, (bool, int, float, str)):
return value
if isinstance(value, dict):
return {str(k): self._serialize_property_value(v) for k, v in value.items()}
if isinstance(value, Iterable) and not isinstance(value, (str, bytes)):
return [self._serialize_property_value(v) for v in value]
return str(value)
async def _asset_payload(self, asset):
timestamp = await asset.timestamp()
return {
"key": asset.key,
"status": asset.status.name,
"message": asset.message,
"type": asset.__class__.__name__,
"group": asset.group,
"subgroup": asset.subgroup,
"allow_retry": asset.allow_retry,
"timestamp": self._serialize_timestamp(timestamp),
"last_build_timestamp": self._serialize_timestamp(asset._last_build_timestamp),
"parents": [p.key for p in self.graph.predecessors(asset)],
"children": [c.key for c in self.graph.successors(asset)],
"properties": self._serialize_property_value(asset.properties),
}
async def _graph_payload(self):
nodes = []
for asset in self.graph:
nodes.append(await self._asset_payload(asset))
edges = []
for parent, child in self.graph.edges():
edges.append({"from": parent.key, "to": child.key})
subgroup_buckets = defaultdict(list)
for asset in self.graph:
if asset.group is None and asset.subgroup is None:
continue
group = '__nogroup__' if asset.group is None else str(asset.group)
subgroup = '__nosubgroup__' if asset.subgroup is None else str(asset.subgroup)
subgroup_buckets[(group, subgroup)].append(asset)
collapse_candidates = []
for (group, subgroup), assets in subgroup_buckets.items():
if len(assets) < 2:
continue
statuses = {a.status for a in assets}
if len(statuses) == 1:
status = list(statuses)[0]
collapse_candidates.append({
"id": f"{group}::{subgroup}::{status.name}",
"group": None if group == '__nogroup__' else group,
"subgroup": None if subgroup == '__nosubgroup__' else subgroup,
"status": status.name,
"node_keys": [a.key for a in assets]
})
return {
"nodes": nodes,
"edges": edges,
"collapse_candidates": collapse_candidates,
"summary": {
"asset_count": self.graph.number_of_nodes(),
"edge_count": self.graph.number_of_edges()
}
}
[docs]
def update_asset_properties(self, asset_key, properties):
asset = self._asset_by_key(asset_key)
if asset is None:
raise ValueError(f"Asset '{asset_key}' not found")
if not isinstance(properties, dict):
raise ValueError("Properties must be a dictionary")
asset.properties = dict(properties)
return asset
[docs]
def schedule(self, schedule_key, action, cron_string='', asset=None, immediate_once=False):
"""Add a recurring (or one-shot) schedule to the graph.
Args:
schedule_key: A unique string identifier for the schedule.
action: ``"refresh"`` to trigger a status refresh, or
``"build"`` to trigger a full build.
cron_string: A standard cron expression (e.g.
``"0 6 * * *"`` for 06:00 every day). Ignored when
*immediate_once* is ``True``.
asset: Optional :class:`~fwirl.Asset` targeted by the action.
When ``None`` the action applies to all assets.
immediate_once: When ``True`` the schedule fires once as soon
as possible and is then automatically removed.
"""
logger.info(f"Adding new schedule with key {schedule_key}")
if schedule_key in self.schedules:
logger.error(f"Tried to add schedule key {schedule_key}, but it already exists. Skipping.")
return
# find asset corresponding to asset_key
jobstr = f"{action} all" if asset is None else f"{action} {asset}"
logger.info(f"Adding scheduled run '{schedule_key}' ({jobstr} at '{cron_string}') to graph {self.key}")
if action == "refresh":
func = self._refresh
kwargs = {"assets": asset}
elif action == "build":
func = self._build
kwargs = {"assets": asset}
else:
logger.error(f"Tried to add schedule key {schedule_key} with unrecognized action {action}. Skipping.")
return
self.schedules[schedule_key] = Schedule(schedule_key, func, kwargs, cron_string, immediate_once)
[docs]
def unschedule(self, schedule_key):
"""Remove a schedule from the graph.
Args:
schedule_key: The string key of the schedule to remove.
"""
logger.info(f"Removing scheduled run '{schedule_key}'")
if schedule_key not in self.schedules:
logger.warning(f"Schedule {schedule_key} not in schedules.")
self.schedules.pop(schedule_key, None)
[docs]
def pause_schedule(self, schedule_key):
"""Pause a schedule so that it no longer fires.
Args:
schedule_key: The string key of the schedule to pause.
"""
logger.info(f"Pausing schedule '{schedule_key}'")
self.schedules[schedule_key].pause()
[docs]
def pause_asset(self, asset):
"""Set an asset's status to :attr:`~fwirl.AssetStatus.Paused`.
A paused asset is skipped during builds and its downstream assets
are marked :attr:`~fwirl.AssetStatus.UpstreamStopped`.
Args:
asset: The :class:`~fwirl.Asset` to pause.
"""
logger.info(f"Pausing asset {asset}")
asset.status = AssetStatus.Paused
[docs]
def set_allow_retry(self, asset, allow_retry):
"""Set whether a failed asset is retried on the next build pass.
See the ``allow_retry`` argument of :class:`~fwirl.Asset`.
Args:
asset: The :class:`~fwirl.Asset` to update.
allow_retry: ``True`` to retry failures automatically, ``False``
to keep the asset ``Failed`` until it is manually unpaused.
"""
logger.info(f"Setting allow_retry={allow_retry} for asset {asset}")
asset.allow_retry = allow_retry
[docs]
def unpause_schedule(self, schedule_key):
"""Resume a previously paused schedule.
Args:
schedule_key: The string key of the schedule to resume.
"""
logger.info(f"Resuming schedule '{schedule_key}'")
self.schedules[schedule_key].unpause()
[docs]
def unpause_asset(self, asset):
"""Mark a paused asset as :attr:`~fwirl.AssetStatus.Stale` so it
will be rebuilt on the next build pass.
Args:
asset: The :class:`~fwirl.Asset` to unpause.
"""
logger.info(f"Unpausing asset {asset}")
asset.status = AssetStatus.Stale
def _get_message(self, timeout):
try:
msg = self.message_queue.get(timeout=timeout)
except Empty:
return None
return msg
# Consolidate status Current nodes to declutter but leave nodes with other statuses
# Returns a modified copy of the original graph
def _preprocess_graph(self):
graph = self.graph.copy()
root_nodes = [node for node, in_degree in graph.in_degree() if in_degree == 0]
# Stack for traversing graph
s = []
visited = []
for r in root_nodes:
s.append(r)
while len(s) > 0:
n = s.pop()
if n in visited:
continue
visited.append(n)
if n.status is not AssetStatus.Current:
continue
# Consolidate node
children = list(graph.successors(n))
parents = list(graph.predecessors(n))
for c in children:
if n not in root_nodes:
for p in parents:
graph.add_edge(p, c)
graph.remove_node(n)
s.append(c)
return graph
# Annotates and colors graph then outputs the pickled and base64 encoded SVG representation
def _generate_graph_viz(self, graph):
agraph = pgv.AGraph(directed=False, strict=True)
for node, attrs in graph.nodes(data=True):
status_color = _NODE_COLORS[node.status]
agraph.add_node(node, color=status_color, fontcolor=status_color)
for u, v, attrs in graph.edges(data=True):
agraph.add_edge(u, v)
agraph.node_attr["shape"] = "square"
svg = agraph.draw(format='svg', prog='dot')
pickled = pickle.dumps(svg)
encoded = base64.b64encode(pickled)
return encoded
async def _process_message(self, msg):
if msg is None: # if message get timed out
return
if msg["type"] == "graph":
preprocessed = self._preprocess_graph()
resp = self._generate_graph_viz(preprocessed)
resp_msg = {'type': 'response', 'response': resp}
publish_msg(msg["resp_queue"], resp_msg)
if msg["type"] == "graph_snapshot":
resp = await self._graph_payload()
publish_msg(msg["resp_queue"], {'type': 'response', 'response': resp})
if msg["type"] == "asset_detail":
asset = self._asset_by_key(msg["asset_key"])
if asset is None:
publish_msg(msg["resp_queue"], {'type': 'response', 'response': None})
else:
resp = await self._asset_payload(asset)
publish_msg(msg["resp_queue"], {'type': 'response', 'response': resp})
if msg["type"] == "update_asset_properties":
try:
self.update_asset_properties(msg["asset_key"], msg.get("properties", {}))
if "resp_queue" in msg:
publish_msg(msg["resp_queue"], {'type': 'response', 'response': {"ok": True}})
except Exception as e:
if "resp_queue" in msg:
publish_msg(msg["resp_queue"], {'type': 'response', 'response': {"ok": False, "error": str(e)}})
if msg["type"] == "summarize":
resp = self.summarize(display=False)
resp_msg = {'type': 'response', 'response': resp}
publish_msg(msg["resp_queue"], resp_msg)
if msg["type"] == "ls":
resp = ''
if msg["assets"]:
resp += self.list_assets(display=False) + ('\n' if msg["jobs"] or msg["schedules"] else '')
if msg["schedules"]:
resp += self.list_schedules(display=False) + ('\n' if msg["jobs"] else '')
if msg["jobs"]:
resp += self.list_jobs(display=False)
resp_msg = {'type' : 'response', 'response': resp}
publish_msg(msg["resp_queue"], resp_msg)
if msg["type"] == "pause":
asset = None
for _asset in self.graph:
if _asset.key == msg["key"]:
asset = _asset
if asset is not None:
self.pause_asset(asset)
if msg["key"] in self.schedules:
self.pause_schedule(msg["key"])
if msg["type"] == "unpause":
asset = None
for _asset in self.graph:
if _asset.key == msg["key"]:
asset = _asset
if asset is not None:
self.unpause_asset(asset)
if msg["key"] in self.schedules:
self.unpause_schedule(msg["key"])
if msg["type"] == "set_allow_retry":
asset = self._asset_by_key(msg["key"])
if asset is not None and isinstance(msg.get("value"), bool):
self.set_allow_retry(asset, msg["value"])
if msg["type"] == "schedule":
asset = None
for _asset in self.graph:
if _asset.key == msg["asset_key"]:
asset = _asset
self.schedule(msg["schedule_key"], msg["action"], cron_string = msg["cron_string"], asset = asset)
if msg["type"] == "unschedule":
self.unschedule(msg["schedule_key"])
if msg["type"] == "build":
asset = None
for _asset in self.graph:
if _asset.key == msg["asset_key"]:
asset = _asset
self.schedule('immediate-build-'+generate_slug(2), "build", '', asset = asset, immediate_once = True)
if msg["type"] == "refresh":
asset = None
for _asset in self.graph:
if _asset.key == msg["asset_key"]:
asset = _asset
self.schedule('immediate-refresh-'+generate_slug(2), "refresh", '', asset = asset, immediate_once = True)
if msg["type"] == "shutdown":
raise ShutdownSignal
[docs]
def run(self):
"""Start the graph server and block until a shutdown signal is received.
Launches a background thread that listens for RabbitMQ messages
addressed to this graph's key, and runs the async job scheduler in
the main thread. The method returns when:
* a ``shutdown`` message is received via RabbitMQ, or
* the user presses ``Ctrl-C``.
Typical usage::
g = fwirl.AssetGraph("my_graph")
g.add_assets([...])
g.schedule("hourly", "build", "0 * * * *")
g.run()
"""
# change the default sigint handler (which does not behave well with the use of asyncio.to_thread below
_original_sigint_handler = signal.getsignal(signal.SIGINT)
def _sigint_handler(sig, frame):
self.message_queue.put(0) # wake up the message queue thread
raise KeyboardInterrupt # then raise the interrupt
signal.signal(signal.SIGINT, _sigint_handler)
registration_file = None
th = None
try:
registration_file = register_graph_process(self.key)
# run the message handling loop
# need a separate thread for this since conn.drain_events() blocks, and kombu isn't compatible with asyncio yet
logger.info(f"Starting fwirl messaging loop")
th = Thread(name = "fwirl_messaging_loop", target=listen, args=(self.key, self.message_queue,), daemon=True)
th.start()
logger.info(f"Starting fwirl job event loop")
# run the task executor/scheduler async
asyncio.run(self._run())
except ValueError as e:
logger.error(str(e))
except KeyboardInterrupt:
logger.info(f"Caught keyboard interrupt; stopping main loop and messaging loop of asset graph {self.key}")
publish_msg(self.key, {"type": "shutdown"})
except ShutdownSignal:
# this only happens if the messaging loop got the "shutdown" message, so it is already shutting itself down, no need to publish shutdown msg
logger.info(f"Caught shutdown signal; stopping main loop of asset graph {self.key}")
finally:
if th is not None:
th.join()
unregister_graph_process(registration_file)
# restore the original sigint handler
signal.signal(signal.SIGINT, _original_sigint_handler)
async def _run(self):
message_task = None
while True:
# remove job from job queue if schedule does not exist or schedule paused
self.job_queue = [job for job in self.job_queue if job[0] in self.schedules and not self.schedules[job[0]].is_paused()]
for sk in self.schedules:
if self.schedules[sk].next() == Schedule.IMMEDIATE:
# if the immediate job is already in the queue, just skip; otherwise add it
if sk not in [job[0] for job in self.job_queue]:
self.job_queue.append((sk, Schedule.IMMEDIATE))
else:
# add all occurrences of each scheduled job onto the queue between the most recent one on the queue and the next future one
# get the latest current scheduled time for this schedule, or now if none scheduled
times_for_sk = [job[1] for job in self.job_queue if job[0] == sk]
if len(times_for_sk) > 0:
latest = max(times_for_sk)
else:
latest = plm.now()
# loop over runs starting from most recent one on the queue until the next one after the current time, adding all
while (not self.schedules[sk].is_paused()) and (latest <= plm.now()) and (self.schedules[sk].next(dt = latest) is not None):
new_time = latest + plm.duration(seconds = self.schedules[sk].next(dt = latest))
self.job_queue.append((sk, new_time))
latest = new_time
# sort upcoming jobs by scheduled time
self.job_queue.sort(key = lambda job : job[1])
# run the first job in the queue if its scheduled time is before now
if len(self.job_queue) > 0 and (self.job_queue[0][1] < plm.now()) and (self.job_running is None):
# pop the next job
next_sk, next_time = self.job_queue.pop(0)
if next_time == Schedule.IMMEDIATE:
logger.info(f"Running event {next_sk} ({self.schedules[next_sk]}) immediately")
else:
# check if it's late (30 second window)
logger.info(f"Running event {next_sk} ({self.schedules[next_sk]}) scheduled for {next_time} now")
if next_time + plm.duration(seconds=30) < plm.now():
logger.warning(f"Scheduled event {next_sk} is late; supposed to run at {next_time}, time now {plm.now()}")
logger.warning(f"Job queue currently has length {len(self.job_queue)}")
# run the job
task = asyncio.create_task(self.schedules[next_sk].generate_coroutine())
self.job_running = (next_sk, self.schedules[next_sk], 'job-'+generate_slug(2), task)
# if it was an immediate job, remove the schedule
if next_time == Schedule.IMMEDIATE:
self.schedules.pop(next_sk)
# if no message task is running, start one
# if job is running, wait indefinitely on message queue
# if no job is running, wait until next job
if message_task is None:
# if jobs are running, just wait indefinitely for next message
# if no jobs are running, wait for next job
timeout = None
if (self.job_running is None) and len(self.job_queue) > 0:
timeout = (self.job_queue[0][1] - plm.now()).seconds
#note: self.job_queue[0][1] should never be IMMEDIATE here, but if it somehow is, assert
assert self.job_queue[0][1] != Schedule.IMMEDIATE
message_task = asyncio.create_task(asyncio.to_thread(self._get_message, timeout))
# wait for next proc
await asyncio.wait([self.job_running[3], message_task] if self.job_running is not None else [message_task], return_when=asyncio.FIRST_COMPLETED)
# if the message task is done, process it
if message_task.done():
msg = await message_task
await self._process_message(msg)
message_task = None
# if the job is done, clear it
if self.job_running is not None and self.job_running[3].done():
res = await self.job_running[3]
self.job_running = None
def _initialize_resources(self, resources):
logger.info(f"Initializing {len(resources)} build resources")
for resource in resources:
try:
logger.debug(f"Initializing resource {resource}")
resource.init()
except Exception as e:
logger.exception(f"Resource {resource} initialization failed; cleaning up and aborting build")
self._cleanup_resources(resources)
return False
return True
def _cleanup_resources(self, resources):
logger.info(f"Cleaning up {len(resources)} build resources")
#cleanup resources
for resource in resources:
try:
logger.debug(f"Cleaning up resource {resource}")
resource.close()
except Exception as e:
logger.exception(f"Resource {resource} cleanup failed; attempting to clean up other resources")
[docs]
def refresh(self, assets=None):
"""Synchronously refresh the status of all (or specified) assets.
Checks each asset's :meth:`~fwirl.Asset.timestamp` against its
dependencies and updates :attr:`~fwirl.Asset.status` accordingly.
Does **not** call :meth:`~fwirl.Asset.build`.
Args:
assets: A single :class:`~fwirl.Asset` or list of assets to
refresh. When ``None`` (the default) every asset in the
graph is refreshed.
"""
asyncio.run(self._refresh(assets))
[docs]
def build(self, assets=None):
"""Synchronously build all (or specified) stale and unavailable assets.
Refreshes statuses first, then calls :meth:`~fwirl.Asset.build` on
every asset that is not current.
Args:
assets: A single :class:`~fwirl.Asset` or list of assets to
build (fwirl will also build all upstream dependencies).
When ``None`` (the default) all assets are considered.
"""
asyncio.run(self._build(assets))
async def _refresh(self, assets = None):
if not assets:
logger.info(f"Refreshing all assets")
sorted_nodes = list(nx.topological_sort(self.graph))
else:
logger.info(f"Refreshing assets downstream of (and including) {assets}")
if isinstance(assets, Asset):
assets = [assets]
nodes_to_refresh = []
for asset in assets:
nodes_to_refresh.append(asset)
nodes_to_refresh.extend(nx.descendants(self.graph, asset))
sg = self.graph.subgraph(nodes_to_refresh)
sorted_nodes = list(nx.topological_sort(sg))
required_resources = set()
for asset in sorted_nodes:
required_resources.update(asset.resources)
if not self._initialize_resources(required_resources):
return
task_map = {}
for asset in sorted_nodes:
coroutine = self._refresh_asset(asset)
task = asyncio.create_task(wait_for_dependencies(
coroutine,
[task_map[a] for a in self.graph.predecessors(asset) if a in task_map],
))
task_map[asset] = task
for asset in sorted_nodes:
await task_map[asset]
self._cleanup_resources(required_resources)
return
async def _build(self, assets = None):
workers_notified = False
while True: # keep rebuilding until graph structure doesn't change
if not assets:
logger.info(f"Building all assets")
sorted_nodes = list(nx.topological_sort(self.graph))
else:
logger.info(f"Building assets upstream of (and including) {assets}")
if isinstance(assets, Asset):
assets = [assets]
nodes_to_build = []
for asset in assets:
nodes_to_build.append(asset)
nodes_to_build.extend(nx.ancestors(self.graph, asset))
sg = self.graph.subgraph(nodes_to_build)
sorted_nodes = list(nx.topological_sort(sg))
required_resources = set()
for asset in sorted_nodes:
required_resources.update(asset.resources)
if not self._initialize_resources(required_resources):
return
# loop over assets, refresh and build
task_map = {}
for asset in sorted_nodes:
coroutine = self._refresh_asset(asset)
_task = asyncio.create_task(wait_for_dependencies(
coroutine,
[task_map[a] for a in self.graph.predecessors(asset) if a in task_map],
))
coroutine = self._build_asset(asset)
task = asyncio.create_task(wait_for_dependencies(coroutine, [_task]))
task_map[asset] = task
all_successful = True
workers_to_notify = set()
for asset in sorted_nodes:
build_result = await task_map[asset]
all_successful = all_successful and (build_result != BuildResult.Failed)
if build_result == BuildResult.Rebuilt:
for worker in self.workers:
if asset in worker.watched_assets:
workers_to_notify.add(worker)
self._cleanup_resources(required_resources)
summary = self.summarize(display=False)
if not all_successful:
logger.info(f"Build failure occurred")
logger.error(summary, colorize=True)
else:
logger.info(f"Build successful")
logger.info(summary, colorize=True)
if len(workers_to_notify) > 0:
logger.info(f"Build triggered graph workers. Restructuring...")
workers_notified = True
new_assets = []
for worker in workers_to_notify:
logger.info(f"Running graph worker {worker}")
new_assets.extend(await worker.restructure(self))
if len(new_assets) > 0:
logger.info(f"Graph structure changed. Triggering build of new assets...")
assets = new_assets
else:
logger.info(f"No new assets added to graph. No rebuild necessary.")
break
else:
logger.info(f"No graph workers triggered during build. No rebuild necessary.")
break
if workers_notified:
logger.info(f"Workers ran; removing any invalid schedules")
found_invalid = False
for sk in self.schedules:
asset = self.schedules[sk].kwargs["assets"]
if asset is not None and asset not in self.graph:
logger.warning(f"Asset {asset} no longer in graph after workers, but schedule {sk} ({self.schedules[sk]}) on it; removing")
self.unschedule(sk)
found_invalid = True
if not found_invalid:
logger.info(f"No invalid schedules found.")
logger.info(f"Build complete.")
async def _refresh_asset(self, asset):
# Status refresh propagates using the following cascade of rules:
# 0. If I'm Failed and self.allow_retry = False, I'm Failed
# 1. If I'm Paused, I'm Paused
# 2. If any of my parents is Paused or UpstreamStopped or Failed, I'm UpstreamStopped.
# 3. If I don't exist, I'm Unavailable.
# 4. If I'm Stale or Failed, I'm Stale.
# 4. If any of my parents is Stale or Unavailable, I'm Stale
# 5. All of my parents are Current. So check timestamps. If my timestamp is earlier than any of them, I'm Stale.
# 6. I'm Current
logger.debug(f"Updating status for asset {asset}")
if (not asset.allow_retry) and (asset.status == AssetStatus.Failed):
logger.info(f"Asset {asset} status: {fmt(asset.status)} (previous failure and asset does not allow retries)")
return
if asset.status == AssetStatus.Paused:
logger.info(f"Asset {asset} status: {fmt(asset.status)} (paused asset)")
return
any_paused_failed = False
any_stale_unav = False
latest_parent_timestamp = AssetStatus.Unavailable
for parent in self.graph.predecessors(asset):
logger.debug(f"Checking parent {parent} with status {parent.status} and timestamp {await parent.timestamp()}")
if parent.status == AssetStatus.Paused or parent.status == AssetStatus.UpstreamStopped or parent.status == AssetStatus.Failed:
any_paused_failed = True
if parent.status == AssetStatus.Stale or parent.status == AssetStatus.Unavailable:
any_stale_unav = True
# don't try to compute timestamps if any parent is paused, since they may not exist (and have no cached ts)
if parent.status != AssetStatus.Unavailable and (not any_paused_failed):
ts = await parent.timestamp()
if latest_parent_timestamp == AssetStatus.Unavailable:
latest_parent_timestamp = ts
else:
latest_parent_timestamp = latest_parent_timestamp if latest_parent_timestamp > ts else ts
logger.debug(f"After processing parent {parent}: any_paused_failed = {any_paused_failed}, any_stale_unav = {any_stale_unav}, latest_ts = {latest_parent_timestamp}")
if any_paused_failed:
asset.status = AssetStatus.UpstreamStopped
logger.info(f"Asset {asset} status: {fmt(asset.status)} (parent paused or failed)")
return
timestamp = await asset.timestamp()
if timestamp == AssetStatus.Unavailable:
asset.status = AssetStatus.Unavailable
logger.info(f"Asset {asset} status: {fmt(asset.status)} (timestamp unavailable)")
return
if asset.status == AssetStatus.Stale:
asset.status = AssetStatus.Stale
logger.info(f"Asset {asset} status: {fmt(asset.status)} (previously marked stale)")
return
if asset.status == AssetStatus.Failed:
asset.status = AssetStatus.Stale
logger.info(f"Asset {asset} status: {fmt(asset.status)} (previously failed but asset allows retries)")
return
if any_stale_unav:
asset.status = AssetStatus.Stale
logger.info(f"Asset {asset} status: {fmt(asset.status)} (parent stale/unavailable)")
return
#check for unavailable; if so, since any_stale_unav = False, this is a root node with no parent, so move on
if (latest_parent_timestamp != AssetStatus.Unavailable) and (timestamp < latest_parent_timestamp):
asset.status = AssetStatus.Stale
logger.info(f"Asset {asset} status: {fmt(asset.status)} (timestamp older than parent)")
return
asset.status = AssetStatus.Current
logger.info(f"Asset {asset} status: {fmt(asset.status)}")
async def _build_asset(self, asset):
parent_status_cts = Counter([p.status for p in self.graph.predecessors(asset)])
logger.debug(f"Checking asset {asset} with status {asset.status} and parent statuses {(' '.join([str(item[0]) +': ' + str(item[1]) + ',' for item in list(parent_status_cts.items())]))[:-1]}")
if (asset.status == AssetStatus.Unavailable or asset.status == AssetStatus.Stale) and all(p.status == AssetStatus.Current for p in self.graph.predecessors(asset)):
logger.info(f"Asset {asset} is {fmt(asset.status)} and all parents are current; rebuilding")
# set building status
asset.status = AssetStatus.Building
asset._last_build_timestamp = plm.now()
# run the build
try:
logger.debug(f"Building {asset}")
await asset.build()
except Exception as e:
asset.status = AssetStatus.Failed
asset.message = "Build failed: exception during build"
logger.exception(f"Asset {asset} build failed (status: {asset.status}): error in-build")
return BuildResult.Failed
# refresh the timestamp cache
ts = await asset.timestamp()
logger.debug(f"New asset {asset} timestamp: {ts}")
# get latest parent timestamp
parents = list(self.graph.predecessors(asset))
parent_timestamps = [(await p.timestamp()) for p in parents]
if len(parents) == 0:
logger.debug(f"Latest parent asset timestamp: (no parents)")
else:
latest_parent_timestamp = max(parent_timestamps)
logger.debug(f"Latest parent asset timestamp: {latest_parent_timestamp}")
if ts < latest_parent_timestamp:
raise RuntimeError("Rebuilt asset has timestamp earlier than latest parent!")
# make sure asset is now available and has a new timestamp
if ts is AssetStatus.Unavailable:
asset.status = AssetStatus.Failed
asset.message = "Build failed: asset unavailable"
logger.error(f"Asset {asset} build failed (status: {asset.status}): asset unavailable after build")
return BuildResult.Failed
# ensure that it now has a newer timestamp than parents
if not all(pt < ts for pt in parent_timestamps):
asset.status = AssetStatus.Failed
asset.message = "Build failed: asset timestamp older than parents"
logger.error(f"Asset {asset} build failed (status: {asset.status}): asset timestamp older than parents")
return BuildResult.Failed
# update to current status
asset.status = AssetStatus.Current
asset.message = "Build complete"
logger.info(f"Asset {asset} build completed successfully (status: {fmt(asset.status)})")
else:
logger.debug(f"Asset {asset} not ready to rebuild. Skipping.")
return BuildResult.Skipped
return BuildResult.Rebuilt
def _collect_groups(self):
logger.debug("Collecting node (sub)groups")
groupings = defaultdict(list)
# ensure __singletons__ group exists
groupings['__singletons__'] = []
for asset in self.graph:
if asset.group == '__singletons__' or asset.subgroup == '__singletons__':
raise ValueError('(Sub)group name used reserved keyword __singletons__')
if asset.group is not None:
if asset.group not in groupings:
groupings[asset.group] = defaultdict(list)
# ensure __singletons__ subgroup exists
groupings[asset.group]['__singletons__'] = []
if asset.subgroup is not None:
groupings[asset.group][asset.subgroup].append(asset)
else:
groupings[asset.group]['__singletons__'].append(asset)
else:
groupings['__singletons__'].append(asset)
# topologically sort the subgraphs
logger.debug("Topologically sorting subgroups")
for gr in groupings:
if gr == '__singletons__':
sg = self.graph.subgraph(groupings[gr])
groupings[gr] = list(nx.topological_sort(sg))
continue
for sb in groupings[gr]:
sg = self.graph.subgraph(groupings[gr][sb])
groupings[gr][sb] = list(nx.topological_sort(sg))
return groupings
[docs]
def summarize(self, display=True):
"""Return (and optionally log) a human-readable status summary.
The summary includes total asset and edge counts broken down by
:class:`~fwirl.AssetStatus`, grouped by the ``group``/``subgroup``
labels assigned to each asset.
Args:
display: When ``True`` (the default) the summary is also
emitted via the loguru logger.
Returns:
str: The formatted summary string.
"""
groupings = self._collect_groups()
logger.debug("Generating summary")
summary = f"\nAsset Graph Summary\n-------------------\nAssets: {self.graph.number_of_nodes()}\n"
for status in AssetStatus:
summary += f" {fmt(status)}: {sum([a.status == status for a in self.graph])} assets\n"
summary += f"Edges: {self.graph.size()}\n"
summary += f"Asset Groups: {len(groupings) - 1} groups, {len(groupings['__singletons__'])} singletons\n"
for gr in groupings:
# report individuals
if gr == '__singletons__':
# report summarized individuals
status_groups = defaultdict(int)
for asset in groupings[gr]:
if asset.status != AssetStatus.Failed:
status_groups[asset.status] += 1
summary += f" Singletons:\n"
for status in status_groups:
summary += f" {fmt(status)}: {status_groups[status]} singletons\n"
# report special cases
for asset in groupings[gr]:
if asset.status == AssetStatus.Failed:
summary += f" {asset}: {fmt(asset.status)}\n"
continue
summary += f" Group {gr}: {len(groupings[gr]) - 1} subgroups, {len(groupings[gr]['__singletons__'])} singletons\n"
# collect and report the summarized subgroups
status_groups = defaultdict(int)
for sb in groupings[gr]:
if sb == '__singletons__':
for asset in groupings[gr][sb]:
if asset.status != AssetStatus.Failed:
status_groups[asset.status] += 1
elif len({asset.status for asset in groupings[gr][sb]}) == 1 and groupings[gr][sb][0].status != AssetStatus.Failed:
status_groups[groupings[gr][sb][0].status] += 1
for sg in status_groups:
summary += f" {fmt(sg)}: {status_groups[sg]} subgroups/singletons\n"
# collect and report special cases
for sb in groupings[gr]:
if sb == '__singletons__':
for asset in groupings[gr][sb]:
if asset.status == AssetStatus.Failed:
summary += f" Singleton {asset}: {fmt(asset.status)}\n"
elif len({asset.status for asset in groupings[gr][sb]}) > 1 or groupings[gr][sb][0].status == AssetStatus.Failed:
summary += f" Subgroup {sb}:\n"
for asset in groupings[gr][sb]:
summary += f" {asset} : {fmt(asset.status)}\n"
if display:
logger.info(summary, colorize=True)
logger.debug("Done generating summary")
return summary
# TODO visualization (this code works but has old viznodes in it
# TODO recode this with new _collect_groups function
#def visualize(self):
# # in order to visualize large graphs
# # collapse all subgroups within a group with all constant state
# logger.info("Visualizing asset graph")
# groupings = {}
# logger.info("Assembling viznodes")
# for asset in self.graph:
# # if the asset is part of a group,subgroup
# if (asset.group is not None) and (asset.subgroup is not None):
# # collect assets in each group/subgroup
# if asset.group not in groupings:
# groupings[asset.group] = {}
# if asset.subgroup not in groupings[asset.group]:
# groupings[asset.group][asset.subgroup] = []
# groupings[asset.group][asset.subgroup].append(asset)
# else:
# #ungrouped assets get their own viznodes
# asset.viznode = ("single",asset.hash,asset.status)
# # for grouped assets, merge into one viznode if status is the same across the subgroup
# for group in groupings:
# for subgroup in groupings[group]:
# if len({asset.status for asset in groupings[group][subgroup]}) == 1:
# for asset in groupings[group][subgroup]:
# asset.viznode = ("group",group,asset.status)
# else:
# for asset in groupings[group][subgroup]:
# asset.viznode = ("single",asset.hash,asset.status)
# logger.info("Building vizgraph")
# # create the vizgraph
# vizgraph = nx.DiGraph()
# # add all viznodes
# for asset in self.graph:
# vizgraph.add_edges_from([(p.viznode, asset.viznode) for p in self.graph.predecessors(asset)])
# # remove any self-edges caused by viznode
# for node in vizgraph:
# try:
# vizgraph.remove_edge(node, node)
# except:
# pass
# logger.info(f"Quotient graph has {vizgraph.number_of_nodes()} nodes and {vizgraph.size()} edges.")
# logger.info("Computing quotient node positions")
# pos = nx.planar_layout(vizgraph)
# pos = nx.spring_layout(vizgraph, pos=pos) #initialize spring with planar
# logger.info("Drawing the graph")
# node_colors = [_NODE_COLORS[node[2]] for node in vizgraph]
# node_sizes = [600 if node[0] == "group" else 100 for node in vizgraph]
# nx.draw(vizgraph, pos=pos, node_color=node_colors, node_size=node_sizes)
# plt.show()