Coverage for python / lsst / pipe / base / quantum_graph / aggregator / _communicators.py: 29%
350 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-15 06:32 -0300
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)
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.output_path is not None 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 self.progress.__enter__()
437 # We make the low-level logger in __enter__ instead of __init__ only
438 # because that's the pattern used by true workers (where it matters).
439 self.log = make_worker_log("supervisor", self.config)
440 return self
442 def __exit__(
443 self,
444 exc_type: type[BaseException] | None,
445 exc_value: BaseException | None,
446 traceback: TracebackType | None,
447 ) -> None:
448 if exc_type is not None:
449 if exc_type is not FatalWorkerError:
450 self.progress.log.critical(f"Caught {exc_type.__name__}; attempting to shut down cleanly.")
451 self._cancel_event.set()
452 self.wait_for_workers_to_finish(already_failing=exc_type is not None)
453 self.progress.__exit__(exc_type, exc_value, traceback)
455 def request_scan(self, quantum_id: uuid.UUID) -> None:
456 """Send a request to the scanners to scan the given quantum.
458 Parameters
459 ----------
460 quantum_id : `uuid.UUID`
461 ID of the quantum to scan.
462 """
463 self._scan_requests.put(_ScanRequest(quantum_id), block=False)
465 def request_write(self, request: ProvenanceQuantumScanData) -> None:
466 """Send a request to the writer to write provenance for the given scan.
468 Parameters
469 ----------
470 request : `ProvenanceQuantumScanData`
471 Information from scanning a quantum (or knowing you don't have to,
472 in the case of blocked quanta).
473 """
474 assert self._write_requests is not None, "Writer should not be used if writing is disabled."
475 self._write_requests.put(request, block=False)
477 def poll(self) -> Iterator[ScanReport]:
478 """Poll for reports from workers while sending scan requests.
480 Yields
481 ------
482 scan_report : `ScanReport`
483 A report from a scanner that a quantum was scanned.
485 Notes
486 -----
487 This iterator blocks until the first scan report is received, and then
488 it continues until the report queue is empty.
489 """
490 block = True
491 msg = _get_from_queue(self._reports, block=block)
492 while msg is not None:
493 match self._handle_progress_reports(msg):
494 case ScanReport() as scan_report:
495 block = False
496 yield scan_report
497 case None:
498 pass
499 case unexpected:
500 raise AssertionError(f"Unexpected message {unexpected!r} to supervisor.")
501 msg = _get_from_queue(self._reports, block=block)
503 def _handle_progress_reports(
504 self, report: Report, already_failing: bool = False
505 ) -> (
506 ScanReport
507 | Literal[
508 _Sentinel.SCANNER_DONE,
509 _Sentinel.INGESTER_DONE,
510 _Sentinel.WRITER_DONE,
511 ]
512 | None
513 ):
514 """Handle reports to the supervisor that can appear at any time, and
515 are typically just updates to the progress we've made.
517 This includes:
519 - exceptions from workers (which raise `FatalWorkerError` here to
520 trigger ``__exit__``);
521 - ingest reports;
522 - write reports;
523 - progress logs.
525 If one of these is handled, `None` is returned; otherwise the original
526 report is returned.
527 """
528 match report:
529 case _WorkerErrorMessage(traceback=traceback, worker=worker):
530 self.progress.log.fatal("Exception raised on %s: \n%s", worker, traceback)
531 if not already_failing:
532 raise FatalWorkerError()
533 case _IngestReport(n_producers=n_producers):
534 self.progress.quantum_ingests.update(n_producers)
535 case _Sentinel.WRITE_REPORT:
536 self.progress.writes.update(1)
537 case _ProgressLog(message=message, level=level):
538 self.progress.log.log(level, "%s [after %0.1fs]", message, self.progress.elapsed_time)
539 case _:
540 return report
541 return None
543 @staticmethod
544 def _expect_empty_queue(queue: Queue[Any]) -> None:
545 """Assert that the given queue is empty."""
546 if (msg := _get_from_queue(queue, block=False, timeout=0)) is not None:
547 raise AssertionError(f"Queue is not empty; found {msg!r}.")
550class WorkerCommunicator:
551 """A base class for non-supervisor workers.
553 Parameters
554 ----------
555 supervisor : `SupervisorCommunicator`
556 Communicator for the supervisor to grab queues and information from.
557 name : `str`
558 Human-readable name for this worker.
560 Notes
561 -----
562 Each worker communicator is constructed in the main process and entered as
563 a context manager on the actual worker process, so attributes that cannot
564 be pickled are constructed in ``__enter__`` instead of ``__init__``.
566 Worker communicators provide access to an `AggregatorConfig` and a logger
567 to their workers. As context managers, they handle exceptions and ensure
568 clean shutdowns, and since most workers need to use a lot of other context
569 managers (for file reading and writing, mostly), they provide an `enter`
570 method to keep every worker from also having to be a context manager just
571 to hold a context manager instance attribute.
573 Worker communicators can also be configured to record and dump profiling
574 information.
575 """
577 def __init__(self, supervisor: SupervisorCommunicator, name: str):
578 self.name = name
579 self.config = supervisor.config
580 self._reports = supervisor._reports
581 self._cancel_event = supervisor._cancel_event
583 def __enter__(self) -> Self:
584 self.log = make_worker_log(self.name, self.config)
585 self.log.verbose("%s has PID %s (parent is %s).", self.name, os.getpid(), os.getppid())
586 self._exit_stack = ExitStack().__enter__()
587 if self.config.n_processes > 1:
588 # Multiprocessing: ignore interrupts so we can shut down cleanly.
589 signal.signal(signal.SIGINT, signal.SIG_IGN)
590 if self.config.worker_profile_dir is not None:
591 # We use time.time because we're interested in wall-clock time,
592 # not just CPU effort, since this is I/O-bound work.
593 self._profiler = cProfile.Profile(timer=time.time)
594 self._profiler.enable()
595 return self
597 def __exit__(
598 self,
599 exc_type: type[BaseException] | None,
600 exc_value: BaseException | None,
601 traceback: TracebackType | None,
602 ) -> bool | None:
603 if self.config.worker_profile_dir is not None and self.config.n_processes > 1:
604 self._profiler.disable()
605 os.makedirs(self.config.worker_profile_dir, exist_ok=True)
606 self._profiler.dump_stats(os.path.join(self.config.worker_profile_dir, f"{self.name}.profile"))
607 if exc_value is not None:
608 assert exc_type is not None, "Should be guaranteed by Python, but MyPy doesn't know that."
609 if exc_type is not FatalWorkerError:
610 self.log.warning("Error raised on this worker.", exc_info=(exc_type, exc_value, traceback))
611 assert exc_type is not None and traceback is not None
612 self._reports.put(
613 _WorkerErrorMessage(
614 self.name,
615 "".join(format_exception(exc_type, exc_value, traceback)),
616 ),
617 block=False,
618 )
619 self.log.debug("Error message sent to supervisor.")
620 else:
621 self.log.warning("Shutting down due to exception raised on another worker.")
622 self._exit_stack.__exit__(exc_type, exc_value, traceback)
623 return True
625 @property
626 def exit_stack(self) -> ExitStack:
627 """A `contextlib.ExitStack` tied to the communicator."""
628 return self._exit_stack
630 def log_progress(self, level: int, message: str) -> None:
631 """Send a high-level log message to the supervisor.
633 Parameters
634 ----------
635 level : `int`
636 Log level. Should be ``VERBOSE`` or higher.
637 message : `str`
638 Log message.
639 """
640 self._reports.put(_ProgressLog(message=message, level=level), block=False)
642 def check_for_cancel(self) -> None:
643 """Check for a cancel signal from the supervisor and raise
644 `FatalWorkerError` if it is present.
645 """
646 if self._cancel_event.is_set():
647 raise FatalWorkerError()
650class ScannerCommunicator(WorkerCommunicator):
651 """A communicator for scanner workers.
653 Parameters
654 ----------
655 supervisor : `SupervisorCommunicator`
656 Communicator for the supervisor to grab queues and information from.
657 scanner_id : `int`
658 Integer ID for this canner.
659 """
661 def __init__(self, supervisor: SupervisorCommunicator, scanner_id: int):
662 super().__init__(supervisor, f"scanner-{scanner_id:03d}")
663 self.scanner_id = scanner_id
664 self._scan_requests = supervisor._scan_requests
665 self._ingest_requests = supervisor._ingest_requests
666 self._write_requests = supervisor._write_requests
667 self._compression_dict = supervisor._compression_dict
668 self._got_no_more_scan_requests: bool = False
669 self._sent_no_more_ingest_requests: bool = False
671 def report_scan(self, msg: ScanReport) -> None:
672 """Report a completed scan to the supervisor.
674 Parameters
675 ----------
676 msg : `ScanReport`
677 Report to send.
678 """
679 self._reports.put(msg, block=False)
681 def request_ingest(self, request: IngestRequest) -> None:
682 """Ask the ingester to ingest a quantum's outputs.
684 Parameters
685 ----------
686 request : `IngestRequest`
687 Description of the datasets to ingest.
689 Notes
690 -----
691 If this request has no datasets, this automatically reports the ingest
692 as complete to the supervisor instead of sending it to the ingester.
693 """
694 if request:
695 self._ingest_requests.put(request, block=False)
696 else:
697 self._reports.put(_IngestReport(1), block=False)
699 def request_write(self, request: ProvenanceQuantumScanData) -> None:
700 """Ask the writer to write provenance for a quantum.
702 Parameters
703 ----------
704 request : `ProvenanceQuantumScanData`
705 Result of scanning a quantum.
706 """
707 assert self._write_requests is not None, "Writer should not be used if writing is disabled."
708 self._write_requests.put(request, block=False)
710 def get_compression_dict(self) -> bytes | None:
711 """Attempt to get the compression dict from the writer.
713 Returns
714 -------
715 data : `bytes` or `None`
716 The `bytes` representation of the compression dictionary, or `None`
717 if the compression dictionary is not yet available.
719 Notes
720 -----
721 A scanner should only call this method before it actually has the
722 compression dict.
723 """
724 if (cdict := _get_from_queue(self._compression_dict)) is not None:
725 return cdict.data
726 return None
728 def poll(self) -> Iterator[uuid.UUID]:
729 """Poll for scan requests to process.
731 Yields
732 ------
733 quantum_id : `uuid.UUID`
734 ID of a new quantum to scan.
736 Notes
737 -----
738 This iterator ends when the supervisor reports that it is done
739 traversing the graph.
740 """
741 while True:
742 self.check_for_cancel()
743 scan_request = _get_from_queue(self._scan_requests, block=True, timeout=self.config.worker_sleep)
744 if scan_request is _Sentinel.NO_MORE_SCAN_REQUESTS:
745 self._got_no_more_scan_requests = True
746 return
747 if scan_request is not None:
748 yield scan_request.quantum_id
750 def __exit__(
751 self,
752 exc_type: type[BaseException] | None,
753 exc_value: BaseException | None,
754 traceback: TracebackType | None,
755 ) -> bool | None:
756 result = super().__exit__(exc_type, exc_value, traceback)
757 self._ingest_requests.put(_Sentinel.NO_MORE_INGEST_REQUESTS, block=False)
758 if self._write_requests is not None:
759 self._write_requests.put(_Sentinel.NO_MORE_WRITE_REQUESTS, block=False)
760 while not self._got_no_more_scan_requests:
761 self.log.debug("Clearing scan request queue (~%d remaining)", self._scan_requests.qsize())
762 if (
763 not self._got_no_more_scan_requests
764 and self._scan_requests.get() is _Sentinel.NO_MORE_SCAN_REQUESTS
765 ):
766 self._got_no_more_scan_requests = True
767 # We let the supervisor clear out the compression dict queue, because
768 # a single scanner can't know if it ever got sent out or not.
769 self.log.verbose("Sending done sentinal.")
770 self._reports.put(_Sentinel.SCANNER_DONE, block=False)
771 return result
774class IngesterCommunicator(WorkerCommunicator):
775 """A communicator for the ingester worker.
777 Parameters
778 ----------
779 supervisor : `SupervisorCommunicator`
780 Communicator for the supervisor to grab queues and information from.
781 """
783 def __init__(self, supervisor: SupervisorCommunicator):
784 super().__init__(supervisor, "ingester")
785 self.n_scanners = supervisor.n_scanners
786 self._ingest_requests = supervisor._ingest_requests
787 self._n_requesters_done = 0
789 def __exit__(
790 self,
791 exc_type: type[BaseException] | None,
792 exc_value: BaseException | None,
793 traceback: TracebackType | None,
794 ) -> bool | None:
795 result = super().__exit__(exc_type, exc_value, traceback)
796 while self._n_requesters_done != self.n_scanners:
797 self.log.debug(
798 "Waiting for %d requesters to be done (currently %d).",
799 self.n_scanners,
800 self._n_requesters_done,
801 )
802 if self._ingest_requests.get(block=True) is _Sentinel.NO_MORE_INGEST_REQUESTS:
803 self._n_requesters_done += 1
804 self.log.verbose("Sending done sentinal.")
805 self._reports.put(_Sentinel.INGESTER_DONE, block=False)
806 return result
808 def report_ingest(self, n_producers: int) -> None:
809 """Report to the supervisor that an ingest batch was completed.
811 Parameters
812 ----------
813 n_producers : `int`
814 Number of producing quanta whose datasets were ingested.
815 """
816 self._reports.put(_IngestReport(n_producers), block=False)
818 def poll(self) -> Iterator[IngestRequest]:
819 """Poll for ingest requests from the scanner workers.
821 Yields
822 ------
823 request : `IngestRequest`
824 A request to ingest datasets produced by a single quantum.
826 Notes
827 -----
828 This iterator ends when all scanners indicate that they are done making
829 ingest requests.
830 """
831 while True:
832 self.check_for_cancel()
833 ingest_request = _get_from_queue(self._ingest_requests, block=True, timeout=_TINY_TIMEOUT)
834 if ingest_request is _Sentinel.NO_MORE_INGEST_REQUESTS:
835 self._n_requesters_done += 1
836 if self._n_requesters_done == self.n_scanners:
837 return
838 else:
839 continue
840 if ingest_request is not None:
841 yield ingest_request
844class WriterCommunicator(WorkerCommunicator):
845 """A communicator for the writer worker.
847 Parameters
848 ----------
849 supervisor : `SupervisorCommunicator`
850 Communicator for the supervisor to grab queues and information from.
851 """
853 def __init__(self, supervisor: SupervisorCommunicator):
854 assert supervisor._write_requests is not None
855 super().__init__(supervisor, "writer")
856 self.n_scanners = supervisor.n_scanners
857 self._write_requests = supervisor._write_requests
858 self._compression_dict = supervisor._compression_dict
859 self._n_requesters = supervisor.n_scanners + 1
860 self._n_requesters_done = 0
861 self._sent_compression_dict = False
863 def __exit__(
864 self,
865 exc_type: type[BaseException] | None,
866 exc_value: BaseException | None,
867 traceback: TracebackType | None,
868 ) -> bool | None:
869 result = super().__exit__(exc_type, exc_value, traceback)
870 if exc_type is None:
871 self.log_progress(logging.INFO, "Provenance quantum graph written successfully.")
872 while self._n_requesters_done != self._n_requesters:
873 self.log.debug(
874 "Waiting for %d requesters to be done (currently %d).",
875 self._n_requesters,
876 self._n_requesters_done,
877 )
878 if self._write_requests.get(block=True) is _Sentinel.NO_MORE_WRITE_REQUESTS:
879 self._n_requesters_done += 1
880 self.log.verbose("Sending done sentinal.")
881 self._reports.put(_Sentinel.WRITER_DONE, block=False)
882 return result
884 def poll(self) -> Iterator[ProvenanceQuantumScanData]:
885 """Poll for writer requests from the scanner workers and supervisor.
887 Yields
888 ------
889 request : `ProvenanceQuantumScanData`
890 The result of a quantum scan.
892 Notes
893 -----
894 This iterator ends when all scanners and the supervisor indicate that
895 they are done making write requests.
896 """
897 while True:
898 self.check_for_cancel()
899 write_request = _get_from_queue(self._write_requests, block=True, timeout=_TINY_TIMEOUT)
900 if write_request is _Sentinel.NO_MORE_WRITE_REQUESTS:
901 self._n_requesters_done += 1
902 if self._n_requesters_done == self._n_requesters:
903 return
904 else:
905 continue
906 if write_request is not None:
907 yield write_request
909 def send_compression_dict(self, cdict_data: bytes) -> None:
910 """Send the compression dictionary to the scanners.
912 Parameters
913 ----------
914 cdict_data : `bytes`
915 The `bytes` representation of the compression dictionary.
916 """
917 self.log.debug("Sending compression dictionary.")
918 for _ in range(self.n_scanners):
919 self._compression_dict.put(_CompressionDictionary(cdict_data), block=False)
920 self._sent_compression_dict = True
922 def report_write(self) -> None:
923 """Report to the supervisor that provenance for a quantum was written
924 to the graph.
925 """
926 self._reports.put(_Sentinel.WRITE_REPORT, block=False)
928 def periodically_check_for_cancel(self, iterable: Iterable[_T], n: int = 100) -> Iterator[_T]:
929 """Iterate while checking for a cancellation signal every ``n``
930 iterations.
932 Parameters
933 ----------
934 iterable : `~collections.abc.Iterable`
935 Object to iterate over.
936 n : `int`
937 Check for cancellation every ``n`` iterations.
939 Returns
940 -------
941 iterator : `~collections.abc.Iterator`
942 Iterator.
943 """
944 i = 0
945 for entry in iterable:
946 yield entry
947 i += 1
948 if i % n == 0:
949 self.check_for_cancel()