lsst-pipe-base 30.0.1rc1__py3-none-any.whl → 30.2025.5200__py3-none-any.whl

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

@@ -27,16 +27,68 @@
 
 from __future__ import annotations
 
-__all__ = ("IngestRequest", "ScanReport")
+__all__ = (
+    "InProgressScan",
+    "IngestRequest",
+    "ScanReport",
+    "ScanStatus",
+    "WriteRequest",
+)
 
 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 .._provenance import ProvenanceQuantumScanStatus
+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.
+    """
 
 
 @dataclasses.dataclass
@@ -46,7 +98,7 @@ class ScanReport:
     quantum_id: uuid.UUID
     """Unique ID of the quantum."""
 
-    status: ProvenanceQuantumScanStatus
+    status: ScanStatus
     """Combined status of the scan and the execution of the quantum."""
 
 
@@ -57,11 +109,69 @@ class IngestRequest:
     producer_id: uuid.UUID
     """ID of the quantum that produced these datasets."""
 
-    refs: list[DatasetRef]
+    datasets: list[PredictedDatasetModel]
     """Registry information about the datasets."""
 
     records: dict[DatastoreName, DatastoreRecordData]
     """Datastore information about the datasets."""
 
     def __bool__(self) -> bool:
-        return bool(self.refs or self.records)
+        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.
+    """

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

@@ -42,18 +42,19 @@ 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
-from ._workers import SpawnWorkerFactory, ThreadWorkerFactory
+from ._structs import ScanReport, ScanStatus, WriteRequest
 from ._writer import Writer
 
 
@@ -115,17 +116,6 @@ 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.
@@ -136,22 +126,18 @@ class Supervisor:
             Information about the scan.
         """
         match scan_report.status:
-            case ProvenanceQuantumScanStatus.SUCCESSFUL | ProvenanceQuantumScanStatus.INIT:
+            case ScanStatus.SUCCESSFUL | ScanStatus.INIT:
                 self.comms.log.debug("Scan complete for %s: quantum succeeded.", scan_report.quantum_id)
                 self.walker.finish(scan_report.quantum_id)
-            case ProvenanceQuantumScanStatus.FAILED:
+            case ScanStatus.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.is_writing_provenance:
-                        self.comms.request_write(
-                            ProvenanceQuantumScanData(
-                                blocked_quantum_id, status=ProvenanceQuantumScanStatus.BLOCKED
-                            )
-                        )
+                    if self.comms.config.output_path is not None:
+                        self.comms.request_write(WriteRequest(blocked_quantum_id, status=ScanStatus.BLOCKED))
                     self.comms.progress.scans.update(1)
                 self.comms.progress.quantum_ingests.update(len(blocked_quanta))
-            case ProvenanceQuantumScanStatus.ABANDONED:
+            case ScanStatus.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
@@ -175,31 +161,55 @@ def aggregate_graph(predicted_path: str, butler_path: str, config: AggregatorCon
         Configuration for the aggregator.
     """
     log = getLogger("lsst.pipe.base.quantum_graph.aggregator")
-    worker_factory = ThreadWorkerFactory() if config.n_processes == 1 else SpawnWorkerFactory()
-    with SupervisorCommunicator(log, config.n_processes, worker_factory, config) as comms:
+    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:
         comms.progress.log.verbose("Starting workers.")
-        if config.is_writing_provenance:
+        if config.output_path is not None:
             writer_comms = WriterCommunicator(comms)
-            comms.workers[writer_comms.name] = worker_factory.make_worker(
+            writer = ctx.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)
-            comms.workers[scanner_comms.name] = worker_factory.make_worker(
+            worker = ctx.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)
-        comms.workers[ingester_comms.name] = worker_factory.make_worker(
+        ingester = ctx.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