graphviper.graph_tools.map

Functions

`make_graph_node_task`(→ Callable)	Adapt `node_task` to the single-`input_params`-dict calling convention.
`monitor_node_task`(node_task, interval)	Wrap a (single-dict-convention) node task so its process CPU / memory /
`map`(→ Dict)	Create a perfectly parallel graph where a node is generated for each item in the node_task_data_mapping using the function specified in the `node_task` parameter.

Module Contents

make_graph_node_task(node_task: Callable) → Callable[source]

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 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).

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:: node_task unchanged (legacy), or a single-dict adapter named <node_task.__name__>_wrap.
Return type:: callable

monitor_node_task(node_task, interval)[source]

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 _sample_process_resources() for the per-process attribution caveat.

map(input_data: Dict | xarray.DataTree, node_task_data_mapping: dict, node_task: Callable[Ellipsis, Any], input_params: dict, in_memory_compute: bool = False, client=None, date_time: str = None, data_loading_task: Callable[Ellipsis, Any] | None = None, disk_chunk_sizes: Dict[str, int] | None = None, load_node_input_params: dict | None = None, monitor_resources_seconds: float | None = None, task_priorities: dict | None = None) → Dict[source]

Create a perfectly parallel graph where a node is generated for each item in the node_task_data_mapping using the function specified in the node_task parameter.

Parameters:

input_data (Union[Dict, ProcessingSet]) – Can either be a ProcessingSet or a Dictionary of xarray.Datasets. 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 notes 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 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 , 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 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 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 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).
task_priorities (dict or None, optional) – {task_id: priority} scheduling priorities for the map tasks (higher runs earlier, matching Dask’s priority annotation convention). Task ids missing from the dict get no priority (Dask default 0). The Dask backend applies them as per-node dask.annotate(priority=...); the MPI backend (graphviper.graph_tools.processes_with_mpi()) dispatches tasks in priority order instead. Use this to control the order concurrent tasks are started in, e.g. so tasks writing to the same on-disk shard are spread out in time. By default None (all tasks equal priority, dispatch in task_id order).

Returns:

A graph-description dictionary with a "map" layer (keys "node_task", "input_params", "task_priorities" when priorities were given, 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 graphviper.graph_tools.generate_dask_workflow() (or graphviper.graph_tools.processes_with_mpi()) to build/execute the actual graph.

Return type:

Dict

Notes

The input_params dictionary will be passed to each instance of the node_task along with the following items from the 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.