Source code for libertem.common.progress

import threading
from typing import TYPE_CHECKING, Iterable, Dict, Callable, List, Any, NamedTuple, Optional
import time

from libertem.common.executor import WorkerQueueEmpty

    from libertem.common.executor import WorkerQueue, TaskCommHandler
    from libertem.udf.base import UDFTask
    from import DataTile
    from import Partition

class CommsDispatcher:
    Monitors a :code:`WorkerQueue` in a background thread
    and launches callbacks in response to messages recieved
    Callbacks are registered as a dictionary of subscriptions
        {topic_name: [callback, ...]]}
    and are called in order they were recieved in the same thread
    that is doing the monitoring. The callbacks should be lightweight
    in order to not build up too many messages in the queue

    The functionality of this class mirrors Dask's structured logs
    feature, which has a similar message => topic => callback
    model running in the client event loop
    def __init__(self, queue: 'WorkerQueue', subscriptions: Dict[str, List[Callable]]):
        self._message_q = queue
        self._subscriptions = subscriptions
        self._thread = None

    def __enter__(self, *args, **kwargs):
        if self._thread is not None:
            raise RuntimeError('Cannot re-enter CommsDispatcher')
        self._thread = threading.Thread(
        self._thread.daemon = True

    def __exit__(self, *args, **kwargs):
        if self._thread is None:
        self._message_q.put(('STOP', {}))
        self._thread = None
        # Drain queue just in case
        while True:
                with self._message_q.get(block=False) as _:
            except WorkerQueueEmpty:

    def monitor_queue(self):
        Monitor the queue for messages
        If there are no subscribers this should drain
        messages from the queue as fast as they are recieved
        while True:
            with self._message_q.get(block=True) as ((topic, msg), _):
                if topic == 'STOP':
                    for callback in self._subscriptions[topic]:
                        callback(topic, msg)
                except KeyError:

[docs]class ProgressState(NamedTuple): """ Container for progress state, used to communicate from ProgressManager to ProgressReporter """ #: float: Number of frames processed num_frames_complete: float #: int: Total number of frames to process num_frames_total: int #: int: Number of partitions completed num_part_complete: int #: int: Number of partitions in-progress num_part_in_progress: int #: int: Total number of partitions num_part_total: int #: str: A unique string identifier for the job associated #: with this progress message progress_id: str
[docs]class ProgressReporter: """ Interface for progress bar display / updating This class will receive :class:`ProgressState` instances to notify it about the start, progression and end of a job submitted to the :code:`UDFRunner`. The implementation should be adapted to display or log the progress as required for the use case. It is possible that multiple jobs are submitted to a single executor at the same time and therefore the implementation should ensure that concurrent instances of the class display correctly, or that the same instance of the class can handle updates from multiple threads concurrently. Each :class:`ProgressState` message contains a field :code:`progress_id` which is unique to each job, and therefore the implementation can use this to distinguish updates from multiple sources. """ def __init__(self): raise NotImplementedError()
[docs] def start(self, state: ProgressState): """ Signal the creation of a new job with the expected number of partitions and frames, and unique progress_id string. """ raise NotImplementedError()
[docs] def update(self, state: ProgressState): """ Signal an intermediate update to the progress of a job """ raise NotImplementedError()
[docs] def end(self, state: ProgressState): """ Signal the end of a given job This method will always be called and any updates recieved after this message should be ignored. """ raise NotImplementedError()
class TQDMProgressReporter(ProgressReporter): """ Progress bar display via tqdm Supports concurrent usage of multiple instances of this class, but does not handle multi-threaded use of the same instance to report multiple jobs. """ def __init__(self): self._bar = None # Integers used to check if bar description should be changed # integers because faster than strcmp, updated in _should_update_description self._desc_key = (-1, -1, -1) def start(self, state: ProgressState): from import tqdm self._bar = tqdm(desc=self._get_description(state), total=state.num_frames_total, leave=True) def update(self, state: ProgressState): return self._update(state, clip=True, refresh=False) def _update(self, state: ProgressState, *, clip: bool, refresh: bool): if state.num_frames_total != # Should never happen but handle just in case = state.num_frames_total self._bar.refresh() increment = self._get_increment(state, clip=clip) if increment > 0: self._bar.update(increment) if self._should_update_description(state): self._bar.set_description(self._get_description(state)) if refresh: self._bar.refresh() def end(self, state: ProgressState): self._update(state, clip=True, refresh=True) self._bar.close() def _should_update_description(self, state: ProgressState) -> bool: """ Check the state to see if the elements used by self._get_description have changed, and if so update our record (self._desc_key) and return True """ new_desc_key = ( state.num_part_complete, state.num_part_in_progress, state.num_part_total ) should_update = new_desc_key != self._desc_key if should_update: self._desc_key = new_desc_key return should_update @staticmethod def _get_description(state: ProgressState) -> str: """ Get the most recent description string for the bar, including partition information If we know that partitions are in progress include this in parentheses after n_completed """ if state.num_part_in_progress: return (f'Partitions {state.num_part_complete}({state.num_part_in_progress})' f'/{state.num_part_total}, Frames') else: return (f'Partitions {state.num_part_complete}' f'/{state.num_part_total}, Frames') def _get_increment(self, state: ProgressState, clip: bool = True): """ Get the increment to apply to the progress bar based on current state and the state of the bar itself (bar.n is the total as-tracked by tqdm) Assumes has first been updated to state.num_frames_total if this is necessary """ increment = int(state.num_frames_complete) - self._bar.n if clip: max_update = - self._bar.n else: max_update = increment + 1 return max(0, min(increment, max_update)) class ProgressManager: """ Handle updating of a progress reporter for a set of :code:`UDFTasks`, to be completed in any order. By default constructs a :code:`TQDMProgressReporter`, if no instance is passed in. The bar displays as such: Partitions: n_complete(n_in_progress) / n_total, ...\ Frames: [XXXXX..] frames_completed / total_frames ... When processing tile stacks, stacks are treated as frames as such: (pseudo_frames = tile.size // sig_size) The bar will render in a Jupyter notebook as a JS widget automatically via """ def __init__( self, tasks: Iterable['UDFTask'], progress_id: str, reporter: Optional[ProgressReporter] = None, ): if not tasks: raise ValueError('Cannot display progress for empty tasks') self._progress_id = progress_id # the number of whole frames we expect each task to process self._task_max = {t.partition.get_ident(): t.task_frames for t in tasks} # _counters is our record of progress on a task, # values are floating whole frames processed # as in tile mode we can process part of a frame self._counters = {k: 0. for k in self._task_max.keys()} self._total_frames = sum(self._task_max.values()) # For converting tiles to pseudo-frames self._sig_size = tasks[0].partition.shape.sig.size # Counters for part display self._complete = set() self._in_progress = set() self._num_total = len(self._counters) if reporter is None: reporter = TQDMProgressReporter() elif not isinstance(reporter, ProgressReporter): # If not a ProgressReporter instance, # instantiate as if it has a bare __init__ # Useful to be able to inject an instance # of ProgressReporter in case we need to setup # or access the reporter somehow (e.g. for testing) reporter = reporter() assert isinstance(reporter, ProgressReporter) self.reporter = reporter reporter.start(self.state) @property def state(self) -> ProgressState: return ProgressState( sum(self._counters.values()), self._total_frames, len(self._complete), len(self._in_progress), self._num_total, self._progress_id, ) def finalize_task(self, task: 'UDFTask'): """ When a task completes and we recieve its results on the main node, this is called to update the partition progress counters and frame counter in case we didn't recieve a complete history of the partition yet """ topic = 'partition_complete' ident = task.partition.get_ident() message = {'ident': task.partition.get_ident()} if ident in self._task_max: self.handle_end_task(topic, message) def close(self): self.reporter.end(self.state) def connect(self, comms: 'TaskCommHandler'): """ Register the callbacks on this class with the TaskCommHandler which will be dispatching messages recieved from the tasks """ comms.subscribe('partition_start', self.handle_start_task) comms.subscribe('partition_complete', self.handle_end_task) comms.subscribe('tile_complete', self.handle_tile_update) def handle_start_task(self, topic: str, message: Dict[str, Any]): """ Increment the num_in_progress counter # NOTE An extension to this would be to track the identities of partitions in progress / completed for a richer display / more accurate accounting """ if topic != 'partition_start': raise RuntimeError('Unrecognized topic') t_id = message['ident'] if t_id not in self._complete: # if not complete handles case task was completed # before we can process its start message, # completion is signalled in the main thread # while start messages are processed in the background self._in_progress.add(t_id) self.reporter.update(self.state) def handle_end_task(self, topic: str, message: Dict[str, Any]): """ Increment the counter for the task to the max value and update the various counters / description """ if topic != 'partition_complete': raise RuntimeError('Unrecognized topic') t_id = message['ident'] remain = self._task_max[t_id] - int(self._counters[t_id]) if remain: self._counters[t_id] = self._task_max[t_id] self._in_progress.discard(t_id) self._complete.add(t_id) self.reporter.update(self.state) def handle_tile_update(self, topic: str, message: Dict[str, Any]): """ Update the frame progress counter for the task and push the increment to the progress reporter Tile stacks are converted to pseudo-frames via the sig_size """ if topic != 'tile_complete': raise RuntimeError('Unrecognized topic') t_id = message['ident'] if self._counters[t_id] >= self._task_max[t_id]: return elements = message['elements'] pframes = elements / self._sig_size self._counters[t_id] += pframes self.reporter.update(self.state) class PartitionTrackerNoOp: """ A no-op class matching the PartitionProgressTracker interface Used when progress == False to avoid any additional overhead """ def signal_start(self, *args, **kwargs): ... def signal_tile_complete(self, *args, **kwargs): ... def signal_complete(self, *args, **kwargs): ... def get_time(): # Exists for testing / mocking return time.time() class PartitionProgressTracker(PartitionTrackerNoOp): """ Tracks the tile processing speed of a Partition and dispatches messages via the worker_context.signal() method under certain conditions Parameters ---------- partition : Partition The partition to track progress for min_message_interval : float, optional The minumum time between messages, by default 1 second. """ def __init__( self, partition: 'Partition', min_message_interval: float = 1., ): self._ident = partition.get_ident() try: self._worker_context = partition._worker_context except AttributeError: self._worker_context = None # Counters to track / rate-limit messages self._elements_complete = 0 self._last_message_t = None self._min_message_interval = min_message_interval def signal_start(self): """ Signal that the partition has begun processing """ if self._worker_context is None: return self._worker_context.signal( self._ident, 'partition_start', {}, ) def signal_tile_complete(self, tile: 'DataTile'): """ Register that tile.size more elements have been processed and if certain condition are met, send a signal """ if self._worker_context is None: return send_elements = self.should_send_progress(tile.size) if send_elements: self._worker_context.signal( self._ident, 'tile_complete', {'elements': send_elements}, ) def signal_complete(self): """ Signal that the partition has completed processing This is not currently called as partition completion is registered on the main node as a fallback """ if self._worker_context is None: return self._worker_context.signal( self._ident, 'partition_complete', {}, ) def should_send_progress(self, elements: int) -> int: """ Given the number elements of data that have been processed since the last message was sent, decide if a signal should be sent to the main node about the partition progress """ current_t = get_time() self._elements_complete += elements if self._last_message_t is None: # Never send a message for the first tile stack # as this might have warmup overheads associated # Include the first elements in the history, however, # to give a better accounting. The first tile stack # is essentially treated as 'free'. self._last_message_t = current_t return 0 time_since_last_m = current_t - self._last_message_t not_rate_limited = time_since_last_m > self._min_message_interval if not_rate_limited: completed_elements = self._elements_complete self._elements_complete = 0 self._last_message_t = current_t return completed_elements return 0