Coverage for python / lsst / pipe / base / quantum_graph / aggregator / _communicators.py: 28%
355 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-31 01:07 -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
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_TINY_TIMEOUT = 0.01
68type Queue[T] = "queue.Queue[T]" | "multiprocessing.Queue[T]"
70type Event = threading.Event | multiprocessing.synchronize.Event
72type Worker = threading.Thread | multiprocessing.context.SpawnProcess
75class WorkerContext(ABC):
76 """A simple abstract interface that can be implemented by both threading
77 and multiprocessing.
78 """
80 @abstractmethod
81 def make_queue(self) -> Queue[Any]:
82 """Make an empty queue that can be used to pass objects between
83 workers in this context.
84 """
85 raise NotImplementedError()
87 @abstractmethod
88 def make_event(self) -> Event:
89 """Make an event that can be used to communicate a boolean state change
90 to workers in this context.
91 """
92 raise NotImplementedError()
94 @abstractmethod
95 def make_worker(
96 self, target: Callable[..., None], args: tuple[Any, ...], name: str | None = None
97 ) -> Worker:
98 """Make a worker that runs the given callable.
100 Parameters
101 ----------
102 target : `~collections.abc.Callable`
103 A callable to invoke on the worker.
104 args : `tuple`
105 Positional arguments to pass to the callable.
106 name : `str`, optional
107 Human-readable name for the worker.
109 Returns
110 -------
111 worker : `threading.Thread` or `multiprocessing.Process`
112 Process or thread. Will need to have its ``start`` method called
113 to actually begin.
114 """
115 raise NotImplementedError()
118class ThreadingContext(WorkerContext):
119 """An implementation of `WorkerContext` backed by the `threading`
120 module.
121 """
123 def make_queue(self) -> Queue[Any]:
124 return queue.Queue()
126 def make_event(self) -> Event:
127 return threading.Event()
129 def make_worker(
130 self, target: Callable[..., None], args: tuple[Any, ...], name: str | None = None
131 ) -> Worker:
132 return threading.Thread(target=target, args=args, name=name)
135class SpawnProcessContext(WorkerContext):
136 """An implementation of `WorkerContext` backed by the `multiprocessing`
137 module, with new processes started by spawning.
138 """
140 def __init__(self) -> None:
141 self._ctx = multiprocessing.get_context("spawn")
143 def make_queue(self) -> Queue[Any]:
144 return self._ctx.Queue()
146 def make_event(self) -> Event:
147 return self._ctx.Event()
149 def make_worker(
150 self, target: Callable[..., None], args: tuple[Any, ...], name: str | None = None
151 ) -> Worker:
152 return self._ctx.Process(target=target, args=args, name=name)
155def _get_from_queue[T](q: Queue[T], block: bool = False, timeout: float | None = None) -> T | None:
156 """Get an object from a queue and return `None` if it is empty.
158 Parameters
159 ----------
160 q : `Queue`
161 Queue to get an object from.
162 block : `bool`
163 Whether to block until an object is available.
164 timeout : `float` or `None`, optional
165 Maximum number of seconds to wait while blocking.
167 Returns
168 -------
169 obj : `object` or `None`
170 Object from the queue, or `None` if it was empty.
171 """
172 try:
173 return q.get(block=block, timeout=timeout)
174 except queue.Empty:
175 return None
178class FatalWorkerError(BaseException):
179 """An exception raised by communicators when one worker (including the
180 supervisor) has caught an exception in order to signal the others to shut
181 down.
182 """
185class _Sentinel(enum.Enum):
186 """Sentinel values used to indicate sequence points or worker shutdown
187 conditions.
188 """
190 NO_MORE_SCAN_REQUESTS = enum.auto()
191 """Sentinel sent from the supervisor to scanners to indicate that there are
192 no more quanta left to be scanned.
193 """
195 NO_MORE_INGEST_REQUESTS = enum.auto()
196 """Sentinel sent from scanners to the ingester to indicate that there are
197 will be no more ingest requests from a particular worker.
198 """
200 NO_MORE_WRITE_REQUESTS = enum.auto()
201 """Sentinel sent from scanners and the supervisor to the writer to
202 indicate that there are will be no more write requests from a particular
203 worker.
204 """
206 WRITE_REPORT = enum.auto()
207 """Sentinel sent from the writer to the supervisor to report that a
208 quantum's provenance was written.
209 """
211 SCANNER_DONE = enum.auto()
212 """Sentinel sent from scanners to the supervisor to report that they are
213 done and shutting down.
214 """
216 INGESTER_DONE = enum.auto()
217 """Sentinel sent from the ingester to the supervisor to report that it is
218 done and shutting down.
219 """
221 WRITER_DONE = enum.auto()
222 """Sentinel sent from the writer to the supervisor to report that it is
223 done and shutting down.
224 """
227@dataclasses.dataclass
228class _WorkerErrorMessage:
229 """An internal worker used to pass information about an error that occurred
230 on a worker back to the supervisor.
232 As a rule, these are unexpected, unrecoverable exceptions.
233 """
235 worker: str
236 """Name of the originating worker."""
238 traceback: str
239 """A logged exception traceback.
241 Note that this is not a `BaseException` subclass that can actually be
242 re-raised on the supervisor; it's just something we can log to make the
243 right traceback appear on the screen. If something silences that printing
244 in favor of its own exception management (pytest!) this information
245 disappears.
246 """
249@dataclasses.dataclass
250class _ScanRequest:
251 """An internal struct passed from the supervisor to the scanners to request
252 a quantum be scanned.
253 """
255 quantum_id: uuid.UUID
256 """ID of the quantum to be scanned."""
259@dataclasses.dataclass
260class _IngestReport:
261 """An internal struct passed from the ingester to the supervisor to report
262 a completed ingest batch.
263 """
265 n_producers: int
266 """Number of producing quanta whose datasets were ingested.
268 We use quanta rather than datasets as the count here because the supervisor
269 knows the total number of quanta in advance but not the total number of
270 datasets to be ingested, so it's a lot easier to attach a denominator
271 and/or progress bar to this number.
272 """
275@dataclasses.dataclass
276class _ProgressLog:
277 """A high-level log message sent from a worker to the supervisor.
279 These are messages that should appear to come from the main
280 'aggregate-graph' logger, not a worker-specific one.
281 """
283 message: str
284 """Log message."""
286 level: int
287 """Log level."""
290@dataclasses.dataclass
291class _CompressionDictionary:
292 """An internal struct used to send the compression dictionary from the
293 writer to the scanners.
294 """
296 data: bytes
297 """The `bytes` representation of a `zstandard.ZstdCompressionDict`.
298 """
301type Report = (
302 ScanReport
303 | _IngestReport
304 | _WorkerErrorMessage
305 | _ProgressLog
306 | Literal[
307 _Sentinel.WRITE_REPORT,
308 _Sentinel.SCANNER_DONE,
309 _Sentinel.INGESTER_DONE,
310 _Sentinel.WRITER_DONE,
311 ]
312)
315def _disable_resources_parallelism() -> None:
316 os.environ["LSST_RESOURCES_NUM_WORKERS"] = "1"
317 os.environ.pop("LSST_RESOURCES_EXECUTOR", None)
318 os.environ["LSST_S3_USE_THREADS"] = "False"
321class SupervisorCommunicator:
322 """A helper object that lets the supervisor direct the other workers.
324 Parameters
325 ----------
326 log : `lsst.utils.logging.LsstLogAdapter`
327 LSST-customized logger.
328 n_scanners : `int`
329 Number of scanner workers.
330 context : `WorkerContext`
331 Abstraction over threading vs. multiprocessing.
332 config : `AggregatorConfig`
333 Configuration for the aggregator.
334 """
336 def __init__(
337 self,
338 log: LsstLogAdapter,
339 n_scanners: int,
340 context: WorkerContext,
341 config: AggregatorConfig,
342 ) -> None:
343 self.config = config
344 self.progress = ProgressManager(log, config)
345 self.n_scanners = n_scanners
346 # The supervisor sends scan requests to scanners on this queue.
347 # When complete, the supervisor sends n_scanners sentinals and each
348 # scanner is careful to only take one before it starts its shutdown.
349 self._scan_requests: Queue[_ScanRequest | Literal[_Sentinel.NO_MORE_SCAN_REQUESTS]] = (
350 context.make_queue()
351 )
352 # The scanners send ingest requests to the ingester on this queue. Each
353 # scanner sends one sentinal when it is done, and the ingester is
354 # careful to wait for n_scanners sentinals to arrive before it starts
355 # its shutdown.
356 self._ingest_requests: Queue[IngestRequest | Literal[_Sentinel.NO_MORE_INGEST_REQUESTS]] = (
357 context.make_queue()
358 )
359 # The scanners send write requests to the writer on this queue (which
360 # will be `None` if we're not writing). The supervisor also sends
361 # write requests for blocked quanta (which we don't scan). Each
362 # scanner and the supervisor send one sentinal when done, and the
363 # writer waits for (n_scanners + 1) sentinals to arrive before it
364 # starts its shutdown.
365 self._write_requests: (
366 Queue[ProvenanceQuantumScanData | Literal[_Sentinel.NO_MORE_WRITE_REQUESTS]] | None
367 ) = context.make_queue() if config.is_writing_provenance else None
368 # All other workers use this queue to send many different kinds of
369 # reports the supervisor. The supervisor waits for a _DONE sentinal
370 # from each worker before it finishes its shutdown.
371 self._reports: Queue[Report] = context.make_queue()
372 # The writer sends the compression dictionary to the scanners on this
373 # queue. It puts n_scanners copies on the queue, and each scanner only
374 # takes one. The compression_dict queue has no sentinal because it is
375 # only used at most once; the supervisor takes responsibility for
376 # clearing it out shutting down.
377 self._compression_dict: Queue[_CompressionDictionary] = context.make_queue()
378 # The supervisor sets this event when it receives an interrupt request
379 # from an exception in the main process (usually KeyboardInterrupt).
380 # Worker communicators check this in their polling loops and raise
381 # FatalWorkerError when they see it set.
382 self._cancel_event: Event = context.make_event()
383 # Track what state we are in closing down, so we can start at the right
384 # point if we're interrupted and __exit__ needs to clean up. Note that
385 # we can't rely on a non-exception __exit__ to do any shutdown work
386 # that might be slow, since a KeyboardInterrupt that occurs when
387 # __exit__ is already running can't be caught inside __exit__.
388 self._sent_no_more_scan_requests = False
389 self._sent_no_more_write_requests = False
390 self._n_scanners_done = 0
391 self._ingester_done = False
392 self._writer_done = self._write_requests is None
394 def wait_for_workers_to_finish(self, already_failing: bool = False) -> None:
395 if not self._sent_no_more_scan_requests:
396 for _ in range(self.n_scanners):
397 self._scan_requests.put(_Sentinel.NO_MORE_SCAN_REQUESTS, block=False)
398 self._sent_no_more_scan_requests = True
399 if not self._sent_no_more_write_requests and self._write_requests is not None:
400 self._write_requests.put(_Sentinel.NO_MORE_WRITE_REQUESTS, block=False)
401 self._sent_no_more_write_requests = True
402 while not (self._ingester_done and self._writer_done and self._n_scanners_done == self.n_scanners):
403 match self._handle_progress_reports(
404 self._reports.get(block=True), already_failing=already_failing
405 ):
406 case None | ScanReport() | _IngestReport():
407 pass
408 case _Sentinel.INGESTER_DONE:
409 self._ingester_done = True
410 self.progress.quantum_ingests.close()
411 case _Sentinel.SCANNER_DONE:
412 self._n_scanners_done += 1
413 self.progress.scans.close()
414 case _Sentinel.WRITER_DONE:
415 self._writer_done = True
416 self.progress.writes.close()
417 case unexpected:
418 raise AssertionError(f"Unexpected message {unexpected!r} to supervisor.")
419 self.log.verbose(
420 "Blocking on reports queue: ingester_done=%s, writer_done=%s, n_scanners_done=%s.",
421 self._ingester_done,
422 self._writer_done,
423 self._n_scanners_done,
424 )
425 while _get_from_queue(self._compression_dict) is not None:
426 self.log.verbose("Flushing compression dict queue.")
427 self.log.verbose("Checking that all queues are empty.")
428 self._expect_empty_queue(self._scan_requests)
429 self._expect_empty_queue(self._ingest_requests)
430 if self._write_requests is not None:
431 self._expect_empty_queue(self._write_requests)
432 self._expect_empty_queue(self._reports)
433 self._expect_empty_queue(self._compression_dict)
435 def __enter__(self) -> Self:
436 _disable_resources_parallelism()
437 self.progress.__enter__()
438 # We make the low-level logger in __enter__ instead of __init__ only
439 # because that's the pattern used by true workers (where it matters).
440 self.log = make_worker_log("supervisor", self.config)
441 return self
443 def __exit__(
444 self,
445 exc_type: type[BaseException] | None,
446 exc_value: BaseException | None,
447 traceback: TracebackType | None,
448 ) -> None:
449 if exc_type is not None:
450 if exc_type is not FatalWorkerError:
451 self.progress.log.critical(f"Caught {exc_type.__name__}; attempting to shut down cleanly.")
452 self._cancel_event.set()
453 self.wait_for_workers_to_finish(already_failing=exc_type is not None)
454 self.progress.__exit__(exc_type, exc_value, traceback)
456 def request_scan(self, quantum_id: uuid.UUID) -> None:
457 """Send a request to the scanners to scan the given quantum.
459 Parameters
460 ----------
461 quantum_id : `uuid.UUID`
462 ID of the quantum to scan.
463 """
464 self._scan_requests.put(_ScanRequest(quantum_id), block=False)
466 def request_write(self, request: ProvenanceQuantumScanData) -> None:
467 """Send a request to the writer to write provenance for the given scan.
469 Parameters
470 ----------
471 request : `ProvenanceQuantumScanData`
472 Information from scanning a quantum (or knowing you don't have to,
473 in the case of blocked quanta).
474 """
475 assert self._write_requests is not None, "Writer should not be used if writing is disabled."
476 self._write_requests.put(request, block=False)
478 def poll(self) -> Iterator[ScanReport]:
479 """Poll for reports from workers while sending scan requests.
481 Yields
482 ------
483 scan_report : `ScanReport`
484 A report from a scanner that a quantum was scanned.
486 Notes
487 -----
488 This iterator blocks until the first scan report is received, and then
489 it continues until the report queue is empty.
490 """
491 block = True
492 msg = _get_from_queue(self._reports, block=block)
493 while msg is not None:
494 match self._handle_progress_reports(msg):
495 case ScanReport() as scan_report:
496 block = False
497 yield scan_report
498 case None:
499 pass
500 case unexpected:
501 raise AssertionError(f"Unexpected message {unexpected!r} to supervisor.")
502 msg = _get_from_queue(self._reports, block=block)
504 def _handle_progress_reports(
505 self, report: Report, already_failing: bool = False
506 ) -> (
507 ScanReport
508 | Literal[
509 _Sentinel.SCANNER_DONE,
510 _Sentinel.INGESTER_DONE,
511 _Sentinel.WRITER_DONE,
512 ]
513 | None
514 ):
515 """Handle reports to the supervisor that can appear at any time, and
516 are typically just updates to the progress we've made.
518 This includes:
520 - exceptions from workers (which raise `FatalWorkerError` here to
521 trigger ``__exit__``);
522 - ingest reports;
523 - write reports;
524 - progress logs.
526 If one of these is handled, `None` is returned; otherwise the original
527 report is returned.
528 """
529 match report:
530 case _WorkerErrorMessage(traceback=traceback, worker=worker):
531 self.progress.log.fatal("Exception raised on %s: \n%s", worker, traceback)
532 if not already_failing:
533 raise FatalWorkerError()
534 case _IngestReport(n_producers=n_producers):
535 self.progress.quantum_ingests.update(n_producers)
536 case _Sentinel.WRITE_REPORT:
537 self.progress.writes.update(1)
538 case _ProgressLog(message=message, level=level):
539 self.progress.log.log(level, "%s [after %0.1fs]", message, self.progress.elapsed_time)
540 case _:
541 return report
542 return None
544 @staticmethod
545 def _expect_empty_queue(queue: Queue[Any]) -> None:
546 """Assert that the given queue is empty."""
547 if (msg := _get_from_queue(queue, block=False, timeout=0)) is not None:
548 raise AssertionError(f"Queue is not empty; found {msg!r}.")
551class WorkerCommunicator:
552 """A base class for non-supervisor workers.
554 Parameters
555 ----------
556 supervisor : `SupervisorCommunicator`
557 Communicator for the supervisor to grab queues and information from.
558 name : `str`
559 Human-readable name for this worker.
561 Notes
562 -----
563 Each worker communicator is constructed in the main process and entered as
564 a context manager on the actual worker process, so attributes that cannot
565 be pickled are constructed in ``__enter__`` instead of ``__init__``.
567 Worker communicators provide access to an `AggregatorConfig` and a logger
568 to their workers. As context managers, they handle exceptions and ensure
569 clean shutdowns, and since most workers need to use a lot of other context
570 managers (for file reading and writing, mostly), they provide an `enter`
571 method to keep every worker from also having to be a context manager just
572 to hold a context manager instance attribute.
574 Worker communicators can also be configured to record and dump profiling
575 information.
576 """
578 def __init__(self, supervisor: SupervisorCommunicator, name: str):
579 self.name = name
580 self.config = supervisor.config
581 self._reports = supervisor._reports
582 self._cancel_event = supervisor._cancel_event
584 def __enter__(self) -> Self:
585 _disable_resources_parallelism()
586 self.log = make_worker_log(self.name, self.config)
587 self.log.verbose("%s has PID %s (parent is %s).", self.name, os.getpid(), os.getppid())
588 self._exit_stack = ExitStack().__enter__()
589 if self.config.n_processes > 1:
590 # Multiprocessing: ignore interrupts so we can shut down cleanly.
591 signal.signal(signal.SIGINT, signal.SIG_IGN)
592 if self.config.worker_profile_dir is not None:
593 # We use time.time because we're interested in wall-clock time,
594 # not just CPU effort, since this is I/O-bound work.
595 self._profiler = cProfile.Profile(timer=time.time)
596 self._profiler.enable()
597 return self
599 def __exit__(
600 self,
601 exc_type: type[BaseException] | None,
602 exc_value: BaseException | None,
603 traceback: TracebackType | None,
604 ) -> bool | None:
605 if self.config.worker_profile_dir is not None and self.config.n_processes > 1:
606 self._profiler.disable()
607 os.makedirs(self.config.worker_profile_dir, exist_ok=True)
608 self._profiler.dump_stats(os.path.join(self.config.worker_profile_dir, f"{self.name}.profile"))
609 if exc_value is not None:
610 assert exc_type is not None, "Should be guaranteed by Python, but MyPy doesn't know that."
611 if exc_type is not FatalWorkerError:
612 self.log.warning("Error raised on this worker.", exc_info=(exc_type, exc_value, traceback))
613 assert exc_type is not None and traceback is not None
614 self._reports.put(
615 _WorkerErrorMessage(
616 self.name,
617 "".join(format_exception(exc_type, exc_value, traceback)),
618 ),
619 block=False,
620 )
621 self.log.debug("Error message sent to supervisor.")
622 else:
623 self.log.warning("Shutting down due to exception raised on another worker.")
624 self._exit_stack.__exit__(exc_type, exc_value, traceback)
625 return True
627 @property
628 def exit_stack(self) -> ExitStack:
629 """A `contextlib.ExitStack` tied to the communicator."""
630 return self._exit_stack
632 def log_progress(self, level: int, message: str) -> None:
633 """Send a high-level log message to the supervisor.
635 Parameters
636 ----------
637 level : `int`
638 Log level. Should be ``VERBOSE`` or higher.
639 message : `str`
640 Log message.
641 """
642 self._reports.put(_ProgressLog(message=message, level=level), block=False)
644 def check_for_cancel(self) -> None:
645 """Check for a cancel signal from the supervisor and raise
646 `FatalWorkerError` if it is present.
647 """
648 if self._cancel_event.is_set():
649 raise FatalWorkerError()
652class ScannerCommunicator(WorkerCommunicator):
653 """A communicator for scanner workers.
655 Parameters
656 ----------
657 supervisor : `SupervisorCommunicator`
658 Communicator for the supervisor to grab queues and information from.
659 scanner_id : `int`
660 Integer ID for this canner.
661 """
663 def __init__(self, supervisor: SupervisorCommunicator, scanner_id: int):
664 super().__init__(supervisor, f"scanner-{scanner_id:03d}")
665 self.scanner_id = scanner_id
666 self._scan_requests = supervisor._scan_requests
667 self._ingest_requests = supervisor._ingest_requests
668 self._write_requests = supervisor._write_requests
669 self._compression_dict = supervisor._compression_dict
670 self._got_no_more_scan_requests: bool = False
671 self._sent_no_more_ingest_requests: bool = False
673 def report_scan(self, msg: ScanReport) -> None:
674 """Report a completed scan to the supervisor.
676 Parameters
677 ----------
678 msg : `ScanReport`
679 Report to send.
680 """
681 self._reports.put(msg, block=False)
683 def request_ingest(self, request: IngestRequest) -> None:
684 """Ask the ingester to ingest a quantum's outputs.
686 Parameters
687 ----------
688 request : `IngestRequest`
689 Description of the datasets to ingest.
691 Notes
692 -----
693 If this request has no datasets, this automatically reports the ingest
694 as complete to the supervisor instead of sending it to the ingester.
695 """
696 if request:
697 self._ingest_requests.put(request, block=False)
698 else:
699 self._reports.put(_IngestReport(1), block=False)
701 def request_write(self, request: ProvenanceQuantumScanData) -> None:
702 """Ask the writer to write provenance for a quantum.
704 Parameters
705 ----------
706 request : `ProvenanceQuantumScanData`
707 Result of scanning a quantum.
708 """
709 assert self._write_requests is not None, "Writer should not be used if writing is disabled."
710 self._write_requests.put(request, block=False)
712 def get_compression_dict(self) -> bytes | None:
713 """Attempt to get the compression dict from the writer.
715 Returns
716 -------
717 data : `bytes` or `None`
718 The `bytes` representation of the compression dictionary, or `None`
719 if the compression dictionary is not yet available.
721 Notes
722 -----
723 A scanner should only call this method before it actually has the
724 compression dict.
725 """
726 if (cdict := _get_from_queue(self._compression_dict)) is not None:
727 return cdict.data
728 return None
730 def poll(self) -> Iterator[uuid.UUID]:
731 """Poll for scan requests to process.
733 Yields
734 ------
735 quantum_id : `uuid.UUID`
736 ID of a new quantum to scan.
738 Notes
739 -----
740 This iterator ends when the supervisor reports that it is done
741 traversing the graph.
742 """
743 while True:
744 self.check_for_cancel()
745 scan_request = _get_from_queue(self._scan_requests, block=True, timeout=self.config.worker_sleep)
746 if scan_request is _Sentinel.NO_MORE_SCAN_REQUESTS:
747 self._got_no_more_scan_requests = True
748 return
749 if scan_request is not None:
750 yield scan_request.quantum_id
752 def __exit__(
753 self,
754 exc_type: type[BaseException] | None,
755 exc_value: BaseException | None,
756 traceback: TracebackType | None,
757 ) -> bool | None:
758 result = super().__exit__(exc_type, exc_value, traceback)
759 self._ingest_requests.put(_Sentinel.NO_MORE_INGEST_REQUESTS, block=False)
760 if self._write_requests is not None:
761 self._write_requests.put(_Sentinel.NO_MORE_WRITE_REQUESTS, block=False)
762 while not self._got_no_more_scan_requests:
763 self.log.debug("Clearing scan request queue (~%d remaining)", self._scan_requests.qsize())
764 if (
765 not self._got_no_more_scan_requests
766 and self._scan_requests.get() is _Sentinel.NO_MORE_SCAN_REQUESTS
767 ):
768 self._got_no_more_scan_requests = True
769 # We let the supervisor clear out the compression dict queue, because
770 # a single scanner can't know if it ever got sent out or not.
771 self.log.verbose("Sending done sentinal.")
772 self._reports.put(_Sentinel.SCANNER_DONE, block=False)
773 return result
776class IngesterCommunicator(WorkerCommunicator):
777 """A communicator for the ingester worker.
779 Parameters
780 ----------
781 supervisor : `SupervisorCommunicator`
782 Communicator for the supervisor to grab queues and information from.
783 """
785 def __init__(self, supervisor: SupervisorCommunicator):
786 super().__init__(supervisor, "ingester")
787 self.n_scanners = supervisor.n_scanners
788 self._ingest_requests = supervisor._ingest_requests
789 self._n_requesters_done = 0
791 def __exit__(
792 self,
793 exc_type: type[BaseException] | None,
794 exc_value: BaseException | None,
795 traceback: TracebackType | None,
796 ) -> bool | None:
797 result = super().__exit__(exc_type, exc_value, traceback)
798 while self._n_requesters_done != self.n_scanners:
799 self.log.debug(
800 "Waiting for %d requesters to be done (currently %d).",
801 self.n_scanners,
802 self._n_requesters_done,
803 )
804 if self._ingest_requests.get(block=True) is _Sentinel.NO_MORE_INGEST_REQUESTS:
805 self._n_requesters_done += 1
806 self.log.verbose("Sending done sentinal.")
807 self._reports.put(_Sentinel.INGESTER_DONE, block=False)
808 return result
810 def report_ingest(self, n_producers: int) -> None:
811 """Report to the supervisor that an ingest batch was completed.
813 Parameters
814 ----------
815 n_producers : `int`
816 Number of producing quanta whose datasets were ingested.
817 """
818 self._reports.put(_IngestReport(n_producers), block=False)
820 def poll(self) -> Iterator[IngestRequest]:
821 """Poll for ingest requests from the scanner workers.
823 Yields
824 ------
825 request : `IngestRequest`
826 A request to ingest datasets produced by a single quantum.
828 Notes
829 -----
830 This iterator ends when all scanners indicate that they are done making
831 ingest requests.
832 """
833 while True:
834 self.check_for_cancel()
835 ingest_request = _get_from_queue(self._ingest_requests, block=True, timeout=_TINY_TIMEOUT)
836 if ingest_request is _Sentinel.NO_MORE_INGEST_REQUESTS:
837 self._n_requesters_done += 1
838 if self._n_requesters_done == self.n_scanners:
839 return
840 else:
841 continue
842 if ingest_request is not None:
843 yield ingest_request
846class WriterCommunicator(WorkerCommunicator):
847 """A communicator for the writer worker.
849 Parameters
850 ----------
851 supervisor : `SupervisorCommunicator`
852 Communicator for the supervisor to grab queues and information from.
853 """
855 def __init__(self, supervisor: SupervisorCommunicator):
856 assert supervisor._write_requests is not None
857 super().__init__(supervisor, "writer")
858 self.n_scanners = supervisor.n_scanners
859 self._write_requests = supervisor._write_requests
860 self._compression_dict = supervisor._compression_dict
861 self._n_requesters = supervisor.n_scanners + 1
862 self._n_requesters_done = 0
863 self._sent_compression_dict = False
865 def __exit__(
866 self,
867 exc_type: type[BaseException] | None,
868 exc_value: BaseException | None,
869 traceback: TracebackType | None,
870 ) -> bool | None:
871 result = super().__exit__(exc_type, exc_value, traceback)
872 if exc_type is None:
873 self.log_progress(logging.INFO, "Provenance quantum graph written successfully.")
874 while self._n_requesters_done != self._n_requesters:
875 self.log.debug(
876 "Waiting for %d requesters to be done (currently %d).",
877 self._n_requesters,
878 self._n_requesters_done,
879 )
880 if self._write_requests.get(block=True) is _Sentinel.NO_MORE_WRITE_REQUESTS:
881 self._n_requesters_done += 1
882 self.log.verbose("Sending done sentinal.")
883 self._reports.put(_Sentinel.WRITER_DONE, block=False)
884 return result
886 def poll(self) -> Iterator[ProvenanceQuantumScanData]:
887 """Poll for writer requests from the scanner workers and supervisor.
889 Yields
890 ------
891 request : `ProvenanceQuantumScanData`
892 The result of a quantum scan.
894 Notes
895 -----
896 This iterator ends when all scanners and the supervisor indicate that
897 they are done making write requests.
898 """
899 while True:
900 self.check_for_cancel()
901 write_request = _get_from_queue(self._write_requests, block=True, timeout=_TINY_TIMEOUT)
902 if write_request is _Sentinel.NO_MORE_WRITE_REQUESTS:
903 self._n_requesters_done += 1
904 if self._n_requesters_done == self._n_requesters:
905 return
906 else:
907 continue
908 if write_request is not None:
909 yield write_request
911 def send_compression_dict(self, cdict_data: bytes) -> None:
912 """Send the compression dictionary to the scanners.
914 Parameters
915 ----------
916 cdict_data : `bytes`
917 The `bytes` representation of the compression dictionary.
918 """
919 self.log.debug("Sending compression dictionary.")
920 for _ in range(self.n_scanners):
921 self._compression_dict.put(_CompressionDictionary(cdict_data), block=False)
922 self._sent_compression_dict = True
924 def report_write(self) -> None:
925 """Report to the supervisor that provenance for a quantum was written
926 to the graph.
927 """
928 self._reports.put(_Sentinel.WRITE_REPORT, block=False)
930 def periodically_check_for_cancel[T](self, iterable: Iterable[T], n: int = 100) -> Iterator[T]:
931 """Iterate while checking for a cancellation signal every ``n``
932 iterations.
934 Parameters
935 ----------
936 iterable : `~collections.abc.Iterable`
937 Object to iterate over.
938 n : `int`
939 Check for cancellation every ``n`` iterations.
941 Returns
942 -------
943 iterator : `~collections.abc.Iterator`
944 Iterator.
945 """
946 i = 0
947 for entry in iterable:
948 yield entry
949 i += 1
950 if i % n == 0:
951 self.check_for_cancel()