import os
import math
import dask
import datetime
import functools
import inspect
import numpy as np
import toolviper.utils.logger as logger
from typing import Dict, Union
from typing import Callable, Any, Tuple, Dict
import xarray as xr
import copy
[docs]
def make_graph_node_task(node_task: Callable) -> Callable:
"""Adapt ``node_task`` to the single-``input_params``-dict calling convention.
A graph node task is always invoked with a single ``input_params`` dict (into
which :func:`map` injects the per-node keys ``task_id``, ``task_coords``,
``data_selection``, ``input_data`` ...). Historically every node task had to
accept that dict as its single argument. This helper lets a node task instead
have a **fully explicit, documented, standalone-callable signature**:
* If ``node_task`` follows the legacy convention -- it declares a parameter
named ``input_params`` (or takes a single argument) -- it is returned
unchanged.
* Otherwise a thin wrapper ``<name>_wrap(input_params)`` is returned that
expands the dict into ``node_task``'s explicit keyword arguments, forwarding
only the keys ``node_task`` declares (extra keys in ``input_params`` are
dropped, unless ``node_task`` accepts ``**kwargs``).
:func:`map` applies this automatically, so callers normally never need it --
they simply pass either an ``input_params``-style or an explicit node task.
Parameters
----------
node_task : callable
The node-task function (legacy single-dict or explicit signature).
Returns
-------
callable
``node_task`` unchanged (legacy), or a single-dict adapter named
``<node_task.__name__>_wrap``.
"""
sig = inspect.signature(node_task)
params = sig.parameters
has_var_kw = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
positional_like = [
p
for p in params.values()
if p.kind
in (
inspect.Parameter.POSITIONAL_ONLY,
inspect.Parameter.POSITIONAL_OR_KEYWORD,
inspect.Parameter.KEYWORD_ONLY,
)
]
# Legacy single-dict node task: it takes the whole input_params dict as one
# argument (named "input_params", or simply its only argument).
if "input_params" in params or (len(positional_like) == 1 and not has_var_kw):
return node_task
accepted = set(params)
@functools.wraps(node_task)
def wrap(input_params):
# Pin the mmap threshold BEFORE any large allocations so they use mmap and
# are returned to the OS immediately on free (no heap fragmentation). Must
# run at the start of the task, not after, or fragmentation is already done.
from toolviper.utils.memory_management import memory_setup, free_memory
memory_setup(131072)
if has_var_kw:
return node_task(**input_params)
return_dict = node_task(
**{k: v for k, v in input_params.items() if k in accepted}
)
free_memory()
return return_dict
wrap.__name__ = node_task.__name__ + "_wrap"
wrap.__qualname__ = getattr(node_task, "__qualname__", node_task.__name__) + "_wrap"
return wrap
def _sample_process_resources(interval, stop_event, samples):
"""Sampler-thread body: append one sample per ``interval`` to ``samples``.
Runs inside the worker process. Counters are PER PROCESS (psutil.Process()),
so attribution to the running task is exact only when the process executes
one task at a time (Dask threads_per_worker=1, or MPI ranks); with several
concurrent task threads the series describes the whole process.
A final sample is taken after the stop event fires so even tasks shorter
than ``interval`` get at least one point.
"""
import time
import psutil
proc = psutil.Process()
proc.cpu_percent(None) # establish the CPU% baseline (first call returns 0)
# Probe I/O counter availability once (Linux/Windows only -- hence the
# getattr; *_chars are the Linux syscall-level counters, which -- unlike
# *_bytes, which only count block-device traffic -- include network
# filesystems such as Lustre).
io_counters = getattr(proc, "io_counters", None)
try:
io0 = io_counters() if io_counters is not None else None
has_io = io0 is not None
has_chars = hasattr(io0, "read_chars")
except Exception:
has_io = has_chars = False
t0 = time.monotonic()
stopping = False
while True:
samples["time_seconds"].append(time.monotonic() - t0)
samples["cpu_percent"].append(proc.cpu_percent(None))
samples["memory_rss_bytes"].append(proc.memory_info().rss)
if has_io:
try:
io = io_counters()
samples["read_bytes"].append(io.read_bytes)
samples["write_bytes"].append(io.write_bytes)
if has_chars:
samples["read_chars"].append(io.read_chars)
samples["write_chars"].append(io.write_chars)
except Exception: # keep the series rectangular if counters vanish
has_io = False
if stopping:
return
# wait() returns True when the task finished -> take one last sample.
stopping = stop_event.wait(interval)
[docs]
def monitor_node_task(node_task, interval):
"""Wrap a (single-dict-convention) node task so its process CPU / memory /
I/O usage is sampled every ``interval`` seconds while it runs, and the
series is attached to the task's return dict as ``"resource_usage"``:
{"sample_interval_seconds": interval,
"start_unixtime": <time.time() when sampling began>,
"time_seconds": [...], "cpu_percent": [...],
"memory_rss_bytes": [...],
# present when the platform exposes them:
"read_bytes": [...], "write_bytes": [...], # block-level, cumulative
"read_chars": [...], "write_chars": [...]} # syscall-level, cumulative
# (includes Lustre/NFS)
Plain lists of numbers -- stdlib-pickle-serialisable, so the MPI backend's
worker->manager result path is safe. Tasks that do not return a dict pass
through unchanged (a debug line notes the dropped series). If psutil is
not installed the task runs unmonitored after a one-line warning.
``cpu_percent`` covers ALL threads of the process (OpenMP included) and can
exceed 100. See :func:`_sample_process_resources` for the per-process
attribution caveat.
"""
@functools.wraps(node_task)
def monitored(input_params):
try:
import psutil # noqa: F401
except ImportError:
logger.warning(
"monitor_resources_seconds is set but psutil is not installed; "
"running the node task unmonitored."
)
return node_task(input_params)
import threading
import time
samples = {
"time_seconds": [],
"cpu_percent": [],
"memory_rss_bytes": [],
"read_bytes": [],
"write_bytes": [],
"read_chars": [],
"write_chars": [],
}
stop_event = threading.Event()
sampler = threading.Thread(
target=_sample_process_resources,
args=(interval, stop_event, samples),
name="graphviper-resource-sampler",
daemon=True,
)
# Wall-clock anchor for the relative time_seconds series: lets an
# analysis place every task on the run's common timeline (cluster-wide
# usage at time t). Cross-node comparability relies on NTP-synced node
# clocks -- fine at typical sampling intervals.
start_unixtime = time.time()
sampler.start()
try:
return_dict = node_task(input_params)
finally:
stop_event.set()
sampler.join(timeout=max(5.0, 2 * interval))
if isinstance(return_dict, dict):
usage = {k: v for k, v in samples.items() if v}
usage["sample_interval_seconds"] = interval
usage["start_unixtime"] = start_unixtime
return_dict["resource_usage"] = usage
else:
logger.debug(
"monitor_node_task: node task did not return a dict; "
"resource-usage series dropped."
)
return return_dict
monitored.__name__ = node_task.__name__ + "_monitored"
monitored.__qualname__ = (
getattr(node_task, "__qualname__", node_task.__name__) + "_monitored"
)
return monitored
[docs]
def map(
input_data: Union[Dict, xr.DataTree],
node_task_data_mapping: dict,
node_task: Callable[..., Any],
input_params: dict,
in_memory_compute: bool = False,
client=None,
date_time: str = None,
data_loading_task: Union[Callable[..., Any], None] = None,
disk_chunk_sizes: Union[Dict[str, int], None] = None,
load_node_input_params: Union[dict, None] = None,
monitor_resources_seconds: Union[float, None] = None,
) -> Dict:
"""Create a perfectly parallel graph where a node is generated for each item in the :ref:`node_task_data_mapping <node task data mapping>` using the function specified in the ``node_task`` parameter.
Parameters
----------
input_data : Union[Dict, ProcessingSet]
Can either be a `ProcessingSet <https://github.com/casangi/xradio/blob/main/src/xradio/correlated_data/processing_set>`_ or a Dictionary of `xarray.Datasets <https://docs.xarray.dev/en/stable/generated/xarray.Dataset.html>`_. Only coordinates are needed so no actual data is loaded into memory (except if ``in_memory_compute`` is True).
node_task_data_mapping :
Node task data mapping dictionary. See :ref:`notes <node task data mapping>` for structure of dictionary.
node_task : Callable[..., Any]
The function that forms the nodes in the graph. It may either follow the
single-dict convention (one parameter, conventionally named
``input_params``, that receives the whole dict) or have a **fully
explicit signature** (the parameters it needs, spelled out); explicit
node tasks are adapted automatically via
:func:`make_graph_node_task`, so the spelled-out parameters are filled
from ``input_params`` (and the per-node keys ``task_id``,
``task_coords``, ``data_selection``, ``input_data`` ...). Extra keys not
declared by an explicit node task are dropped.
input_params : Dict
The input parameters to be passed to ``node_task``.
See notes for input parameters requirements.
in_memory_compute : optional
Setting ``in_memory_compute`` can lead to memory issues since all data is loaded into memory. Consequently, this option should only be used for testing purposes. If true the lazy arrays in ``input_data`` are loaded into memory using `xarray.Dataset.load <https://docs.xarray.dev/en/stable/generated/xarray.Dataset.load.html>`_ , by default False.
See notes of how to access data in ``node_task``.
client : optional
The Dask client is only required if local caching is enabled see :func:`toolviper.dask.client.slurm_cluster_client` , by default None.
date_time : str, optional
Used to annotate local cache, by default None.
data_loading_task : Callable[..., Any] or None, optional
Node task used by the optional data-loading layer to read one native
on-disk chunk of ``input_data`` at a time, shared across all map tasks
that fall within that chunk. The load layer is only built when **both**
``data_loading_task`` and ``disk_chunk_sizes`` are provided, by default None.
disk_chunk_sizes : Dict[str, int] or None, optional
``{dim: native_chunk_size}`` for the parallel dimensions whose on-disk
chunking should be used to coalesce I/O (see :func:`graphviper.graph_tools.coordinate_utils.get_disk_chunk_sizes`).
Required, together with ``data_loading_task``, to build the load layer, by default None.
load_node_input_params : dict or None, optional
Extra parameters merged into every load-node parameter dict (e.g.
``processing_set_data_group_name``); ``None`` is treated as ``{}``, by default None.
monitor_resources_seconds : float or None, optional
If set, sample the worker process's CPU / memory / I/O usage every this
many seconds while each map node task runs, and attach the series to the
task's return dict under ``"resource_usage"`` (see
:func:`monitor_node_task` for the exact keys and the per-process
attribution caveat -- accurate per-task attribution needs one concurrent
task per process, e.g. Dask ``threads_per_worker=1`` or MPI ranks).
Requires ``psutil`` (task runs unmonitored with a warning if missing).
By default None (no monitoring, zero overhead).
Returns
-------
Dict:
A graph-description dictionary with a ``"map"`` layer (keys ``"node_task"``,
``"input_params"``, and, when a load layer is built, ``"load_node_ids"`` and
``"relative_data_selections"``) and, when a load layer is built, a ``"load"``
layer. This dict is consumed by :func:`graphviper.graph_tools.generate_dask_workflow`
(or :func:`graphviper.graph_tools.processes_with_mpi`) to build/execute the actual graph.
Notes
-----
The ``input_params`` dictionary will be passed to each instance of the ``node_task`` along with the following items from the :ref:`node_task_data_mapping <node task data mapping>`:
- chunk_indices
- parallel_dims
- data_selection
- task_coords
- task_id
If ``in_memory_compute`` is ``False`` then ``input_params["input_data_store"]`` must be provided with a ``MutableMapping`` where a ``Zarr Group`` has been stored or a path to a directory in file system where a ``Zarr DirectoryStore`` has been stored.
If local caching is enabled the following will also be included with the `input_params` dictionary:
- date_time
- viper_local_dir
Example of how to access data in ``node_task``::
if input_params["input_data"] is None: #in_memory_compute==False
ps = load_data(
input_params["input_data_store"], data_selection=input_params["data_selection"]
)
else: #in_memory_compute==True
ps = input_params["input_data"]
, where ``load_data`` is appropriate function for your data.
"""
n_tasks = len(node_task_data_mapping)
# Allow the node task to have an explicit, documented signature instead of a
# single ``input_params`` dict: legacy single-dict node tasks are returned
# unchanged, explicit ones are wrapped so they are still called with one dict.
node_task = make_graph_node_task(node_task)
# Optional per-task resource sampling (CPU / memory / I/O series appended to
# the task's return dict). Applied AFTER the signature adapter so the sampler
# brackets the whole task body regardless of node-task style.
if monitor_resources_seconds:
node_task = monitor_node_task(node_task, monitor_resources_seconds)
# Get local_cache configuration if enabled in toolviper.dask.client.slurm_cluster_client.
# local_cache will be True if enabled.
(
local_cache,
viper_local_dir,
date_time,
tasks_to_node_map,
nodes_ip_list,
) = _local_cache_configuration(n_tasks, client, date_time)
input_param_list = []
# Create a node in Dask graph for each task_id in node_task_data_mapping
for task_id, node_task_parameters in node_task_data_mapping.items():
# logger.debug('Task id: ' + str(task_id) + ", " + str(node_task_parameters.keys()))
input_params.update(node_task_parameters)
input_params["task_id"] = task_id
if in_memory_compute: # Data gets loaded into memory.
input_params["input_data"] = _select_data(
input_data, input_params["data_selection"]
)
else:
input_params["input_data"] = None
if local_cache:
input_params["date_time"] = date_time
input_params["viper_local_dir"] = viper_local_dir
node_ip = nodes_ip_list[tasks_to_node_map[task_id]]
input_params["node_ip"] = node_ip
# with dask.annotate(resources={node_ip: 1}):
# graph_list.append(dask.delayed(node_task)(dask.delayed(input_params)))
else:
input_params["date_time"] = None
input_param_list.append(copy.deepcopy(input_params))
graph = {"map": {"node_task": node_task, "input_params": input_param_list}}
if data_loading_task is not None and disk_chunk_sizes is not None:
if load_node_input_params is None:
load_node_input_params = {}
load_input_params_list, load_node_ids, relative_data_selections = (
_build_load_stage(
input_param_list, disk_chunk_sizes, load_node_input_params
)
)
graph["load"] = {
"node_task": data_loading_task,
"input_params": load_input_params_list,
}
graph["map"]["load_node_ids"] = load_node_ids
graph["map"]["relative_data_selections"] = relative_data_selections
return graph
def _build_load_stage(
input_param_list: list,
disk_chunk_sizes: Dict[str, int],
load_node_input_params: dict,
):
"""Groups mapping tasks by their native on-disk chunk and builds load node parameters.
For each parallel dimension listed in ``disk_chunk_sizes`` the mapping tasks are
grouped by the disk chunk they fall in. Tasks that span a chunk boundary, or have
no disk-chunked dimension selections, are excluded from the load layer
(``load_node_id == -1``) and will fall back to the normal per-task data loading.
Parameters
----------
input_param_list : list of dict
Per-task parameter dicts produced by :func:`map`. Each dict must contain
``"input_data_store"`` and ``"data_selection"``.
disk_chunk_sizes : dict
``{dim: native_chunk_size}`` for every parallel dimension whose on-disk
chunking should be used to coalesce I/O.
load_node_input_params : dict
Extra parameters merged into every load node parameter dict
(e.g. ``processing_set_data_group_name``).
Returns
-------
load_input_params_list : list of dict
One parameter dict per unique disk-chunk group (= one load node).
Each dict contains ``"input_data_store"``, ``"data_selection"`` (at the
disk-chunk granularity), and any entries from ``load_node_input_params``.
load_node_ids : list of int
Index into ``load_input_params_list`` for each mapping task, or ``-1`` when
no load node applies to that task.
relative_data_selections : list of dict or None
Per-task data selection expressed as indices *relative to the start of the
pre-loaded disk chunk*, or ``None`` when ``load_node_id == -1``.
"""
disk_chunk_group_to_id: Dict = {}
load_input_params_list: list = []
load_node_ids: list = []
relative_data_selections: list = []
for task_params in input_param_list:
data_selection = task_params.get("data_selection", {})
group_key_parts = []
disk_level_selection: Dict = {}
relative_sel: Dict = {}
skip = False
for xds_name, xds_sel in data_selection.items():
disk_level_selection[xds_name] = {}
relative_sel[xds_name] = {}
for dim, sel in xds_sel.items():
if (
dim in disk_chunk_sizes
and isinstance(sel, slice)
and sel.start is not None
and sel.stop is not None
and sel.start >= 0
and sel.stop >= 0
):
chunk_size = disk_chunk_sizes[dim]
abs_start = sel.start
abs_stop = sel.stop
start_disk_chunk = abs_start // chunk_size
# abs_stop is exclusive, so the last included index is abs_stop - 1
end_disk_chunk = max(0, abs_stop - 1) // chunk_size
if start_disk_chunk != end_disk_chunk:
# Task spans a disk chunk boundary — skip the optimization.
logger.debug(
f"Load layer: task spans disk chunk boundary for dim '{dim}' "
f"(sel={sel}, chunk_size={chunk_size}); falling back to per-task loading."
)
skip = True
break
disk_start = start_disk_chunk * chunk_size
disk_stop = disk_start + chunk_size
disk_level_selection[xds_name][dim] = slice(disk_start, disk_stop)
relative_sel[xds_name][dim] = slice(
abs_start - disk_start,
abs_stop - disk_start,
)
group_key_parts.append((xds_name, dim, start_disk_chunk))
else:
# Dim not disk-chunked: load node loads everything; the
# original absolute slice is still valid relative to the
# loaded data (which starts at index 0 for this dim).
disk_level_selection[xds_name][dim] = slice(None)
relative_sel[xds_name][dim] = sel
if skip:
break
if skip or not group_key_parts:
load_node_ids.append(-1)
relative_data_selections.append(None)
continue
group_key = frozenset(group_key_parts)
if group_key not in disk_chunk_group_to_id:
new_id = len(load_input_params_list)
disk_chunk_group_to_id[group_key] = new_id
load_params: Dict = {
"input_data_store": task_params["input_data_store"],
"data_selection": disk_level_selection,
}
load_params.update(load_node_input_params)
load_input_params_list.append(load_params)
load_node_ids.append(disk_chunk_group_to_id[group_key])
relative_data_selections.append(relative_sel)
return load_input_params_list, load_node_ids, relative_data_selections
def _select_data(input_data, data_selection):
if isinstance(input_data, xr.DataTree):
input_data_sel = xr.DataTree()
else:
input_data_sel = {}
for xds_name, xds_isel in data_selection.items():
input_data_sel[xds_name] = input_data[xds_name].isel(xds_isel).load()
return input_data_sel
def _local_cache_configuration(n_tasks, client, date_time):
if "VIPER_LOCAL_DIR" in os.environ:
local_cache = True
viper_local_dir = os.environ["VIPER_LOCAL_DIR"]
if date_time is None:
date_time = datetime.datetime.utcnow().strftime("%y%m%d%H%M%S")
else:
local_cache = False
viper_local_dir = None
date_time = None
tasks_to_node_map = None
nodes_ip_list = None
return local_cache, viper_local_dir, date_time, tasks_to_node_map, nodes_ip_list
workers_info = client.scheduler_info()["workers"]
nodes_ip_list = _get_unique_resource_ip(workers_info)
n_nodes = len(nodes_ip_list)
tasks_per_compute_node = math.floor(n_tasks / n_nodes + 0.5)
if tasks_per_compute_node == 0:
tasks_per_compute_node = 1
tasks_to_node_map = np.repeat(np.arange(n_nodes), tasks_per_compute_node)
if len(tasks_to_node_map) < n_tasks:
n_pad = n_tasks - len(tasks_to_node_map)
tasks_to_node_map = np.concatenate(
[tasks_to_node_map, np.array([tasks_to_node_map[-1]] * n_pad)]
)
return local_cache, viper_local_dir, date_time, tasks_to_node_map, nodes_ip_list
def _get_unique_resource_ip(workers_info):
nodes = []
for worker, wi in workers_info.items():
worker_ip = worker[worker.rfind("/") + 1 : worker.rfind(":")]
assert worker_ip in list(wi["resources"].keys()), (
"local_cache enabled but workers have not been annotated. Make sure that local_cache has been set to True "
"during client setup."
)
if worker_ip not in nodes:
nodes.append(worker_ip)
return nodes