50 graphviper changes needed for astroviper demonstrator

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 (

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 ``<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 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 <node task data mapping>` 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.
     (

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)