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):
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,
):
# 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):
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 {
"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
):
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
):
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
