Coverage for python / lsst / pipe / base / quantum_graph / aggregator / _communicators.py: 28%
356 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-24 01:12 -0800
1# This file is part of pipe_base.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28from __future__ import annotations
30__all__ = (
31 "FatalWorkerError",
32 "IngesterCommunicator",
33 "ScannerCommunicator",
34 "SpawnProcessContext",
35 "SupervisorCommunicator",
36 "ThreadingContext",
37 "WorkerContext",
38)
40import cProfile
41import dataclasses
42import enum
43import logging
44import multiprocessing.context
45import multiprocessing.synchronize
46import os
47import queue
48import signal
49import threading
50import time
51import uuid
52from abc import ABC, abstractmethod
53from collections.abc import Callable, Iterable, Iterator
54from contextlib import ExitStack
55from traceback import format_exception
56from types import TracebackType
57from typing import Any, Literal, Self, TypeAlias, TypeVar, Union
59from lsst.utils.logging import LsstLogAdapter
61from .._provenance import ProvenanceQuantumScanData
62from ._config import AggregatorConfig
63from ._progress import ProgressManager, make_worker_log
64from ._structs import IngestRequest, ScanReport
66_T = TypeVar("_T")
68_TINY_TIMEOUT = 0.01
70# multiprocessing.Queue is a type according to the standard library type stubs,
71# but it's really a function at runtime. But since the Python <= 3.11 type
72# alias syntax uses the real runtime things we need to use strings, and hence
73# we need to use Union. With Python 3.12's 'type' statement this gets cleaner.
74Queue: TypeAlias = Union["queue.Queue[_T]", "multiprocessing.Queue[_T]"]
76Event: TypeAlias = threading.Event | multiprocessing.synchronize.Event
78Worker: TypeAlias = threading.Thread | multiprocessing.context.SpawnProcess
81class WorkerContext(ABC):
82 """A simple abstract interface that can be implemented by both threading
83 and multiprocessing.
84 """
86 @abstractmethod
87 def make_queue(self) -> Queue[Any]:
88 """Make an empty queue that can be used to pass objects between
89 workers in this context.
90 """
91 raise NotImplementedError()
93 @abstractmethod
94 def make_event(self) -> Event:
95 """Make an event that can be used to communicate a boolean state change
96 to workers in this context.
97 """
98 raise NotImplementedError()
100 @abstractmethod
101 def make_worker(
102 self, target: Callable[..., None], args: tuple[Any, ...], name: str | None = None
103 ) -> Worker:
104 """Make a worker that runs the given callable.
106 Parameters
107 ----------
108 target : `~collections.abc.Callable`
109 A callable to invoke on the worker.
110 args : `tuple`
111 Positional arguments to pass to the callable.
112 name : `str`, optional
113 Human-readable name for the worker.
115 Returns
116 -------
117 worker : `threading.Thread` or `multiprocessing.Process`
118 Process or thread. Will need to have its ``start`` method called
119 to actually begin.
120 """
121 raise NotImplementedError()
124class ThreadingContext(WorkerContext):
125 """An implementation of `WorkerContext` backed by the `threading`
126 module.
127 """
129 def make_queue(self) -> Queue[Any]:
130 return queue.Queue()
132 def make_event(self) -> Event:
133 return threading.Event()
135 def make_worker(
136 self, target: Callable[..., None], args: tuple[Any, ...], name: str | None = None
137 ) -> Worker:
138 return threading.Thread(target=target, args=args, name=name)
141class SpawnProcessContext(WorkerContext):
142 """An implementation of `WorkerContext` backed by the `multiprocessing`
143 module, with new processes started by spawning.
144 """
146 def __init__(self) -> None:
147 self._ctx = multiprocessing.get_context("spawn")
149 def make_queue(self) -> Queue[Any]:
150 return self._ctx.Queue()
152 def make_event(self) -> Event:
153 return self._ctx.Event()
155 def make_worker(
156 self, target: Callable[..., None], args: tuple[Any, ...], name: str | None = None
157 ) -> Worker:
158 return self._ctx.Process(target=target, args=args, name=name)
161def _get_from_queue(q: Queue[_T], block: bool = False, timeout: float | None = None) -> _T | None:
162 """Get an object from a queue and return `None` if it is empty.
164 Parameters
165 ----------
166 q : `Queue`
167 Queue to get an object from.
168 block : `bool`
169 Whether to block until an object is available.
170 timeout : `float` or `None`, optional
171 Maximum number of seconds to wait while blocking.
173 Returns
174 -------
175 obj : `object` or `None`
176 Object from the queue, or `None` if it was empty.
177 """
178 try:
179 return q.get(block=block, timeout=timeout)
180 except queue.Empty:
181 return None
184class FatalWorkerError(BaseException):
185 """An exception raised by communicators when one worker (including the
186 supervisor) has caught an exception in order to signal the others to shut
187 down.
188 """
191class _Sentinel(enum.Enum):
192 """Sentinel values used to indicate sequence points or worker shutdown
193 conditions.
194 """
196 NO_MORE_SCAN_REQUESTS = enum.auto()
197 """Sentinel sent from the supervisor to scanners to indicate that there are
198 no more quanta left to be scanned.
199 """
201 NO_MORE_INGEST_REQUESTS = enum.auto()
202 """Sentinel sent from scanners to the ingester to indicate that there are
203 will be no more ingest requests from a particular worker.
204 """
206 NO_MORE_WRITE_REQUESTS = enum.auto()
207 """Sentinel sent from scanners and the supervisor to the writer to
208 indicate that there are will be no more write requests from a particular
209 worker.
210 """
212 WRITE_REPORT = enum.auto()
213 """Sentinel sent from the writer to the supervisor to report that a
214 quantum's provenance was written.
215 """
217 SCANNER_DONE = enum.auto()
218 """Sentinel sent from scanners to the supervisor to report that they are
219 done and shutting down.
220 """
222 INGESTER_DONE = enum.auto()
223 """Sentinel sent from the ingester to the supervisor to report that it is
224 done and shutting down.
225 """
227 WRITER_DONE = enum.auto()
228 """Sentinel sent from the writer to the supervisor to report that it is
229 done and shutting down.
230 """
233@dataclasses.dataclass
234class _WorkerErrorMessage:
235 """An internal worker used to pass information about an error that occurred
236 on a worker back to the supervisor.
238 As a rule, these are unexpected, unrecoverable exceptions.
239 """
241 worker: str
242 """Name of the originating worker."""
244 traceback: str
245 """A logged exception traceback.
247 Note that this is not a `BaseException` subclass that can actually be
248 re-raised on the supervisor; it's just something we can log to make the
249 right traceback appear on the screen. If something silences that printing
250 in favor of its own exception management (pytest!) this information
251 disappears.
252 """
255@dataclasses.dataclass
256class _ScanRequest:
257 """An internal struct passed from the supervisor to the scanners to request
258 a quantum be scanned.
259 """
261 quantum_id: uuid.UUID
262 """ID of the quantum to be scanned."""
265@dataclasses.dataclass
266class _IngestReport:
267 """An internal struct passed from the ingester to the supervisor to report
268 a completed ingest batch.
269 """
271 n_producers: int
272 """Number of producing quanta whose datasets were ingested.
274 We use quanta rather than datasets as the count here because the supervisor
275 knows the total number of quanta in advance but not the total number of
276 datasets to be ingested, so it's a lot easier to attach a denominator
277 and/or progress bar to this number.
278 """
281@dataclasses.dataclass
282class _ProgressLog:
283 """A high-level log message sent from a worker to the supervisor.
285 These are messages that should appear to come from the main
286 'aggregate-graph' logger, not a worker-specific one.
287 """
289 message: str
290 """Log message."""
292 level: int
293 """Log level."""
296@dataclasses.dataclass
297class _CompressionDictionary:
298 """An internal struct used to send the compression dictionary from the
299 writer to the scanners.
300 """
302 data: bytes
303 """The `bytes` representation of a `zstandard.ZstdCompressionDict`.
304 """
307Report: TypeAlias = (
308 ScanReport
309 | _IngestReport
310 | _WorkerErrorMessage
311 | _ProgressLog
312 | Literal[
313 _Sentinel.WRITE_REPORT,
314 _Sentinel.SCANNER_DONE,
315 _Sentinel.INGESTER_DONE,
316 _Sentinel.WRITER_DONE,
317 ]
318)
321def _disable_resources_parallelism() -> None:
322 os.environ["LSST_RESOURCES_NUM_WORKERS"] = "1"
323 os.environ.pop("LSST_RESOURCES_EXECUTOR", None)
324 os.environ["LSST_S3_USE_THREADS"] = "False"
327class SupervisorCommunicator:
328 """A helper object that lets the supervisor direct the other workers.
330 Parameters
331 ----------
332 log : `lsst.utils.logging.LsstLogAdapter`
333 LSST-customized logger.
334 n_scanners : `int`
335 Number of scanner workers.
336 context : `WorkerContext`
337 Abstraction over threading vs. multiprocessing.
338 config : `AggregatorConfig`
339 Configuration for the aggregator.
340 """
342 def __init__(
343 self,
344 log: LsstLogAdapter,
345 n_scanners: int,
346 context: WorkerContext,
347 config: AggregatorConfig,
348 ) -> None:
349 self.config = config
350 self.progress = ProgressManager(log, config)
351 self.n_scanners = n_scanners
352 # The supervisor sends scan requests to scanners on this queue.
353 # When complete, the supervisor sends n_scanners sentinals and each
354 # scanner is careful to only take one before it starts its shutdown.
355 self._scan_requests: Queue[_ScanRequest | Literal[_Sentinel.NO_MORE_SCAN_REQUESTS]] = (
356 context.make_queue()
357 )
358 # The scanners send ingest requests to the ingester on this queue. Each
359 # scanner sends one sentinal when it is done, and the ingester is
360 # careful to wait for n_scanners sentinals to arrive before it starts
361 # its shutdown.
362 self._ingest_requests: Queue[IngestRequest | Literal[_Sentinel.NO_MORE_INGEST_REQUESTS]] = (
363 context.make_queue()
364 )
365 # The scanners send write requests to the writer on this queue (which
366 # will be `None` if we're not writing). The supervisor also sends
367 # write requests for blocked quanta (which we don't scan). Each
368 # scanner and the supervisor send one sentinal when done, and the
369 # writer waits for (n_scanners + 1) sentinals to arrive before it
370 # starts its shutdown.
371 self._write_requests: (
372 Queue[ProvenanceQuantumScanData | Literal[_Sentinel.NO_MORE_WRITE_REQUESTS]] | None
373 ) = context.make_queue() if config.is_writing_provenance else None
374 # All other workers use this queue to send many different kinds of
375 # reports the supervisor. The supervisor waits for a _DONE sentinal
376 # from each worker before it finishes its shutdown.
377 self._reports: Queue[Report] = context.make_queue()
378 # The writer sends the compression dictionary to the scanners on this
379 # queue. It puts n_scanners copies on the queue, and each scanner only
380 # takes one. The compression_dict queue has no sentinal because it is
381 # only used at most once; the supervisor takes responsibility for
382 # clearing it out shutting down.
383 self._compression_dict: Queue[_CompressionDictionary] = context.make_queue()
384 # The supervisor sets this event when it receives an interrupt request
385 # from an exception in the main process (usually KeyboardInterrupt).
386 # Worker communicators check this in their polling loops and raise
387 # FatalWorkerError when they see it set.
388 self._cancel_event: Event = context.make_event()
389 # Track what state we are in closing down, so we can start at the right
390 # point if we're interrupted and __exit__ needs to clean up. Note that
391 # we can't rely on a non-exception __exit__ to do any shutdown work
392 # that might be slow, since a KeyboardInterrupt that occurs when
393 # __exit__ is already running can't be caught inside __exit__.
394 self._sent_no_more_scan_requests = False
395 self._sent_no_more_write_requests = False
396 self._n_scanners_done = 0
397 self._ingester_done = False
398 self._writer_done = self._write_requests is None
400 def wait_for_workers_to_finish(self, already_failing: bool = False) -> None:
401 if not self._sent_no_more_scan_requests:
402 for _ in range(self.n_scanners):
403 self._scan_requests.put(_Sentinel.NO_MORE_SCAN_REQUESTS, block=False)
404 self._sent_no_more_scan_requests = True
405 if not self._sent_no_more_write_requests and self._write_requests is not None:
406 self._write_requests.put(_Sentinel.NO_MORE_WRITE_REQUESTS, block=False)
407 self._sent_no_more_write_requests = True
408 while not (self._ingester_done and self._writer_done and self._n_scanners_done == self.n_scanners):
409 match self._handle_progress_reports(
410 self._reports.get(block=True), already_failing=already_failing
411 ):
412 case None | ScanReport() | _IngestReport():
413 pass
414 case _Sentinel.INGESTER_DONE:
415 self._ingester_done = True
416 self.progress.quantum_ingests.close()
417 case _Sentinel.SCANNER_DONE:
418 self._n_scanners_done += 1
419 self.progress.scans.close()
420 case _Sentinel.WRITER_DONE:
421 self._writer_done = True
422 self.progress.writes.close()
423 case unexpected:
424 raise AssertionError(f"Unexpected message {unexpected!r} to supervisor.")
425 self.log.verbose(
426 "Blocking on reports queue: ingester_done=%s, writer_done=%s, n_scanners_done=%s.",
427 self._ingester_done,
428 self._writer_done,
429 self._n_scanners_done,
430 )
431 while _get_from_queue(self._compression_dict) is not None:
432 self.log.verbose("Flushing compression dict queue.")
433 self.log.verbose("Checking that all queues are empty.")
434 self._expect_empty_queue(self._scan_requests)
435 self._expect_empty_queue(self._ingest_requests)
436 if self._write_requests is not None:
437 self._expect_empty_queue(self._write_requests)
438 self._expect_empty_queue(self._reports)
439 self._expect_empty_queue(self._compression_dict)
441 def __enter__(self) -> Self:
442 _disable_resources_parallelism()
443 self.progress.__enter__()
444 # We make the low-level logger in __enter__ instead of __init__ only
445 # because that's the pattern used by true workers (where it matters).
446 self.log = make_worker_log("supervisor", self.config)
447 return self
449 def __exit__(
450 self,
451 exc_type: type[BaseException] | None,
452 exc_value: BaseException | None,
453 traceback: TracebackType | None,
454 ) -> None:
455 if exc_type is not None:
456 if exc_type is not FatalWorkerError:
457 self.progress.log.critical(f"Caught {exc_type.__name__}; attempting to shut down cleanly.")
458 self._cancel_event.set()
459 self.wait_for_workers_to_finish(already_failing=exc_type is not None)
460 self.progress.__exit__(exc_type, exc_value, traceback)
462 def request_scan(self, quantum_id: uuid.UUID) -> None:
463 """Send a request to the scanners to scan the given quantum.
465 Parameters
466 ----------
467 quantum_id : `uuid.UUID`
468 ID of the quantum to scan.
469 """
470 self._scan_requests.put(_ScanRequest(quantum_id), block=False)
472 def request_write(self, request: ProvenanceQuantumScanData) -> None:
473 """Send a request to the writer to write provenance for the given scan.
475 Parameters
476 ----------
477 request : `ProvenanceQuantumScanData`
478 Information from scanning a quantum (or knowing you don't have to,
479 in the case of blocked quanta).
480 """
481 assert self._write_requests is not None, "Writer should not be used if writing is disabled."
482 self._write_requests.put(request, block=False)
484 def poll(self) -> Iterator[ScanReport]:
485 """Poll for reports from workers while sending scan requests.
487 Yields
488 ------
489 scan_report : `ScanReport`
490 A report from a scanner that a quantum was scanned.
492 Notes
493 -----
494 This iterator blocks until the first scan report is received, and then
495 it continues until the report queue is empty.
496 """
497 block = True
498 msg = _get_from_queue(self._reports, block=block)
499 while msg is not None:
500 match self._handle_progress_reports(msg):
501 case ScanReport() as scan_report:
502 block = False
503 yield scan_report
504 case None:
505 pass
506 case unexpected:
507 raise AssertionError(f"Unexpected message {unexpected!r} to supervisor.")
508 msg = _get_from_queue(self._reports, block=block)
510 def _handle_progress_reports(
511 self, report: Report, already_failing: bool = False
512 ) -> (
513 ScanReport
514 | Literal[
515 _Sentinel.SCANNER_DONE,
516 _Sentinel.INGESTER_DONE,
517 _Sentinel.WRITER_DONE,
518 ]
519 | None
520 ):
521 """Handle reports to the supervisor that can appear at any time, and
522 are typically just updates to the progress we've made.
524 This includes:
526 - exceptions from workers (which raise `FatalWorkerError` here to
527 trigger ``__exit__``);
528 - ingest reports;
529 - write reports;
530 - progress logs.
532 If one of these is handled, `None` is returned; otherwise the original
533 report is returned.
534 """
535 match report:
536 case _WorkerErrorMessage(traceback=traceback, worker=worker):
537 self.progress.log.fatal("Exception raised on %s: \n%s", worker, traceback)
538 if not already_failing:
539 raise FatalWorkerError()
540 case _IngestReport(n_producers=n_producers):
541 self.progress.quantum_ingests.update(n_producers)
542 case _Sentinel.WRITE_REPORT:
543 self.progress.writes.update(1)
544 case _ProgressLog(message=message, level=level):
545 self.progress.log.log(level, "%s [after %0.1fs]", message, self.progress.elapsed_time)
546 case _:
547 return report
548 return None
550 @staticmethod
551 def _expect_empty_queue(queue: Queue[Any]) -> None:
552 """Assert that the given queue is empty."""
553 if (msg := _get_from_queue(queue, block=False, timeout=0)) is not None:
554 raise AssertionError(f"Queue is not empty; found {msg!r}.")
557class WorkerCommunicator:
558 """A base class for non-supervisor workers.
560 Parameters
561 ----------
562 supervisor : `SupervisorCommunicator`
563 Communicator for the supervisor to grab queues and information from.
564 name : `str`
565 Human-readable name for this worker.
567 Notes
568 -----
569 Each worker communicator is constructed in the main process and entered as
570 a context manager on the actual worker process, so attributes that cannot
571 be pickled are constructed in ``__enter__`` instead of ``__init__``.
573 Worker communicators provide access to an `AggregatorConfig` and a logger
574 to their workers. As context managers, they handle exceptions and ensure
575 clean shutdowns, and since most workers need to use a lot of other context
576 managers (for file reading and writing, mostly), they provide an `enter`
577 method to keep every worker from also having to be a context manager just
578 to hold a context manager instance attribute.
580 Worker communicators can also be configured to record and dump profiling
581 information.
582 """
584 def __init__(self, supervisor: SupervisorCommunicator, name: str):
585 self.name = name
586 self.config = supervisor.config
587 self._reports = supervisor._reports
588 self._cancel_event = supervisor._cancel_event
590 def __enter__(self) -> Self:
591 _disable_resources_parallelism()
592 self.log = make_worker_log(self.name, self.config)
593 self.log.verbose("%s has PID %s (parent is %s).", self.name, os.getpid(), os.getppid())
594 self._exit_stack = ExitStack().__enter__()
595 if self.config.n_processes > 1:
596 # Multiprocessing: ignore interrupts so we can shut down cleanly.
597 signal.signal(signal.SIGINT, signal.SIG_IGN)
598 if self.config.worker_profile_dir is not None:
599 # We use time.time because we're interested in wall-clock time,
600 # not just CPU effort, since this is I/O-bound work.
601 self._profiler = cProfile.Profile(timer=time.time)
602 self._profiler.enable()
603 return self
605 def __exit__(
606 self,
607 exc_type: type[BaseException] | None,
608 exc_value: BaseException | None,
609 traceback: TracebackType | None,
610 ) -> bool | None:
611 if self.config.worker_profile_dir is not None and self.config.n_processes > 1:
612 self._profiler.disable()
613 os.makedirs(self.config.worker_profile_dir, exist_ok=True)
614 self._profiler.dump_stats(os.path.join(self.config.worker_profile_dir, f"{self.name}.profile"))
615 if exc_value is not None:
616 assert exc_type is not None, "Should be guaranteed by Python, but MyPy doesn't know that."
617 if exc_type is not FatalWorkerError:
618 self.log.warning("Error raised on this worker.", exc_info=(exc_type, exc_value, traceback))
619 assert exc_type is not None and traceback is not None
620 self._reports.put(
621 _WorkerErrorMessage(
622 self.name,
623 "".join(format_exception(exc_type, exc_value, traceback)),
624 ),
625 block=False,
626 )
627 self.log.debug("Error message sent to supervisor.")
628 else:
629 self.log.warning("Shutting down due to exception raised on another worker.")
630 self._exit_stack.__exit__(exc_type, exc_value, traceback)
631 return True
633 @property
634 def exit_stack(self) -> ExitStack:
635 """A `contextlib.ExitStack` tied to the communicator."""
636 return self._exit_stack
638 def log_progress(self, level: int, message: str) -> None:
639 """Send a high-level log message to the supervisor.
641 Parameters
642 ----------
643 level : `int`
644 Log level. Should be ``VERBOSE`` or higher.
645 message : `str`
646 Log message.
647 """
648 self._reports.put(_ProgressLog(message=message, level=level), block=False)
650 def check_for_cancel(self) -> None:
651 """Check for a cancel signal from the supervisor and raise
652 `FatalWorkerError` if it is present.
653 """
654 if self._cancel_event.is_set():
655 raise FatalWorkerError()
658class ScannerCommunicator(WorkerCommunicator):
659 """A communicator for scanner workers.
661 Parameters
662 ----------
663 supervisor : `SupervisorCommunicator`
664 Communicator for the supervisor to grab queues and information from.
665 scanner_id : `int`
666 Integer ID for this canner.
667 """
669 def __init__(self, supervisor: SupervisorCommunicator, scanner_id: int):
670 super().__init__(supervisor, f"scanner-{scanner_id:03d}")
671 self.scanner_id = scanner_id
672 self._scan_requests = supervisor._scan_requests
673 self._ingest_requests = supervisor._ingest_requests
674 self._write_requests = supervisor._write_requests
675 self._compression_dict = supervisor._compression_dict
676 self._got_no_more_scan_requests: bool = False
677 self._sent_no_more_ingest_requests: bool = False
679 def report_scan(self, msg: ScanReport) -> None:
680 """Report a completed scan to the supervisor.
682 Parameters
683 ----------
684 msg : `ScanReport`
685 Report to send.
686 """
687 self._reports.put(msg, block=False)
689 def request_ingest(self, request: IngestRequest) -> None:
690 """Ask the ingester to ingest a quantum's outputs.
692 Parameters
693 ----------
694 request : `IngestRequest`
695 Description of the datasets to ingest.
697 Notes
698 -----
699 If this request has no datasets, this automatically reports the ingest
700 as complete to the supervisor instead of sending it to the ingester.
701 """
702 if request:
703 self._ingest_requests.put(request, block=False)
704 else:
705 self._reports.put(_IngestReport(1), block=False)
707 def request_write(self, request: ProvenanceQuantumScanData) -> None:
708 """Ask the writer to write provenance for a quantum.
710 Parameters
711 ----------
712 request : `ProvenanceQuantumScanData`
713 Result of scanning a quantum.
714 """
715 assert self._write_requests is not None, "Writer should not be used if writing is disabled."
716 self._write_requests.put(request, block=False)
718 def get_compression_dict(self) -> bytes | None:
719 """Attempt to get the compression dict from the writer.
721 Returns
722 -------
723 data : `bytes` or `None`
724 The `bytes` representation of the compression dictionary, or `None`
725 if the compression dictionary is not yet available.
727 Notes
728 -----
729 A scanner should only call this method before it actually has the
730 compression dict.
731 """
732 if (cdict := _get_from_queue(self._compression_dict)) is not None:
733 return cdict.data
734 return None
736 def poll(self) -> Iterator[uuid.UUID]:
737 """Poll for scan requests to process.
739 Yields
740 ------
741 quantum_id : `uuid.UUID`
742 ID of a new quantum to scan.
744 Notes
745 -----
746 This iterator ends when the supervisor reports that it is done
747 traversing the graph.
748 """
749 while True:
750 self.check_for_cancel()
751 scan_request = _get_from_queue(self._scan_requests, block=True, timeout=self.config.worker_sleep)
752 if scan_request is _Sentinel.NO_MORE_SCAN_REQUESTS:
753 self._got_no_more_scan_requests = True
754 return
755 if scan_request is not None:
756 yield scan_request.quantum_id
758 def __exit__(
759 self,
760 exc_type: type[BaseException] | None,
761 exc_value: BaseException | None,
762 traceback: TracebackType | None,
763 ) -> bool | None:
764 result = super().__exit__(exc_type, exc_value, traceback)
765 self._ingest_requests.put(_Sentinel.NO_MORE_INGEST_REQUESTS, block=False)
766 if self._write_requests is not None:
767 self._write_requests.put(_Sentinel.NO_MORE_WRITE_REQUESTS, block=False)
768 while not self._got_no_more_scan_requests:
769 self.log.debug("Clearing scan request queue (~%d remaining)", self._scan_requests.qsize())
770 if (
771 not self._got_no_more_scan_requests
772 and self._scan_requests.get() is _Sentinel.NO_MORE_SCAN_REQUESTS
773 ):
774 self._got_no_more_scan_requests = True
775 # We let the supervisor clear out the compression dict queue, because
776 # a single scanner can't know if it ever got sent out or not.
777 self.log.verbose("Sending done sentinal.")
778 self._reports.put(_Sentinel.SCANNER_DONE, block=False)
779 return result
782class IngesterCommunicator(WorkerCommunicator):
783 """A communicator for the ingester worker.
785 Parameters
786 ----------
787 supervisor : `SupervisorCommunicator`
788 Communicator for the supervisor to grab queues and information from.
789 """
791 def __init__(self, supervisor: SupervisorCommunicator):
792 super().__init__(supervisor, "ingester")
793 self.n_scanners = supervisor.n_scanners
794 self._ingest_requests = supervisor._ingest_requests
795 self._n_requesters_done = 0
797 def __exit__(
798 self,
799 exc_type: type[BaseException] | None,
800 exc_value: BaseException | None,
801 traceback: TracebackType | None,
802 ) -> bool | None:
803 result = super().__exit__(exc_type, exc_value, traceback)
804 while self._n_requesters_done != self.n_scanners:
805 self.log.debug(
806 "Waiting for %d requesters to be done (currently %d).",
807 self.n_scanners,
808 self._n_requesters_done,
809 )
810 if self._ingest_requests.get(block=True) is _Sentinel.NO_MORE_INGEST_REQUESTS:
811 self._n_requesters_done += 1
812 self.log.verbose("Sending done sentinal.")
813 self._reports.put(_Sentinel.INGESTER_DONE, block=False)
814 return result
816 def report_ingest(self, n_producers: int) -> None:
817 """Report to the supervisor that an ingest batch was completed.
819 Parameters
820 ----------
821 n_producers : `int`
822 Number of producing quanta whose datasets were ingested.
823 """
824 self._reports.put(_IngestReport(n_producers), block=False)
826 def poll(self) -> Iterator[IngestRequest]:
827 """Poll for ingest requests from the scanner workers.
829 Yields
830 ------
831 request : `IngestRequest`
832 A request to ingest datasets produced by a single quantum.
834 Notes
835 -----
836 This iterator ends when all scanners indicate that they are done making
837 ingest requests.
838 """
839 while True:
840 self.check_for_cancel()
841 ingest_request = _get_from_queue(self._ingest_requests, block=True, timeout=_TINY_TIMEOUT)
842 if ingest_request is _Sentinel.NO_MORE_INGEST_REQUESTS:
843 self._n_requesters_done += 1
844 if self._n_requesters_done == self.n_scanners:
845 return
846 else:
847 continue
848 if ingest_request is not None:
849 yield ingest_request
852class WriterCommunicator(WorkerCommunicator):
853 """A communicator for the writer worker.
855 Parameters
856 ----------
857 supervisor : `SupervisorCommunicator`
858 Communicator for the supervisor to grab queues and information from.
859 """
861 def __init__(self, supervisor: SupervisorCommunicator):
862 assert supervisor._write_requests is not None
863 super().__init__(supervisor, "writer")
864 self.n_scanners = supervisor.n_scanners
865 self._write_requests = supervisor._write_requests
866 self._compression_dict = supervisor._compression_dict
867 self._n_requesters = supervisor.n_scanners + 1
868 self._n_requesters_done = 0
869 self._sent_compression_dict = False
871 def __exit__(
872 self,
873 exc_type: type[BaseException] | None,
874 exc_value: BaseException | None,
875 traceback: TracebackType | None,
876 ) -> bool | None:
877 result = super().__exit__(exc_type, exc_value, traceback)
878 if exc_type is None:
879 self.log_progress(logging.INFO, "Provenance quantum graph written successfully.")
880 while self._n_requesters_done != self._n_requesters:
881 self.log.debug(
882 "Waiting for %d requesters to be done (currently %d).",
883 self._n_requesters,
884 self._n_requesters_done,
885 )
886 if self._write_requests.get(block=True) is _Sentinel.NO_MORE_WRITE_REQUESTS:
887 self._n_requesters_done += 1
888 self.log.verbose("Sending done sentinal.")
889 self._reports.put(_Sentinel.WRITER_DONE, block=False)
890 return result
892 def poll(self) -> Iterator[ProvenanceQuantumScanData]:
893 """Poll for writer requests from the scanner workers and supervisor.
895 Yields
896 ------
897 request : `ProvenanceQuantumScanData`
898 The result of a quantum scan.
900 Notes
901 -----
902 This iterator ends when all scanners and the supervisor indicate that
903 they are done making write requests.
904 """
905 while True:
906 self.check_for_cancel()
907 write_request = _get_from_queue(self._write_requests, block=True, timeout=_TINY_TIMEOUT)
908 if write_request is _Sentinel.NO_MORE_WRITE_REQUESTS:
909 self._n_requesters_done += 1
910 if self._n_requesters_done == self._n_requesters:
911 return
912 else:
913 continue
914 if write_request is not None:
915 yield write_request
917 def send_compression_dict(self, cdict_data: bytes) -> None:
918 """Send the compression dictionary to the scanners.
920 Parameters
921 ----------
922 cdict_data : `bytes`
923 The `bytes` representation of the compression dictionary.
924 """
925 self.log.debug("Sending compression dictionary.")
926 for _ in range(self.n_scanners):
927 self._compression_dict.put(_CompressionDictionary(cdict_data), block=False)
928 self._sent_compression_dict = True
930 def report_write(self) -> None:
931 """Report to the supervisor that provenance for a quantum was written
932 to the graph.
933 """
934 self._reports.put(_Sentinel.WRITE_REPORT, block=False)
936 def periodically_check_for_cancel(self, iterable: Iterable[_T], n: int = 100) -> Iterator[_T]:
937 """Iterate while checking for a cancellation signal every ``n``
938 iterations.
940 Parameters
941 ----------
942 iterable : `~collections.abc.Iterable`
943 Object to iterate over.
944 n : `int`
945 Check for cancellation every ``n`` iterations.
947 Returns
948 -------
949 iterator : `~collections.abc.Iterator`
950 Iterator.
951 """
952 i = 0
953 for entry in iterable:
954 yield entry
955 i += 1
956 if i % n == 0:
957 self.check_for_cancel()