diff --git a/src/graphviper/graph_tools/__init__.py b/src/graphviper/graph_tools/__init__.py index c396084..40f88dd 100644 --- a/src/graphviper/graph_tools/__init__.py +++ b/src/graphviper/graph_tools/__init__.py @@ -5,7 +5,7 @@ interpolate_data_coords_onto_parallel_coords, get_disk_chunk_sizes, ) -from .map import map +from .map import map, make_graph_node_task from .reduce import reduce from .generate_dask_workflow import generate_dask_workflow from .generate_airflow_workflow import ( diff --git a/src/graphviper/graph_tools/map.py b/src/graphviper/graph_tools/map.py index 4106640..3f72ead 100644 --- a/src/graphviper/graph_tools/map.py +++ b/src/graphviper/graph_tools/map.py @@ -2,6 +2,8 @@ import math import dask import datetime +import functools +import inspect import numpy as np import toolviper.utils.logger as logger @@ -12,6 +14,80 @@ import copy +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 ``_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 + ``_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 map( input_data: Union[Dict, xr.DataTree], node_task_data_mapping: dict, @@ -33,7 +109,15 @@ def map( node_task_data_mapping : Node task data mapping dictionary. See :ref:`notes ` for structure of dictionary. node_task : Callable[..., Any] - The function that forms the nodes in the graph. The function must have a single input parameter that must be a dictionary. The ``input_params``, along with graph and coordinate-related parameters, will be passed to this function. + 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. @@ -81,6 +165,11 @@ def map( """ 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) + # Get local_cache configuration if enabled in toolviper.dask.client.slurm_cluster_client. # local_cache will be True if enabled. ( diff --git a/tests/test_graph_tools.py b/tests/test_graph_tools.py index 3b4f1e4..d6840a9 100644 --- a/tests/test_graph_tools.py +++ b/tests/test_graph_tools.py @@ -387,3 +387,82 @@ def my_load_fn(load_params): chunk_indx 11 (3, 2) """ + + +def test_make_graph_node_task(): + """make_graph_node_task: legacy single-dict tasks pass through; explicit + signatures get a single-dict adapter that expands/filters the dict.""" + from graphviper.graph_tools.map import make_graph_node_task + + # Legacy single-dict node task (named input_params) -> returned unchanged. + def legacy(input_params): + return input_params["a"] + + assert make_graph_node_task(legacy) is legacy + + # A single argument of any name is also treated as the legacy dict. + def legacy_other_name(params): + return params + + assert make_graph_node_task(legacy_other_name) is legacy_other_name + + # Explicit signature -> wrapped; the dict is expanded and extra keys dropped. + def explicit(a, b, c=3): + return (a, b, c) + + wrapped = make_graph_node_task(explicit) + assert wrapped is not explicit + assert wrapped.__name__ == "explicit_wrap" + assert wrapped({"a": 1, "b": 2, "unused": 99}) == (1, 2, 3) + + # Explicit signature with **kwargs -> every key is forwarded. + def explicit_var_kw(a, **kw): + return (a, kw) + + wrapped_kw = make_graph_node_task(explicit_var_kw) + assert wrapped_kw({"a": 1, "x": 2}) == (1, {"x": 2}) + + +def test_map_wraps_explicit_node_task(): + """map() auto-wraps an explicit-signature node task: graphviper users need + not know about the single-dict calling convention.""" + from graphviper.graph_tools.map import map + + seen = [] + + def explicit_task(task_id, scale, image_size=10): + seen.append((task_id, scale, image_size)) + return task_id * scale + + node_task_data_mapping = { + 1: { + "chunk_indices": (0,), + "parallel_dims": ["x"], + "data_selection": {}, + "task_coords": {}, + }, + 2: { + "chunk_indices": (1,), + "parallel_dims": ["x"], + "data_selection": {}, + "task_coords": {}, + }, + } + input_params = {"scale": 7, "image_size": 3, "extra_unused": "drop me"} + + graph = map( + input_data={}, + node_task_data_mapping=node_task_data_mapping, + node_task=explicit_task, + input_params=input_params, + in_memory_compute=False, + client=None, + ) + + stored = graph["map"]["node_task"] + assert stored.__name__ == "explicit_task_wrap" + + results = [stored(ip) for ip in graph["map"]["input_params"]] + assert sorted(results) == [7, 14] + # image_size was forwarded; extra_unused / input_data / chunk_indices dropped. + assert all(image_size == 3 for (_, _, image_size) in seen)