Pipeline: add an optional 'if' and 'after' parameters for conditional tasks and better control flow

caput/pipeline/_pipeline.py

@@ -539,8 +539,9 @@ def _validate_task_inputs(self):
         for task_spec in self.task_specs:
             in_ = task_spec.get("in", None)
             requires = task_spec.get("requires", None)
+            after = task_spec.get("after", None)
 
-            for key, value in (["in", in_], ["requires", requires]):
+            for key, value in (["in", in_], ["requires", requires], ["after", after]):
                 if value is None:
                     continue
                 if not isinstance(value, list):
@@ -566,7 +567,7 @@ def _get_task_from_spec(self, task_spec):
         """
         # Check that only the expected keys are in the task spec.
         for key in task_spec.keys():
-            if key not in ["type", "params", "requires", "in", "out"]:
+            if key not in ["type", "params", "requires", "in", "out", "if", "after"]:
                 raise config.CaputConfigError(
                     f"Task got an unexpected key '{key}' in 'tasks' list."
                 )
@@ -577,9 +578,13 @@ def _get_task_from_spec(self, task_spec):
         except KeyError as e:
             raise config.CaputConfigError("'type' not specified for task.") from e
 
+        # If the task specifies a conditional and that condition is false,
+        # replace this task with a :py:class:`NoOp`
+        if "if" in task_spec and not task_spec["if"]:
+            task_cls = NoOp
         # Find the tasks class either in the local set, or by importing a fully
         # qualified class name
-        if task_path in local_tasks:
+        elif task_path in local_tasks:
             task_cls = local_tasks[task_path]
         else:
             try:
@@ -673,9 +678,10 @@ def _check_duplicate(key0: str, key1: str, d0: dict, d1: dict):
         requires = _check_duplicate("requires", "requires", task_spec, kwargs)
         in_ = _check_duplicate("in", "in_", task_spec, kwargs)
         out = _check_duplicate("out", "out", task_spec, kwargs)
+        after = _check_duplicate("after", "after", task_spec, kwargs)
 
         try:
-            task._setup_keys(in_, out, requires)
+            task._setup_keys(in_, out, requires, after)
         # Want to blindly catch errors
         except Exception as exc:
             raise config.CaputConfigError(f"Error adding task {task!s}\n") from exc
@@ -691,9 +697,9 @@ def _check_duplicate(key0: str, key1: str, d0: dict, d1: dict):
 def _format_product_keys(keys):
     """Formats the pipeline task product keys.
 
-    In the pipeline config task list, the values of 'requires', 'in' and 'out'
-    are keys representing data products.  This function gets that key from the
-    task's entry of the task list, defaults to zero, and ensures it's formated
+    In the pipeline config task list, the values of 'requires', 'in', 'out', and
+    'after' are keys representing data products.  This function gets that key from
+    the task's entry of the task list, defaults to zero, and ensures it's formated
     as a sequence of strings.
     """
     if keys is None:
@@ -750,12 +756,17 @@ class attributes will be overridden with instance attributes with the same
         If true, signals to the pipeline runner to make a call to `breakpoint` each
         time this task is run. This will drop the interpreter into pdb, allowing for
         interactive debugging of the current pipeline and task state. Default is False.
+    after_first_only : bool
+        If true, keys in the `after` queue only have to be received once, even if `next`
+        iterates multiple times. Otherwise, and iteration of `after` keys must be received
+        prior to each iteration. Default is False.
     """
 
     broadcast_inputs = config.Property(proptype=bool, default=False)
     limit_outputs = config.Property(proptype=int, default=None)
     base_priority = config.Property(proptype=int, default=0)
     breakpoint = config.Property(proptype=bool, default=False)
+    after_first_only = config.Property(proptype=bool, default=False)
 
     # Overridable Attributes
     # -----------------------
@@ -855,6 +866,11 @@ def _pipeline_is_available(self):
             # This task hasn't been initialized
             return False
 
+        if self._after is not None and not bool(
+            min((q.qsize() for q in self._after), default=1)
+        ):
+            return False
+
         if self._pipeline_state == "setup":
             # True if all `requires` items have been provided
             # This also returns True is `self._requires` is empty
@@ -939,8 +955,8 @@ def _from_config(cls, config):
 
         return self
 
-    def _setup_keys(self, in_, out=None, requires=None):  # noqa: D417
-        r"""Setup the 'requires', 'in' and 'out' keys for this task.
+    def _setup_keys(self, in_, out=None, requires=None, after=None):  # noqa: D417
+        r"""Setup the 'requires', 'in', 'out', and 'after' keys for this task.
 
         Parameters
         ----------
@@ -950,6 +966,9 @@ def _setup_keys(self, in_, out=None, requires=None):  # noqa: D417
             The names of the output data products.
         requires : Sequence[str] | str | None
             The names of the required data products for the `setup` method.
+        after : Sequence[str] | str | None
+            The names of data products which must be produced for this task
+            to proceed, even if they are not used by this task
 
         Raises
         ------
@@ -961,6 +980,7 @@ def _setup_keys(self, in_, out=None, requires=None):  # noqa: D417
         requires = _format_product_keys(requires)
         in_ = _format_product_keys(in_)
         out = _format_product_keys(out)
+        after = _format_product_keys(after)
 
         # Inspect the `setup` method to see how many arguments it takes.
         setup_argspec = inspect.getfullargspec(self.setup)
@@ -1024,6 +1044,10 @@ def _setup_keys(self, in_, out=None, requires=None):  # noqa: D417
         # produce multiple values, queue up items which may be used in the
         # future
         self._in = [queue.Queue() for _ in range(n_in)]
+        # Store after keys
+        self._after_keys = after
+        # Set up queues for wait keys
+        self._after = [queue.Queue() for _ in range(len(after))]
         # Store output keys
         self._out_keys = out
         # Keep track of the number of times this task has produced output
@@ -1078,6 +1102,7 @@ def _pipeline_advance_state(self):
                     )
 
             self._in = None
+            self._after = None
             self._pipeline_state = "finish"
 
         elif self._pipeline_state == "finish":
@@ -1125,6 +1150,9 @@ def _pipeline_next(self):
             else:  # noqa RET506
                 # Get the next set of data to be run.
                 args = tuple(in_.get() for in_ in self._in)
+                # If `after` keys are not pinned, remove them from the queue
+                if not self.after_first_only:
+                    _ = [q.get() for q in self._after]
 
             # Call the next iteration of `next`. If it is done running,
             # advance the task state and continue
@@ -1226,4 +1254,36 @@ def _pipeline_queue_product(self, key, product):
 
                 result = True
 
+        if key in self._after_keys:
+            indices = (ii for ii, k in enumerate(self._after_keys) if k == key)
+
+            for ii in indices:
+                logger.debug(
+                    f"{self!s} stowing data product with key {key} for `after`."
+                )
+                if self._after is None:
+                    raise exceptions.PipelineRuntimeError(
+                        f"Tried to queue 'after' data product, but task state is {self._pipeline_state}"
+                    )
+                # data product isn't actually needed - just record that it was revceived
+                self._after[ii].put(True)
+
+                result = True
+
         return result
+
+
+class NoOp(Task):
+    """Pass along an input.
+
+    Primarily useful to support optional tasks.
+    """
+
+    def next(self, *input):
+        """Forward any input."""
+        return input
+
+    @classmethod
+    def _from_config(cls, config):
+        """Ignore any config parameters because this task only cares about keys."""
+        return super()._from_config({})

caput/pipeline/tasklib/base.py

@@ -371,8 +371,9 @@ def next(self, *input):
         if not isinstance(output, tuple):
             output = (output,)
 
-        # Insert the input tags into the output containers
-        for opt in output:
+        # Insert the input tags into the output containers. Output
+        # may contain nested iterables
+        for opt in _iterate_nested(output):
             opt.attrs["input_tags"] = input_tags
 
         # Process each output individually
@@ -417,29 +418,44 @@ def finish(self):
         return output if len(output) > 1 else output[0]
 
     def _process_output(self, output, ii=0):
-        """Check a single output container and write it if needed."""
-        if not isinstance(output, MemDiskGroup):
-            raise exceptions.PipelineRuntimeError(
-                f"Task must output a valid memdata container; given {type(output)}"
-            )
+        """Check a single output and write it if needed.
 
-        # Set the tag according to the format
-        idict = self._interpolation_dict(output, ii)
+        The output is allowed to be a list of containers, where each
+        container is handled in the same way.
+        """
+        if not isinstance(output, (list, tuple)):
+            output = [output]
+            _single_output = True
+        else:
+            _single_output = False
 
-        # Set the attributes in the output container (including from the `tag` config
-        # option)
-        attrs_to_set = {} if self.attrs is None else self.attrs.copy()
-        attrs_to_set["tag"] = self.tag
-        for attrname, attrval in attrs_to_set.items():
-            if isinstance(attrval, str):
-                attrval = attrval.format(**idict)
-            output.attrs[attrname] = attrval
+        for opt in output:
+            if not isinstance(opt, MemDiskGroup):
+                raise exceptions.PipelineRuntimeError(
+                    f"Task must output a valid memdata container; given {type(opt)}"
+                )
+
+            # Set the tag according to the format
+            idict = self._interpolation_dict(opt, ii)
+
+            # Set the attributes in the output container (including from the `tag` config
+            # option)
+            attrs_to_set = {} if self.attrs is None else self.attrs.copy()
+            attrs_to_set["tag"] = self.tag
+            for attrname, attrval in attrs_to_set.items():
+                if isinstance(attrval, str):
+                    attrval = attrval.format(**idict)
+                opt.attrs[attrname] = attrval
+
+            # Check for NaN's etc
+            opt = self._nan_process_output(opt)
 
-        # Check for NaN's etc
-        output = self._nan_process_output(output)
+            # Write the output if needed
+            self._save_output(opt, ii)
 
-        # Write the output if needed
-        self._save_output(output, ii)
+        if _single_output:
+            # make sure that we recover the original output type
+            output = output[0]
 
         return output
 
@@ -628,6 +644,14 @@ def _nan_check_walk(self, cont):
         return self.comm.allreduce(found, op=MPI.MAX)
 
 
+def _iterate_nested(x):
+    for item in x:
+        if isinstance(item, (list, tuple, set)):
+            yield from _iterate_nested(item)
+        else:
+            yield item
+
+
 # Alias for backwards compatibility
 class SingleTask(ContainerTask):
     """Alias for :py:meth:`~caput.pipeline.tasklib.ContainerTask` for backwards compatibility."""

caput/pipeline/tasklib/io.py

@@ -36,6 +36,8 @@
 
 import numpy as np
 
+from caput.config import CaputConfigError
+
 from ... import config
 from ...containers import Container
 from ...memdata import MemDataset, MemDiskGroup, fileformats
@@ -595,6 +597,108 @@ def setup(self, files):
         self.files = files
 
 
+class LoadAllFiles(LoadFilesFromParams):
+    """Load all files input files in a single pass and return as a list of containers."""
+
+    def process(self):
+        """Load all files before returning.
+
+        Returns
+        -------
+        list[Container]
+            List of loaded containers.
+        """
+        filelist = []
+
+        while True:
+            try:
+                filelist.append(super().process())
+            except PipelineStopIteration:
+                break
+
+        return filelist
+
+
+class LoadFileBatches(BaseLoadFiles):
+    """Load batches of files and pass along an entire batch.
+
+    File batches are passed as a dict where each entry is a list of files.
+    On each pass, this task loads one item from each entry.
+
+    In addition, each entry is allowed to be a list of files (globs not allowed).
+    In this case, the entire sub-list is loaded and passed as a list.
+
+    The top-level lists must all have equal length.
+
+    Attributes
+    ----------
+    file_batches : dict[str, PathLik | list[PathLike]]
+        Batches of files to be loaded together.
+    """
+
+    file_batches = config.Property(proptype=dict)
+
+    _file_ind = 0
+    _max_len = None
+
+    def setup(self):
+        """Ensure that all batch lists have the same length."""
+        for key, entry in self.file_batches.items():
+            if self._max_len is None:
+                self._max_len = len(entry)
+
+            if len(entry) != self._max_len:
+                raise CaputConfigError(
+                    f"Batch entry `{key}` has a different length from the rest: "
+                    f"{len(entry)} != {self._max_len}"
+                )
+        # `SelectionsMixin` setup to resolve selections
+        super().setup()
+
+    def process(self):
+        """Load a batch of files and return.
+
+        Returns
+        -------
+        list[Container | list[Container]]
+            Single batch of files
+        """
+        if self._file_ind == self._max_len:
+            raise PipelineStopIteration
+
+        outputs = {key: [] for key in self.file_batches.keys()}
+
+        for key in self.file_batches.keys():
+            entry = self.file_batches[key][self._file_ind]
+
+            # package each entry as a list so we can always just iterate,
+            # even if there's only one item
+            if not isinstance(entry, list):
+                entry = [entry]
+
+            for ii, item in enumerate(entry):
+                message = f"[batch {self._file_ind + 1}/{self._max_len}; item {ii + 1}/{len(entry)} for entry `{key}`]"
+
+                cont = self._load_file(item, extra_message=message)
+                # repeat attr checks
+                if "tag" not in cont.attrs:
+                    # Get the first part of the actual filename and use it as the tag
+                    tag = os.path.splitext(os.path.basename(item))[0]
+
+                    cont.attrs["tag"] = tag
+
+                outputs[key].append(cont)
+
+        self._file_ind += 1
+
+        # construct the output list and return
+        for key, item in outputs.items():
+            if len(item) == 1:
+                outputs[key] = item[0]
+
+        return tuple(outputs.values())
+
+
 class Save(ContainerTask):
     """Save out the input, and pass it on.