diff --git a/caput/pipeline/_pipeline.py b/caput/pipeline/_pipeline.py index 36d921fb..e94b6017 100644 --- a/caput/pipeline/_pipeline.py +++ b/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({}) diff --git a/caput/pipeline/tasklib/base.py b/caput/pipeline/tasklib/base.py index 9af355f6..33cd008a 100644 --- a/caput/pipeline/tasklib/base.py +++ b/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.""" diff --git a/caput/pipeline/tasklib/io.py b/caput/pipeline/tasklib/io.py index 16a69955..2b04b101 100644 --- a/caput/pipeline/tasklib/io.py +++ b/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.