50 graphviper changes needed for astroviper demonstrator

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "graphviper"
-version = "0.0.40"
+version = "0.0.41"
 description = "Visibility and Image Parallel Execution Reduction for Radio Astronomy "
 authors = [
     {name = "Jan-Willem Steeb", email="jsteeb@nrao.edu"},

src/graphviper/graph_tools/__init__.py

@@ -3,6 +3,7 @@
     make_frequency_coord,
     make_parallel_coord,
     interpolate_data_coords_onto_parallel_coords,
+    get_disk_chunk_sizes,
 )
 from .map import map
 from .reduce import reduce

src/graphviper/graph_tools/coordinate_utils.py

@@ -711,6 +711,97 @@ def interpolate_data_coords_onto_parallel_coords(
     return node_task_data_mapping
 
 
+def get_disk_chunk_sizes(
+    input_data: Union[Dict, xr.DataTree],
+    parallel_coords: Dict,
+) -> Dict[str, int]:
+    """Determine the native on-disk chunk size for each parallel coordinate dimension.
+
+    When a Zarr store is opened with ``xarray`` using a dask backend
+    (``chunks={}``), the resulting dask arrays carry the native zarr chunk sizes.
+    This function reads those chunk sizes from the coordinate arrays so that the
+    caller can pass them as ``disk_chunk_sizes`` to
+    :func:`graphviper.graph_tools.map.map`, enabling the data loading layer to
+    coalesce I/O at disk-chunk granularity.
+
+    Only the dimensions listed in *parallel_coords* are inspected; non-parallel
+    dimensions are ignored.
+
+    Parameters
+    ----------
+    input_data : Union[Dict, xr.DataTree]
+        The opened processing set — either an ``xr.DataTree`` or a plain dict of
+        ``xr.Dataset``.  The datasets must have been opened with a dask backend
+        (``chunks={}`` or equivalent) so that their coordinates are backed by
+        dask arrays with chunk information.  If a coordinate is numpy-backed (no
+        chunk information), the full coordinate length is used as the chunk size,
+        which is equivalent to treating the entire axis as a single disk chunk.
+    parallel_coords : Dict
+        The parallel coordinates dict as produced by :func:`make_parallel_coord`.
+        Only the keys (dimension names) are used.
+
+    Returns
+    -------
+    Dict[str, int]
+        ``{dim: native_chunk_size}`` for every dimension in *parallel_coords*
+        that is found in at least one dataset.  The returned chunk size is the
+        **minimum first-chunk size** observed across all datasets, which is the
+        conservative choice when datasets in the same processing set have
+        different on-disk chunking for the same dimension.
+
+    Notes
+    -----
+    Zarr stores the chunk size uniformly except for the last chunk, which may be
+    smaller if the array length is not a multiple of the chunk size.  This
+    function always reads the *first* chunk size (index 0) to get the nominal
+    chunk size and ignores the trailing remainder chunk.
+
+    Examples
+    --------
+    >>> ps_xdt = open_processing_set(ps_store)
+    >>> parallel_coords = {"frequency": make_parallel_coord(coord=img_xds.frequency, n_chunks=n_chunks)}
+    >>> disk_chunk_sizes = get_disk_chunk_sizes(ps_xdt, parallel_coords)
+    >>> # e.g. {"frequency": 200}
+    >>> viper_graph = map(..., disk_chunk_sizes=disk_chunk_sizes)
+    """
+    disk_chunk_sizes: Dict[str, int] = {}
+
+    for dim in parallel_coords:
+        min_chunk_size: Optional[int] = None
+
+        for xds_name, xds in input_data.items():
+            # DataTree nodes expose their dataset via .ds; plain dicts yield
+            # the dataset directly.
+            ds = xds.ds if isinstance(xds, xr.DataTree) else xds
+
+            if dim not in ds.coords:
+                continue
+
+            coord_array = ds[dim]
+            coord_data = coord_array.data  # dask array or numpy array
+
+            if hasattr(coord_data, "chunks") and coord_data.chunks:
+                # Dask array: chunks is a tuple-of-tuples, one tuple per axis.
+                # For a 1-D coordinate there is exactly one axis (index 0).
+                first_chunk = int(coord_data.chunks[0][0])
+            else:
+                # Numpy-backed (no dask): treat the whole axis as one chunk.
+                first_chunk = int(coord_array.size)
+
+            if min_chunk_size is None or first_chunk < min_chunk_size:
+                min_chunk_size = first_chunk
+
+        if min_chunk_size is not None:
+            disk_chunk_sizes[dim] = min_chunk_size
+        else:
+            logger.warning(
+                f"get_disk_chunk_sizes: dimension '{dim}' not found in any dataset; "
+                "skipping."
+            )
+
+    return disk_chunk_sizes
+
+
 # def interpolate_data_coords_onto_parallel_coords(
 #     parallel_coords: dict,
 #     input_data: Union[Dict, xr.DataTree],

src/graphviper/graph_tools/generate_dask_workflow.py

@@ -1,4 +1,7 @@
+import contextlib
+
 import dask
+from collections import defaultdict
 
 
 def _tree_combine(list_to_combine, reduce_node_task, input_params):
@@ -21,26 +24,155 @@ def _single_node(graph, reduce_node_task, input_params):
     return dask.delayed(reduce_node_task)(graph, input_params)
 
 
-def generate_dask_workflow(viper_graph):
-
-    dask_graph = []
-    for input_params in viper_graph["map"]["input_params"]:
-        dask_graph.append(
-            dask.delayed(viper_graph["map"]["node_task"])(dask.delayed(input_params))
-        )
-
-    if "reduce" in viper_graph:
-        if viper_graph["reduce"]["mode"] == "tree":
-            dask_graph = _tree_combine(
-                dask_graph,
-                viper_graph["reduce"]["node_task"],
-                viper_graph["reduce"]["input_params"],
-            )
-        elif viper_graph["reduce"]["mode"] == "single_node":
-            dask_graph = _single_node(
-                dask_graph,
-                viper_graph["reduce"]["node_task"],
-                viper_graph["reduce"]["input_params"],
-            )
-
-    return dask_graph
+def _prepare_task_input(loaded_data, relative_data_selection, input_params):
+    """Sub-select the task slice from a pre-loaded disk chunk and inject it into
+    the task parameter dict.
+
+    This is the single per-task node that sits between a shared load node and a
+    map node.  Combining sub-selection and injection into one step avoids the
+    extra intermediate dask node that would otherwise be required if they were
+    separate delayed calls.
+
+    The load node is shared across all map tasks that fall within the same
+    on-disk chunk; this function is NOT shared — one instance runs per map task.
+
+    Parameters
+    ----------
+    loaded_data : dict
+        ``{xds_name: xarray.Dataset}`` returned by a load node covering a full
+        disk chunk.
+    relative_data_selection : dict
+        ``{xds_name: {dim: slice}}`` with indices relative to the start of the
+        pre-loaded chunk.  ``slice(None)`` entries are skipped.
+    input_params : dict
+        Per-task parameter dict produced by :func:`graphviper.graph_tools.map.map`.
+
+    Returns
+    -------
+    dict
+        Shallow copy of *input_params* with ``"input_data"`` set to the
+        sub-selected dataset dict.
+    """
+    input_data = {}
+    for xds_name, xds_isel in relative_data_selection.items():
+        if xds_name not in loaded_data:
+            continue
+        ds = loaded_data[xds_name]
+        effective_sel = {
+            dim: sel for dim, sel in xds_isel.items() if sel != slice(None)
+        }
+        if effective_sel:
+            ds = ds.isel(effective_sel)
+        input_data[xds_name] = ds
+
+    result = dict(input_params)
+    result["input_data"] = input_data
+    return result
+
+
+def generate_dask_workflow(viper_graph, use_resource_restrictions=False):
+    """Build the Dask graph (``dask.delayed`` objects) for a viper map/reduce graph.
+
+    Parameters
+    ----------
+    viper_graph : dict
+        The graph description produced by :func:`graphviper.graph_tools.map.map`
+        (and optionally :func:`graphviper.graph_tools.reduce.reduce`).
+    use_resource_restrictions : bool, default False
+        When ``True`` every task is annotated with ``resources={"slots": 1}`` so
+        the distributed scheduler runs at most one node task per worker "slot".
+        This only works if the workers actually advertise a ``"slots"`` resource
+        (e.g. start the cluster with ``resources={"slots": N}``). If they do not,
+        the tasks are unschedulable and ``dask.compute`` hangs forever with no
+        scheduler activity. It is therefore **off by default** so graphs run on
+        any worker/scheduler (plain ``Client``, ``local_client``, synchronous)
+        without special cluster configuration.
+
+    Returns
+    -------
+    list or dask.delayed
+        The list of ``dask.delayed`` map nodes, or a single reduced node when a
+        reduce stage is present.
+    """
+    # Optional per-worker concurrency limit via a "slots" resource. Enabled only
+    # when requested AND the cluster advertises matching slots; a no-op context
+    # otherwise so the graph schedules on any worker.
+    resource_annotation = (
+        dask.annotate(resources={"slots": 1})
+        if use_resource_restrictions
+        else contextlib.nullcontext()
+    )
+
+    with resource_annotation:
+        dask_graph = []
+
+        if "load" in viper_graph:
+            # ------------------------------------------------------------------
+            # Build load nodes — one dask.delayed per unique disk-chunk group.
+            # Dask will compute each load node once even when multiple map nodes
+            # depend on it, avoiding redundant disk reads.
+            #
+            # Each load node and the prepare/map tasks it feeds are annotated
+            # with ``viper_load_group`` and ``viper_map_pair`` so that the
+            # ViperGraphPlugin scheduler plugin can assign priorities that
+            # (a) minimise the number of concurrently-loaded disk chunks and
+            # (b) schedule reduction-adjacent task pairs together.
+            # ------------------------------------------------------------------
+            load_fn = viper_graph["load"]["node_task"]
+            load_nodes = []
+            for load_id, lp in enumerate(viper_graph["load"]["input_params"]):
+                delayed_lp = dask.delayed(lp)
+                with dask.annotate(viper_load_group=load_id):
+                    load_nodes.append(dask.delayed(load_fn)(delayed_lp))
+
+            load_node_ids = viper_graph["map"]["load_node_ids"]
+            relative_data_selections = viper_graph["map"]["relative_data_selections"]
+            map_fn = viper_graph["map"]["node_task"]
+
+            # Count tasks per load group so pair_id can be computed incrementally.
+            group_task_counter: dict = defaultdict(int)
+
+            for i, input_params in enumerate(viper_graph["map"]["input_params"]):
+                load_node_id = load_node_ids[i]
+                delayed_params = dask.delayed(input_params)
+                if load_node_id == -1:
+                    # No load node for this task — fall back to per-task loading.
+                    dask_graph.append(dask.delayed(map_fn)(delayed_params))
+                else:
+                    load_node = load_nodes[load_node_id]
+                    rel_sel = relative_data_selections[i]
+                    # pair_id groups consecutive tasks within the same load group
+                    # into the pairs that will be combined at the first level of
+                    # the binary tree reduction.
+                    pair_id = group_task_counter[load_node_id] // 2
+                    group_task_counter[load_node_id] += 1
+                    with dask.annotate(
+                        viper_load_group=load_node_id, viper_map_pair=pair_id
+                    ):
+                        task_input = dask.delayed(_prepare_task_input)(
+                            load_node, rel_sel, delayed_params
+                        )
+                        dask_graph.append(dask.delayed(map_fn)(task_input))
+        else:
+            for input_params in viper_graph["map"]["input_params"]:
+                dask_graph.append(
+                    dask.delayed(viper_graph["map"]["node_task"])(
+                        dask.delayed(input_params)
+                    )
+                )
+
+        if "reduce" in viper_graph:
+            if viper_graph["reduce"]["mode"] == "tree":
+                dask_graph = _tree_combine(
+                    dask_graph,
+                    viper_graph["reduce"]["node_task"],
+                    viper_graph["reduce"]["input_params"],
+                )
+            elif viper_graph["reduce"]["mode"] == "single_node":
+                dask_graph = _single_node(
+                    dask_graph,
+                    viper_graph["reduce"]["node_task"],
+                    viper_graph["reduce"]["input_params"],
+                )
+
+        return dask_graph