lsst-pipe-base 30.0.0rc2__py3-none-any.whl → 30.0.1rc1__py3-none-any.whl

lsst/pipe/base/quantum_graph/aggregator/_structs.py

@@ -27,68 +27,16 @@
 
 from __future__ import annotations
 
-__all__ = (
-    "InProgressScan",
-    "IngestRequest",
-    "ScanReport",
-    "ScanStatus",
-    "WriteRequest",
-)
+__all__ = ("IngestRequest", "ScanReport")
 
 import dataclasses
-import enum
 import uuid
 
+from lsst.daf.butler import DatasetRef
 from lsst.daf.butler.datastore.record_data import DatastoreRecordData
 
 from .._common import DatastoreName
-from .._predicted import PredictedDatasetModel
-from .._provenance import (
-    ProvenanceLogRecordsModel,
-    ProvenanceQuantumAttemptModel,
-    ProvenanceTaskMetadataModel,
-)
-
-
-class ScanStatus(enum.Enum):
-    """Status enum for quantum scanning.
-
-    Note that this records the status for the *scanning* which is distinct
-    from the status of the quantum's execution.
-    """
-
-    INCOMPLETE = enum.auto()
-    """The quantum is not necessarily done running, and cannot be scanned
-    conclusively yet.
-    """
-
-    ABANDONED = enum.auto()
-    """The quantum's execution appears to have failed but we cannot rule out
-    the possibility that it could be recovered, but we've also waited long
-    enough (according to `ScannerTimeConfigDict.retry_timeout`) that it's time
-    to stop trying for now.
-
-    This state means a later run with `ScannerConfig.assume_complete` is
-    required.
-    """
-
-    SUCCESSFUL = enum.auto()
-    """The quantum was conclusively scanned and was executed successfully,
-    unblocking scans for downstream quanta.
-    """
-
-    FAILED = enum.auto()
-    """The quantum was conclusively scanned and failed execution, blocking
-    scans for downstream quanta.
-    """
-
-    BLOCKED = enum.auto()
-    """A quantum upstream of this one failed."""
-
-    INIT = enum.auto()
-    """Init quanta need special handling, because they don't have logs and
-    metadata.
-    """
+from .._provenance import ProvenanceQuantumScanStatus
 
 
 @dataclasses.dataclass
@@ -98,7 +46,7 @@ class ScanReport:
     quantum_id: uuid.UUID
     """Unique ID of the quantum."""
 
-    status: ScanStatus
+    status: ProvenanceQuantumScanStatus
     """Combined status of the scan and the execution of the quantum."""
 
 
@@ -109,69 +57,11 @@ class IngestRequest:
     producer_id: uuid.UUID
     """ID of the quantum that produced these datasets."""
 
-    datasets: list[PredictedDatasetModel]
+    refs: list[DatasetRef]
     """Registry information about the datasets."""
 
     records: dict[DatastoreName, DatastoreRecordData]
     """Datastore information about the datasets."""
 
     def __bool__(self) -> bool:
-        return bool(self.datasets or self.records)
-
-
-@dataclasses.dataclass
-class InProgressScan:
-    """A struct that represents a quantum that is being scanned."""
-
-    quantum_id: uuid.UUID
-    """Unique ID for the quantum."""
-
-    status: ScanStatus
-    """Combined status for the scan and the execution of the quantum."""
-
-    attempts: list[ProvenanceQuantumAttemptModel] = dataclasses.field(default_factory=list)
-    """Provenance information about each attempt to run the quantum."""
-
-    outputs: dict[uuid.UUID, bool] = dataclasses.field(default_factory=dict)
-    """Unique IDs of the output datasets mapped to whether they were actually
-    produced.
-    """
-
-    metadata: ProvenanceTaskMetadataModel = dataclasses.field(default_factory=ProvenanceTaskMetadataModel)
-    """Task metadata information for each attempt.
-    """
-
-    logs: ProvenanceLogRecordsModel = dataclasses.field(default_factory=ProvenanceLogRecordsModel)
-    """Log records for each attempt.
-    """
-
-
-@dataclasses.dataclass
-class WriteRequest:
-    """A struct that represents a request to write provenance for a quantum."""
-
-    quantum_id: uuid.UUID
-    """Unique ID for the quantum."""
-
-    status: ScanStatus
-    """Combined status for the scan and the execution of the quantum."""
-
-    existing_outputs: set[uuid.UUID] = dataclasses.field(default_factory=set)
-    """Unique IDs of the output datasets that were actually written."""
-
-    quantum: bytes = b""
-    """Serialized quantum provenance model.
-
-    This may be empty for quanta that had no attempts.
-    """
-
-    metadata: bytes = b""
-    """Serialized task metadata."""
-
-    logs: bytes = b""
-    """Serialized logs."""
-
-    is_compressed: bool = False
-    """Whether the `quantum`, `metadata`, and `log` attributes are
-    compressed.
-    """
+        return bool(self.refs or self.records)

lsst/pipe/base/quantum_graph/aggregator/_supervisor.py

@@ -42,19 +42,18 @@ from lsst.utils.usage import get_peak_mem_usage
 from ...graph_walker import GraphWalker
 from ...pipeline_graph import TaskImportMode
 from .._predicted import PredictedQuantumGraphComponents, PredictedQuantumGraphReader
+from .._provenance import ProvenanceQuantumScanData, ProvenanceQuantumScanStatus
 from ._communicators import (
     IngesterCommunicator,
     ScannerCommunicator,
-    SpawnProcessContext,
     SupervisorCommunicator,
-    ThreadingContext,
-    Worker,
     WriterCommunicator,
 )
 from ._config import AggregatorConfig
 from ._ingester import Ingester
 from ._scanner import Scanner
-from ._structs import ScanReport, ScanStatus, WriteRequest
+from ._structs import ScanReport
+from ._workers import SpawnWorkerFactory, ThreadWorkerFactory
 from ._writer import Writer
 
 
@@ -116,6 +115,17 @@ class Supervisor:
                 self.comms.request_scan(ready_set.pop())
             for scan_return in self.comms.poll():
                 self.handle_report(scan_return)
+        if self.comms.config.incomplete:
+            quantum_or_quanta = "quanta" if self.n_abandoned != 1 else "quantum"
+            self.comms.progress.log.info(
+                "%d %s incomplete/failed abandoned; re-run with incomplete=False to finish.",
+                self.n_abandoned,
+                quantum_or_quanta,
+            )
+        self.comms.progress.log.info(
+            "Scanning complete after %0.1fs; waiting for workers to finish.",
+            self.comms.progress.elapsed_time,
+        )
 
     def handle_report(self, scan_report: ScanReport) -> None:
         """Handle a report from a scanner.
@@ -126,18 +136,22 @@ class Supervisor:
             Information about the scan.
         """
         match scan_report.status:
-            case ScanStatus.SUCCESSFUL | ScanStatus.INIT:
+            case ProvenanceQuantumScanStatus.SUCCESSFUL | ProvenanceQuantumScanStatus.INIT:
                 self.comms.log.debug("Scan complete for %s: quantum succeeded.", scan_report.quantum_id)
                 self.walker.finish(scan_report.quantum_id)
-            case ScanStatus.FAILED:
+            case ProvenanceQuantumScanStatus.FAILED:
                 self.comms.log.debug("Scan complete for %s: quantum failed.", scan_report.quantum_id)
                 blocked_quanta = self.walker.fail(scan_report.quantum_id)
                 for blocked_quantum_id in blocked_quanta:
-                    if self.comms.config.output_path is not None:
-                        self.comms.request_write(WriteRequest(blocked_quantum_id, status=ScanStatus.BLOCKED))
+                    if self.comms.config.is_writing_provenance:
+                        self.comms.request_write(
+                            ProvenanceQuantumScanData(
+                                blocked_quantum_id, status=ProvenanceQuantumScanStatus.BLOCKED
+                            )
+                        )
                     self.comms.progress.scans.update(1)
                 self.comms.progress.quantum_ingests.update(len(blocked_quanta))
-            case ScanStatus.ABANDONED:
+            case ProvenanceQuantumScanStatus.ABANDONED:
                 self.comms.log.debug("Abandoning scan for %s: quantum has not succeeded (yet).")
                 self.walker.fail(scan_report.quantum_id)
                 self.n_abandoned += 1
@@ -161,55 +175,31 @@ def aggregate_graph(predicted_path: str, butler_path: str, config: AggregatorCon
         Configuration for the aggregator.
     """
     log = getLogger("lsst.pipe.base.quantum_graph.aggregator")
-    ctx = ThreadingContext() if config.n_processes == 1 else SpawnProcessContext()
-    scanners: list[Worker] = []
-    ingester: Worker
-    writer: Worker | None = None
-    with SupervisorCommunicator(log, config.n_processes, ctx, config) as comms:
+    worker_factory = ThreadWorkerFactory() if config.n_processes == 1 else SpawnWorkerFactory()
+    with SupervisorCommunicator(log, config.n_processes, worker_factory, config) as comms:
         comms.progress.log.verbose("Starting workers.")
-        if config.output_path is not None:
+        if config.is_writing_provenance:
             writer_comms = WriterCommunicator(comms)
-            writer = ctx.make_worker(
+            comms.workers[writer_comms.name] = worker_factory.make_worker(
                 target=Writer.run,
                 args=(predicted_path, writer_comms),
                 name=writer_comms.name,
             )
-            writer.start()
         for scanner_id in range(config.n_processes):
             scanner_comms = ScannerCommunicator(comms, scanner_id)
-            worker = ctx.make_worker(
+            comms.workers[scanner_comms.name] = worker_factory.make_worker(
                 target=Scanner.run,
                 args=(predicted_path, butler_path, scanner_comms),
                 name=scanner_comms.name,
             )
-            worker.start()
-            scanners.append(worker)
         ingester_comms = IngesterCommunicator(comms)
-        ingester = ctx.make_worker(
+        comms.workers[ingester_comms.name] = worker_factory.make_worker(
             target=Ingester.run,
             args=(predicted_path, butler_path, ingester_comms),
             name=ingester_comms.name,
         )
-        ingester.start()
         supervisor = Supervisor(predicted_path, comms)
         supervisor.loop()
-        log.info(
-            "Scanning complete after %0.1fs; waiting for workers to finish.",
-            comms.progress.elapsed_time,
-        )
-        comms.wait_for_workers_to_finish()
-        if supervisor.n_abandoned:
-            raise RuntimeError(
-                f"{supervisor.n_abandoned} {'quanta' if supervisor.n_abandoned > 1 else 'quantum'} "
-                "abandoned because they did not succeed.  Re-run with assume_complete=True after all retry "
-                "attempts have been exhausted."
-            )
-    for w in scanners:
-        w.join()
-    ingester.join()
-    if writer is not None and writer.is_alive():
-        log.info("Waiting for writer process to close (garbage collecting can be very slow).")
-        writer.join()
     # We can't get memory usage for children until they've joined.
     parent_mem, child_mem = get_peak_mem_usage()
     # This is actually an upper bound on the peak (since the peaks could be

lsst/pipe/base/quantum_graph/aggregator/_workers.py

@@ -0,0 +1,303 @@
+# This file is part of pipe_base.
+#
+# Developed for the LSST Data Management System.
+# This product includes software developed by the LSST Project
+# (http://www.lsst.org).
+# See the COPYRIGHT file at the top-level directory of this distribution
+# for details of code ownership.
+#
+# This software is dual licensed under the GNU General Public License and also
+# under a 3-clause BSD license. Recipients may choose which of these licenses
+# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
+# respectively.  If you choose the GPL option then the following text applies
+# (but note that there is still no warranty even if you opt for BSD instead):
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+from __future__ import annotations
+
+__all__ = ("Event", "Queue", "SpawnWorkerFactory", "ThreadWorkerFactory", "Worker", "WorkerFactory")
+
+import multiprocessing.context
+import multiprocessing.synchronize
+import queue
+import threading
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import Any, Literal, overload
+
+_TINY_TIMEOUT = 0.01
+
+type Event = threading.Event | multiprocessing.synchronize.Event
+
+
+class Worker(ABC):
+    """A thin abstraction over `threading.Thread` and `multiprocessing.Process`
+    that also provides a variable to track whether it reported successful
+    completion.
+    """
+
+    def __init__(self) -> None:
+        self.successful = False
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Name of the worker, as assigned at creation."""
+        raise NotImplementedError()
+
+    @abstractmethod
+    def join(self, timeout: float | None = None) -> None:
+        """Wait for the worker to finish.
+
+        Parameters
+        ----------
+        timeout : `float`, optional
+            How long to wait in seconds.  If the timeout is exceeded,
+            `is_alive` can be used to see whether the worker finished or not.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def is_alive(self) -> bool:
+        """Return whether the worker is still running."""
+        raise NotImplementedError()
+
+    def kill(self) -> None:
+        """Kill the worker, if possible."""
+
+
+class Queue[T](ABC):
+    """A thin abstraction over `queue.Queue` and `multiprocessing.Queue` that
+    provides better control over disorderly shutdowns.
+    """
+
+    @overload
+    def get(self, *, block: Literal[True]) -> T: ...
+
+    @overload
+    def get(self, *, timeout: float | None = None, block: bool = False) -> T | None: ...
+
+    @abstractmethod
+    def get(self, *, timeout: float | None = None, block: bool = False) -> T | None:
+        """Get an object or return `None` if the queue is empty.
+
+        Parameters
+        ----------
+        timeout : `float` or `None`, optional
+            Maximum number of seconds to wait while blocking.
+        block : `bool`, optional
+            Whether to block until an object is available.
+
+        Returns
+        -------
+        obj : `object` or `None`
+            Object from the queue, or `None` if it was empty.  Note that this
+            is different from the behavior of the built-in Python queues,
+            which raise `queue.Empty` instead.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def put(self, item: T) -> None:
+        """Add an object to the queue.
+
+        Parameters
+        ----------
+        item : `object`
+            Item to add.
+        """
+        raise NotImplementedError()
+
+    def clear(self) -> bool:
+        """Clear out all objects currently on the queue.
+
+        This does not guarantee that more objects will not be added later.
+        """
+        found_anything: bool = False
+        while self.get() is not None:
+            found_anything = True
+        return found_anything
+
+    def kill(self) -> None:
+        """Prepare a queue for a disorderly shutdown, without assuming that
+        any other workers using it are still alive and functioning.
+        """
+
+
+class WorkerFactory(ABC):
+    """A simple abstract interface that can be implemented by both threading
+    and multiprocessing.
+    """
+
+    @abstractmethod
+    def make_queue(self) -> Queue[Any]:
+        """Make an empty queue that can be used to pass objects between
+        workers created by this factory.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def make_event(self) -> Event:
+        """Make an event that can be used to communicate a boolean state change
+        to workers created by this factory.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def make_worker(
+        self, target: Callable[..., None], args: tuple[Any, ...], name: str | None = None
+    ) -> Worker:
+        """Make a worker that runs the given callable.
+
+        Parameters
+        ----------
+        target : `~collections.abc.Callable`
+            A callable to invoke on the worker.
+        args : `tuple`
+            Positional arguments to pass to the callable.
+        name : `str`, optional
+            Human-readable name for the worker.
+
+        Returns
+        -------
+        worker : `Worker`
+            Process or thread that is already running the given callable.
+        """
+        raise NotImplementedError()
+
+
+class _ThreadWorker(Worker):
+    """An implementation of `Worker` backed by the `threading` module."""
+
+    def __init__(self, thread: threading.Thread):
+        super().__init__()
+        self._thread = thread
+
+    @property
+    def name(self) -> str:
+        return self._thread.name
+
+    def join(self, timeout: float | None = None) -> None:
+        self._thread.join(timeout=timeout)
+
+    def is_alive(self) -> bool:
+        return self._thread.is_alive()
+
+
+class _ThreadQueue[T](Queue[T]):
+    def __init__(self) -> None:
+        self._impl = queue.Queue[T]()
+
+    @overload
+    def get(self, *, block: Literal[True]) -> T: ...
+
+    @overload
+    def get(self, *, timeout: float | None = None, block: bool = False) -> T | None: ...
+
+    def get(self, *, timeout: float | None = None, block: bool = False) -> T | None:
+        try:
+            return self._impl.get(block=block, timeout=timeout)
+        except queue.Empty:
+            return None
+
+    def put(self, item: T) -> None:
+        self._impl.put(item, block=False)
+
+
+class ThreadWorkerFactory(WorkerFactory):
+    """An implementation of `WorkerFactory` backed by the `threading`
+    module.
+    """
+
+    def make_queue(self) -> Queue[Any]:
+        return _ThreadQueue()
+
+    def make_event(self) -> Event:
+        return threading.Event()
+
+    def make_worker(
+        self, target: Callable[..., None], args: tuple[Any, ...], name: str | None = None
+    ) -> Worker:
+        thread = threading.Thread(target=target, args=args, name=name)
+        thread.start()
+        return _ThreadWorker(thread)
+
+
+class _ProcessWorker(Worker):
+    """An implementation of `Worker` backed by the `multiprocessing` module."""
+
+    def __init__(self, process: multiprocessing.context.SpawnProcess):
+        super().__init__()
+        self._process = process
+
+    @property
+    def name(self) -> str:
+        return self._process.name
+
+    def join(self, timeout: float | None = None) -> None:
+        self._process.join(timeout=timeout)
+
+    def is_alive(self) -> bool:
+        return self._process.is_alive()
+
+    def kill(self) -> None:
+        """Kill the worker, if possible."""
+        self._process.kill()
+
+
+class _ProcessQueue[T](Queue[T]):
+    def __init__(self, impl: multiprocessing.Queue):
+        self._impl = impl
+
+    @overload
+    def get(self, *, block: Literal[True]) -> T: ...
+
+    @overload
+    def get(self, *, timeout: float | None = None, block: bool = False) -> T | None: ...
+
+    def get(self, *, timeout: float | None = None, block: bool = False) -> T | None:
+        try:
+            return self._impl.get(block=block, timeout=timeout)
+        except queue.Empty:
+            return None
+
+    def put(self, item: T) -> None:
+        self._impl.put(item, block=False)
+
+    def kill(self) -> None:
+        self._impl.cancel_join_thread()
+        self._impl.close()
+
+
+class SpawnWorkerFactory(WorkerFactory):
+    """An implementation of `WorkerFactory` backed by the `multiprocessing`
+    module, with new processes started by spawning.
+    """
+
+    def __init__(self) -> None:
+        self._ctx = multiprocessing.get_context("spawn")
+
+    def make_queue(self) -> Queue[Any]:
+        return _ProcessQueue(self._ctx.Queue())
+
+    def make_event(self) -> Event:
+        return self._ctx.Event()
+
+    def make_worker(
+        self, target: Callable[..., None], args: tuple[Any, ...], name: str | None = None
+    ) -> Worker:
+        process = self._ctx.Process(target=target, args=args, name=name)
+        process.start()
+        return _ProcessWorker(process)