import code
import logging
import os
from functools import partial
from graphlib import TopologicalSorter
from collections import defaultdict
from .......utils import path
from .......utils.check.errclass import CifError, CifImportError
from .individual_transform import do_individual_transform
def _reinit_worker_logging(level):
"""Re-establish a minimal CIF logging configuration in a multiprocessing worker.
Each worker process starts without CIF's init_log having been called.
Some import in dask's worker bootstrap triggers logging.basicConfig(),
which installs a plain WARNING-level handler before any user code runs.
This initializer replaces that with a ColorFormatter-equipped stream
handler at the same level as the main process.
The file handler is intentionally omitted: concurrent writes from
multiple worker processes to the same log file are unsafe without an
explicit multiprocessing-safe log handler (e.g. QueueHandler).
"""
from .......utils.check.coloredlog import ColorFormatter
root = logging.getLogger()
for hdl in root.handlers[:]:
root.removeHandler(hdl)
handler = logging.StreamHandler()
handler.setFormatter(ColorFormatter(fmt="#(level)%(message)s"))
root.addHandler(handler)
root.setLevel(level)
try:
import dask
delayed = dask.delayed
DASK_AVAILABLE = True
try:
import dask.distributed # pre-import in main thread: prevents a
# first-time import race when xarray's get_write_lock() calls
# _get_scheduler() from a worker thread during to_netcdf()
except ImportError:
pass
except ImportError:
dask = None
DASK_AVAILABLE = False
def delayed(func):
"""Fallback for ``dask.delayed`` when Dask is not installed: runs eagerly.
Args:
func: The function that would otherwise be wrapped as a Dask
delayed task.
Returns:
callable: ``func`` itself, unmodified, so decorated functions
still work (just executed synchronously) without Dask.
"""
return func # passthrough: runs eagerly
[docs]
def plot_task_graph(tasks, dependencies, rundir):
"""Render the resolved transform DAG as an interactive HTML graph.
dask's own visualize() rasterizes through Graphviz, which becomes
unreadable past a few dozen nodes. pyvis (vis.js) instead produces a
pannable/zoomable/draggable HTML page with hover tooltips, which stays
usable with the hundreds of transform tasks a typical CIF run builds.
networkx and pyvis are optional (``pip install networkx pyvis`` or the
``graph`` extra): imported lazily here so their absence never breaks a
run that doesn't request plotting.
Args:
tasks (dict): mapping of task name -> delayed object, as built by
init_dask's main loop.
dependencies (dict): mapping of task name -> list of precursor
task names, as built by init_dask before the main loop.
rundir (str): directory to write ``dask_graph.html`` into.
"""
try:
import networkx as nx
from pyvis.network import Network
except ImportError:
logging.warning(
"plot_dask_graph is True but 'networkx' and/or 'pyvis' are not "
"installed; skipping the graph plot. Install them with "
"'pip install networkx pyvis' (or the 'graph' extra) to enable it."
)
return
direction_colors = {"forward": "#4C72B0", "adjoint": "#C44E52"}
graph = nx.DiGraph()
for name in tasks:
ddi, transform, direction = name
graph.add_node(
str(name),
label=f"{transform}\n{direction}",
title=f"{transform} ({direction})<br>{ddi}",
color=direction_colors.get(direction, "#888888"),
)
for name in tasks:
for dep in dependencies.get(name, []):
if dep in tasks:
graph.add_edge(str(dep), str(name))
net = Network(
height="900px", width="100%",
directed=True, notebook=False,
cdn_resources="in_line",
)
net.from_nx(graph)
net.show_buttons(filter_=["physics"])
out_file = os.path.join(rundir, "dask_graph.html")
net.write_html(out_file, notebook=False, open_browser=False)
logging.info(f"Dask task graph written to {out_file}")
[docs]
def init_dask(
self, pipe_links,
mode="fwd",
do_simu=True,
onlyinit=False,
check_transforms=False,
adj_test_threshold=10,
save_debug=False,
ignore_exceptions=False,
ref_fwd_dir="",
run_id=0,
):
"""Build and execute the Dask task graph for one direction of the pipe.
Inverts ``pipe_links`` (which maps each node to its successors) into
a ``dependencies`` mapping (node -> precursors), topologically sorts
it, augments it with extra serialization dependencies coming from
``force_dump``/``force_loadout`` I/O constraints
(:func:`update_dependencies_with_io`), then walks the sorted nodes to
build one Dask delayed task per transform (:func:`add_transform`),
skipping dead branches and nodes with no resolved dependency. Also
attaches a "dry run" metadata dependency on the opposite-direction
counterpart of each node so that side-effect-only initialization runs
before the real computation.
If ``self.plot_dask_graph`` is set, writes an interactive HTML graph
of the resulting task DAG (:func:`plot_task_graph`).
Finally triggers the computation of the graph's final task
(``final_toobsvect`` in forward/tangent-linear mode,
``final_fromcontrol`` in adjoint mode), using the scheduler resolved
from ``self.dask_mode`` (``_resolve_dask_scheduler``), restoring the
root logger's level afterwards in case a worker left it altered.
Args:
self: The obs operator, exposing ``workdir``, ``transform_pipe``,
``datei``, ``controlvect``, ``obsvect``, ``plot_dask_graph``
and ``dask_mode``.
pipe_links: Mapping of each ``(ddi, transform, direction)`` node to
the list of its successor nodes, as built by
:func:`..fwd_pipe.fwd_adj_pipe`.
mode: Execution mode: ``"fwd"``, ``"tl"`` or ``"adj"``.
do_simu: Whether to actually run the underlying model simulations.
onlyinit: Whether to only initialize (not execute) every transform.
check_transforms: Whether to keep debug information needed by the
adjoint/tangent-linear consistency test.
adj_test_threshold: Threshold passed through to transforms for the
adjoint/tangent-linear consistency test.
save_debug: Whether to dump debug datasets for each transform.
ignore_exceptions: Whether to swallow exceptions raised by
transforms unless required outputs are missing.
ref_fwd_dir: Path to the reference forward run directory, used by
adjoint/tangent-linear transforms that need forward outputs.
run_id: Numeric id used to name this run's sub-directory under
``{self.workdir}/obsoperator/{mode}_{run_id}``.
Raises:
CifImportError: If Dask is not installed.
Returns:
None. The task graph is executed for its side effects (writing
to the shared datastore via each transform's ``forward``/
``adjoint`` method).
"""
# Import error if no dask
if not DASK_AVAILABLE:
raise CifImportError(
"Dask is not available. Install it with 'pip install dask'"
)
# Create of sub- working directory for the present run
workdir = self.workdir
rundir = f"{workdir}/obsoperator/{mode}_{run_id:04d}/"
path.init_dir(rundir)
# First invert the dependency graph
# Pipe_links give successors, whereas we need precursors
dependencies = defaultdict(list)
for parent, children in pipe_links.items():
for child in children:
dependencies[child].append(parent)
ts = TopologicalSorter(dependencies)
static_order = list(ts.static_order())
# Start/end points
end_point = ("", "final_fromcontrol", "adjoint") if mode == "adj" \
else ("", "final_toobsvect", "forward")
start_point = ("", "final_fromcontrol", "forward") if mode == "adj" \
else ("", "final_toobsvect", "adjoint")
# Update dependencies with I/O dependencies
dependencies = update_dependencies_with_io(
self, static_order, end_point, start_point, dependencies
)
# Initialize the dask graph
tasks = {}
# Now initialize the delayed functions
for name in static_order:
# Skip if dead branch
if name != end_point:
if pipe_links.get(name, []) == []:
continue
# Look up the actual delayed objects for the dependencies
dep_objects = {
d: tasks[d]["main"] for d in dependencies[name]
if d in tasks}
meta_objects = {
d: tasks[d]["meta"] for d in dependencies[name]
if d in tasks
}
# Stop here if no dependencies was ever initialize
# That mean a dead branch
if name != start_point:
if dep_objects == {}:
continue
# Update meta with dry run transform
if name != end_point:
if mode == "fwd" and name[2] == "forward":
dry_run = (name[0], name[1], "adjoint")
meta_objects[dry_run] = tasks[dry_run]["meta"]
elif mode == "adj" and name[2] == "adjoint":
dry_run = (name[0], name[1], "forward")
meta_objects[dry_run] = tasks[dry_run]["meta"]
# Create the new task using the registry
tasks[name] = add_transform(
name, dep_objects, meta_objects,
self.transform_pipe,
self.datei, rundir, workdir,
self.controlvect, self.obsvect,
mode=mode,
do_simu=do_simu,
onlyinit=onlyinit,
check_transforms=check_transforms,
adj_test_threshold=adj_test_threshold,
save_debug=save_debug,
ignore_exceptions=ignore_exceptions,
ref_fwd_dir=ref_fwd_dir,
)
# Plot the task graph if requested
if self.plot_dask_graph:
plot_task_graph(tasks, dependencies, rundir)
def _resolve_dask_scheduler(dask_mode):
"""Translate a user-facing Dask mode into a scheduler name Dask accepts.
Args:
dask_mode: A falsy value, one of the aliases ``"sync"``,
``"single-threaded"``, ``"threaded"``, ``"multiprocessing"``,
or any scheduler name/object already understood by
:meth:`dask.compute` (e.g. ``"threads"``, ``"processes"``,
``"synchronous"``, or a distributed client).
Returns:
The resolved scheduler, or None if ``dask_mode`` is falsy (use
Dask's default scheduler).
"""
if not dask_mode:
return None
aliases = {
"sync": "synchronous",
"single-threaded": "synchronous",
"threaded": "threads",
"multiprocessing": "processes",
}
return aliases.get(dask_mode, dask_mode)
scheduler = _resolve_dask_scheduler(self.dask_mode)
# Now do the computation. Accept classical dask scheduler options such as
# "threads", "processes", "single-threaded" / "synchronous",
# as well as custom scheduler objects supported by dask.compute.
if mode in ["fwd", "tl"]:
final_task = tasks[("", "final_toobsvect", "forward")]
else:
final_task = tasks[("", "final_fromcontrol", "adjoint")]
# Safety net: some transform (e.g. GribDataset's cfgrib open) may
# temporarily raise the root logger's level and, on an unguarded
# exception, leave it raised. Since the root logger is shared process-
# wide (not thread-local), that silently kills debug/info/warning
# output for the rest of the run, on every thread, well past this
# compute() call. Restore the level unconditionally once dask is done.
root_logger = logging.getLogger()
orig_level = root_logger.level
try:
if scheduler is None:
final_task.compute()
elif scheduler == "processes":
worker_log_init = partial(_reinit_worker_logging, logging.getLogger().level)
final_task.compute(scheduler=scheduler, initializer=worker_log_init)
else:
final_task.compute(scheduler=scheduler)
finally:
root_logger.setLevel(orig_level)
return
[docs]
def entry_point():
"""Return an empty result shaped like a real transform task's output.
Used as the result of the virtual ``final_toobsvect``/
``final_fromcontrol`` graph nodes, which have no actual transform to
run but must still produce a value in the shape downstream/upstream
tasks expect.
Returns:
dict: ``{"main": {"inputs": {}, "outputs": {}}, "meta": []}``.
"""
return {
"main": {"inputs": {}, "outputs": {}},
"meta": []
}
@delayed
def add_transform(
name, dep_objects, meta_objects,
transform_pipe, datei_ref,
rundir, workdir,
controlvect, obsvect,
mode="fwd", init_input=None,
do_simu=True,
onlyinit=False,
check_transforms=False,
adj_test_threshold=10,
save_debug=False,
ignore_exceptions=False,
ref_fwd_dir="",
**kwargs
):
"""Dask-delayed task wrapping the execution of a single transform node.
Registered as a :func:`dask.delayed` function (or run eagerly if Dask
is unavailable, see the module-level ``delayed`` fallback). For the
virtual ``final_toobsvect``/``final_fromcontrol`` nodes, returns an
empty placeholder result (:func:`entry_point`) instead of running a
transform. Otherwise, delegates to
:func:`..individual_transform.do_individual_transform`.
Args:
name: ``(ddi, transform, direction)`` tuple identifying the task.
dep_objects: Dict of ``{dependency_name: delayed_result["main"]}``
for this task's precursor tasks, used by Dask to infer the
graph edges and pass their resolved values in.
meta_objects: Dict of ``{dependency_name: delayed_result["meta"]}``
for this task's precursor tasks (plus its own dry-run
counterpart when applicable).
transform_pipe: Object exposing the transform instances and their
shared ``mapper``.
datei_ref: Reference start date of the run.
rundir: Base run directory for this Dask run.
workdir: Working directory passed through to the transform.
controlvect: Control vector object passed through to the transform.
obsvect: Observation vector object passed through to the transform.
mode: Execution mode: ``"fwd"``, ``"tl"`` or ``"adj"``.
init_input: Unused here, kept for interface consistency.
do_simu: Whether to actually run the underlying model simulation.
onlyinit: Whether to only initialize (not execute) the transform.
check_transforms: Whether to keep debug information needed by the
adjoint/tangent-linear consistency test.
adj_test_threshold: Threshold passed through to the transform for
the adjoint/tangent-linear consistency test.
save_debug: Whether to dump debug datasets for this transform.
ignore_exceptions: Whether to swallow exceptions raised by the
transform unless required outputs are missing.
ref_fwd_dir: Path to the reference forward run directory.
**kwargs: Additional keyword arguments forwarded to
:func:`do_individual_transform`.
Returns:
dict: The task result, as returned by :func:`entry_point` for
virtual nodes, or by
:func:`..individual_transform.do_individual_transform` otherwise.
"""
if name[1] in ["final_toobsvect", "final_fromcontrol"]:
return entry_point()
else:
func = do_individual_transform
return func(
name, dep_objects, meta_objects,
transform_pipe, datei_ref,
rundir, workdir,
controlvect, obsvect,
mode=mode,
do_simu=do_simu,
onlyinit=onlyinit,
check_transforms=check_transforms,
adj_test_threshold=adj_test_threshold,
save_debug=save_debug,
ignore_exceptions=ignore_exceptions,
ref_fwd_dir=ref_fwd_dir,
**kwargs
)
[docs]
def update_dependencies_with_io(
self, static_order,
end_point, start_point, dependencies
):
"""Add extra Dask dependencies to serialize transforms sharing dump/load files.
Some transforms write/read shared files on disk (via ``dump2inputs``
in forward mode when a trid is flagged ``force_dump``, or
``loadfromoutputs`` in adjoint mode when a trid is flagged
``force_loadout``). Because Dask has no visibility into that
filesystem side effect, this walks every node's mapper for such
flags and, for each declared dependency
(``dumpin_dependencies``/``loadout_dependencies``), adds an explicit
edge in ``dependencies`` between the corresponding
``dump2inputs``/``loadfromoutputs`` transform ids (resolved via
``dump2inputs_ids``/``loadfromoutputs_ids`` in the mapper) so that
Dask executes them in the required order.
Args:
self: The obs operator, exposing ``transform_pipe.mapper``.
static_order: Topologically sorted list of
``(ddi, transform, direction)`` nodes to scan.
end_point: The virtual end-of-pipe node, skipped during the scan.
start_point: The virtual start-of-pipe node, skipped during the scan.
dependencies: Mapping of each node to the list of its precursor
nodes, mutated in place with the extra I/O-derived edges.
Raises:
CifError: If a transform declares a dump/load dependency on
another trid for which no corresponding
``dump2inputs``/``loadfromoutputs`` transform was created.
Returns:
dict: The same ``dependencies`` mapping, augmented in place.
"""
for transform in static_order:
if transform in [end_point, start_point]:
continue
transf = transform[1]
transf_mapper = self.transform_pipe.mapper[transf]
# Now resolve dependencies for inputs
direction = transform[2]
ddi = transform[0]
if direction == "forward":
for trid in transf_mapper["inputs"]:
if not transf_mapper["inputs"][trid].get("force_dump", False):
continue
# Add extra dependency to ensure serialization of dump2inputs
# if required by the transform
dump2inputs_deps = transf_mapper["inputs"][trid].get(
'dumpin_dependencies', [])
if dump2inputs_deps == []:
continue
# Loop over dependencies
dump2inputs_ids = transf_mapper.get("dump2inputs_ids", {})
for dep in dump2inputs_deps:
if dep not in dump2inputs_ids:
raise CifError(
f"Transform {transform} requires dump2inputs "
f"for {trid} to depend on {dep}, but no "
f"dump2inputs transform was created for {dep}.")
# Add dependency
dependencies[(ddi, dump2inputs_ids[trid], direction)].append(
(ddi, dump2inputs_ids[dep], direction))
if direction == "adjoint":
for trid in transf_mapper["outputs"]:
if not transf_mapper["outputs"][trid].get("force_loadout", False):
continue
# Add extra dependency to ensure serialization of loadfromoutputs
# if required by the transform
loadout_deps = transf_mapper["outputs"][trid].get(
'loadout_dependencies', [])
if loadout_deps == []:
continue
# Loop over dependencies
loadfromoutputs_ids = transf_mapper.get("loadfromoutputs_ids", {})
for dep in loadout_deps:
if dep not in loadfromoutputs_ids:
raise CifError(
f"Transform {transform} requires loadfromoutputs "
f"for {trid} to depend on {dep}, but no "
f"loadfromoutputs transform was created for {dep}.")
# Add dependency
dependencies[(ddi, loadfromoutputs_ids[trid], direction)].append(
(ddi, loadfromoutputs_ids[dep], direction))
return dependencies
