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

Returns:

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