X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=executors.py;h=b4cb06b4c9ad80816f0bb13e05e23863846fa8c1;hb=b6dd2da4aecbc5d4690023d5a5e4f63052425444;hp=28507b0e8a547e9d1edc615daa5de7945edb457d;hpb=1a092f01f1a945a330fe19e96a8de507823fc28e;p=python_utils.git diff --git a/executors.py b/executors.py index 28507b0..b4cb06b 100644 --- a/executors.py +++ b/executors.py @@ -2,33 +2,32 @@ from __future__ import annotations -from abc import ABC, abstractmethod import concurrent.futures as fut -from collections import defaultdict -from dataclasses import dataclass import logging -import numpy import os import platform import random import subprocess import threading import time -from typing import Any, Callable, Dict, List, Optional, Set import warnings +from abc import ABC, abstractmethod +from collections import defaultdict +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Set import cloudpickle # type: ignore +import numpy from overrides import overrides -from ansi import bg, fg, underline, reset import argparse_utils import config -from decorator_utils import singleton -from exec_utils import run_silently, cmd_in_background, cmd_with_timeout import histogram as hist +from ansi import bg, fg, reset, underline +from decorator_utils import singleton +from exec_utils import cmd_in_background, cmd_with_timeout, run_silently from thread_utils import background_thread - logger = logging.getLogger(__name__) parser = config.add_commandline_args( @@ -84,10 +83,10 @@ class BaseExecutor(ABC): pass @abstractmethod - def shutdown(self, wait: bool = True) -> None: + def shutdown(self, *, wait: bool = True, quiet: bool = False) -> None: pass - def shutdown_if_idle(self) -> bool: + def shutdown_if_idle(self, *, quiet: bool = False) -> bool: """Shutdown the executor and return True if the executor is idle (i.e. there are no pending or active tasks). Return False otherwise. Note: this should only be called by the launcher @@ -95,7 +94,7 @@ class BaseExecutor(ABC): """ if self.task_count == 0: - self.shutdown() + self.shutdown(wait=True, quiet=quiet) return True return False @@ -155,11 +154,12 @@ class ThreadExecutor(BaseExecutor): return result @overrides - def shutdown(self, wait=True) -> None: + def shutdown(self, *, wait: bool = True, quiet: bool = False) -> None: if not self.already_shutdown: logger.debug(f'Shutting down threadpool executor {self.title}') - print(self.histogram.__repr__(label_formatter='%d')) self._thread_pool_executor.shutdown(wait) + if not quiet: + print(self.histogram.__repr__(label_formatter='%ds')) self.already_shutdown = True @@ -197,11 +197,12 @@ class ProcessExecutor(BaseExecutor): return result @overrides - def shutdown(self, wait=True) -> None: + def shutdown(self, *, wait: bool = True, quiet: bool = False) -> None: if not self.already_shutdown: logger.debug(f'Shutting down processpool executor {self.title}') self._process_executor.shutdown(wait) - print(self.histogram.__repr__(label_formatter='%d')) + if not quiet: + print(self.histogram.__repr__(label_formatter='%ds')) self.already_shutdown = True def __getstate__(self): @@ -246,7 +247,7 @@ class BundleDetails: end_ts: float slower_than_local_p95: bool slower_than_global_p95: bool - src_bundle: BundleDetails + src_bundle: Optional[BundleDetails] is_cancelled: threading.Event was_cancelled: bool backup_bundles: Optional[List[BundleDetails]] @@ -286,7 +287,7 @@ class RemoteExecutorStatus: self.worker_count: int = total_worker_count self.known_workers: Set[RemoteWorkerRecord] = set() self.start_time: float = time.time() - self.start_per_bundle: Dict[str, float] = defaultdict(float) + self.start_per_bundle: Dict[str, Optional[float]] = defaultdict(float) self.end_per_bundle: Dict[str, float] = defaultdict(float) self.finished_bundle_timings_per_worker: Dict[ RemoteWorkerRecord, List[float] @@ -343,7 +344,9 @@ class RemoteExecutorStatus: self.end_per_bundle[uuid] = ts self.in_flight_bundles_by_worker[worker].remove(uuid) if not was_cancelled: - bundle_latency = ts - self.start_per_bundle[uuid] + start = self.start_per_bundle[uuid] + assert start is not None + bundle_latency = ts - start x = self.finished_bundle_timings_per_worker.get(worker, list()) x.append(bundle_latency) self.finished_bundle_timings_per_worker[worker] = x @@ -737,7 +740,7 @@ class RemoteExecutor(BaseExecutor): worker = None while worker is None: worker = self.find_available_worker_or_block(avoid_machine) - assert worker + assert worker is not None # Ok, found a worker. bundle.worker = worker @@ -834,9 +837,10 @@ class RemoteExecutor(BaseExecutor): return self.wait_for_process(p, bundle, 0) def wait_for_process( - self, p: subprocess.Popen, bundle: BundleDetails, depth: int + self, p: Optional[subprocess.Popen], bundle: BundleDetails, depth: int ) -> Any: machine = bundle.machine + assert p is not None pid = p.pid if depth > 3: logger.error( @@ -979,10 +983,12 @@ class RemoteExecutor(BaseExecutor): # Tell the original to stop if we finished first. if not was_cancelled: + orig_bundle = bundle.src_bundle + assert orig_bundle is not None logger.debug( - f'{bundle}: Notifying original {bundle.src_bundle.uuid} we beat them to it.' + f'{bundle}: Notifying original {orig_bundle.uuid} we beat them to it.' ) - bundle.src_bundle.is_cancelled.set() + orig_bundle.is_cancelled.set() self.release_worker(bundle, was_cancelled=was_cancelled) return result @@ -1066,7 +1072,9 @@ class RemoteExecutor(BaseExecutor): # they will move the result_file to this machine and let # the original pick them up and unpickle them. - def emergency_retry_nasty_bundle(self, bundle: BundleDetails) -> fut.Future: + def emergency_retry_nasty_bundle( + self, bundle: BundleDetails + ) -> Optional[fut.Future]: is_original = bundle.src_bundle is None bundle.worker = None avoid_last_machine = bundle.machine @@ -1107,13 +1115,14 @@ class RemoteExecutor(BaseExecutor): return self._helper_executor.submit(self.launch, bundle) @overrides - def shutdown(self, wait=True) -> None: + def shutdown(self, *, wait: bool = True, quiet: bool = False) -> None: if not self.already_shutdown: logging.debug(f'Shutting down RemoteExecutor {self.title}') self.heartbeat_stop_event.set() self.heartbeat_thread.join() self._helper_executor.shutdown(wait) - print(self.histogram.__repr__(label_formatter='%d')) + if not quiet: + print(self.histogram.__repr__(label_formatter='%ds')) self.already_shutdown = True @@ -1212,11 +1221,11 @@ class DefaultExecutors(object): def shutdown(self) -> None: if self.thread_executor is not None: - self.thread_executor.shutdown() + self.thread_executor.shutdown(wait=True, quiet=True) self.thread_executor = None if self.process_executor is not None: - self.process_executor.shutdown() + self.process_executor.shutdown(wait=True, quiet=True) self.process_executor = None if self.remote_executor is not None: - self.remote_executor.shutdown() + self.remote_executor.shutdown(wait=True, quiet=True) self.remote_executor = None