# SPDX-FileCopyrightText: 2024-2025 ModelCloud.ai # SPDX-FileCopyrightText: 2024-2025 qubitium@modelcloud.ai # SPDX-License-Identifier: Apache-2.0 # Contact: qubitium@modelcloud.ai, x.com/qubitium from __future__ import annotations import contextlib import os import queue import sys import threading import time import traceback from concurrent.futures import Future from dataclasses import dataclass from datetime import datetime, timezone from enum import Enum from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Tuple, Union import torch try: from device_smi import Device # type: ignore except Exception: # pragma: no cover - defensive: optional dependency may be unavailable Device = None from .. import DEBUG_ON from ..utils import torch as torch_utils from ..utils.ctx import ctx from ..utils.device_telemetry import capture_device_telemetry from ..utils.logger import setup_logger from ..utils.torch import HAS_NPU, torch_empty_cache_any log = setup_logger() def _best_effort_stderr_write(message: str) -> None: with contextlib.suppress(Exception): sys.stderr.write(f"{message}\n") sys.stderr.flush() # Debug logging is very chatty and can alter timings subtly in tests. # We gate all extra diagnostics behind the DEBUG env (1/true/yes/on). # DeviceLike allows ergonomic call sites: 'cuda:0', 0, torch.device('cuda', 0), etc. DeviceLike = Union[str, int, torch.device] WarmupFn = Callable[[torch.device, "WarmUpCtx"], None] class WarmUpCtx(str, Enum): THREAD = "thread" DEVICE = "device" THREAD_AND_DEVICE = "thread_and_device" def invocation_order(self) -> Tuple["WarmUpCtx", ...]: if self is WarmUpCtx.THREAD_AND_DEVICE: return (WarmUpCtx.DEVICE, WarmUpCtx.THREAD) return (self,) @dataclass(frozen=True) class WarmupTask: fn: WarmupFn scope: WarmUpCtx = WarmUpCtx.DEVICE # --------------------------- Backend availability helpers --------------------------- # We keep these helpers small and side-effect free—only feature checks—so we can # query once and rely on final results without redundant availability if-ladders. def _mps_available() -> bool: return ( hasattr(torch, "backends") and hasattr(torch.backends, "mps") and torch.backends.mps.is_available() ) # --------------------------- Device coercion & context helpers --------------------------- def _coerce_device(d: DeviceLike) -> torch.device: """ Convert a DeviceLike into a concrete torch.device. For integers, we interpret as accelerator indices if present, otherwise map to CPU/MPS. """ if isinstance(d, torch.device): return d if isinstance(d, int): if torch.cuda.is_available(): return torch.device("cuda", d) if hasattr(torch, "xpu") and torch.xpu.is_available(): return torch.device("xpu", d) if HAS_NPU: return torch.device("npu", d) if _mps_available(): return torch.device("mps") return torch.device("cpu") # Accept strings like 'cuda:0', 'xpu:1', 'npu:2', 'cpu', 'mps' return torch.device(d) @contextlib.contextmanager def _device_ctx(dev: torch.device): """ Set the caller thread’s *current* device while running a task so handles/streams line up correctly. For CUDA/XPU we set the per-thread current device; CPU/MPS do not require pinning here. """ if dev.type == "cuda": target = dev if dev.index is not None else "cuda" with torch.cuda.device(target): yield elif dev.type == "xpu" and hasattr(torch, "xpu"): target = dev if dev.index is not None else "xpu" with torch.xpu.device(target): yield elif dev.type == "npu" and HAS_NPU: target = dev if dev.index is not None else "npu" with torch.npu.device(target): yield else: yield def _activate_thread_device(dev: torch.device): """ Pin the worker thread to its device once, before entering its main loop. CUDA/XPU require per-thread device activation for correct handle usage. """ if dev.type == "cuda": target = dev if dev.index is not None else "cuda" torch.cuda.set_device(target) elif dev.type == "xpu" and hasattr(torch, "xpu"): target = dev if dev.index is not None else "xpu" torch.xpu.set_device(target) elif dev.type == "npu" and HAS_NPU: target = dev if dev.index is not None else "npu" torch.npu.set_device(target) # mps/cpu: nothing to pin def _pop_public_kwarg(kwargs: Dict[str, Any], public_name: str, private_name: str): if private_name in kwargs: raise TypeError(f"'{private_name}' is no longer supported; use '{public_name}'") if public_name in kwargs: return kwargs.pop(public_name) return None # --------------------------- Read-Write Lock (writer-preference) --------------------------- # We implement a writer-preference RWLock. Multiple readers may hold the lock, # but when a writer is waiting we block new readers, ensuring GC (writer) can # eventually acquire exclusivity even under task pressure. class _RWLock: """ Reader-writer lock with writer preference. - Multiple readers may hold the lock simultaneously. - A single writer holds exclusivity. - When a writer is waiting, new readers will block. - Writer is re-entrant for its owning thread. """ def __init__(self): self._cond = threading.Condition() self._readers = 0 self._writer: Optional[int] = None # thread id that owns write self._writer_depth = 0 self._writers_waiting = 0 # --- Write (exclusive) --- def acquire_write(self): me = threading.get_ident() with self._cond: if self._writer == me: # Re-entrant writer self._writer_depth += 1 return self._writers_waiting += 1 try: while self._writer is not None or self._readers > 0: self._cond.wait() self._writer = me self._writer_depth = 1 finally: self._writers_waiting -= 1 def release_write(self): me = threading.get_ident() with self._cond: if self._writer != me: raise RuntimeError("release_write called by non-owner") self._writer_depth -= 1 if self._writer_depth == 0: self._writer = None self._cond.notify_all() @contextlib.contextmanager def writer(self): self.acquire_write() try: yield finally: self.release_write() # --- Read (shared) --- def acquire_read(self): me = threading.get_ident() with self._cond: # The writer may re-enter as a reader; this keeps invariants simple # for code that wants to read while already holding write. if self._writer == me: self._readers += 1 return # If a writer is waiting, block new readers to give it priority. while self._writer is not None or self._writers_waiting > 0: self._cond.wait() self._readers += 1 def release_read(self): with self._cond: if self._readers <= 0: raise RuntimeError("release_read without acquire_read") self._readers -= 1 if self._readers == 0: self._cond.notify_all() @contextlib.contextmanager def reader(self): self.acquire_read() try: yield finally: self.release_read() class _LockGroup(contextlib.AbstractContextManager): """ Acquire multiple device **write** locks in deterministic order to avoid deadlocks. Helpful for GC passes or any multi-device exclusive operation. """ def __init__(self, ordered_pairs: List[tuple[str, _RWLock]]): self._pairs = ordered_pairs def __enter__(self): for name, lk in self._pairs: if DEBUG_ON: log.debug(f"_LockGroup: acquiring write lock for {name}") lk.acquire_write() if DEBUG_ON: log.debug(f"_LockGroup: acquired write lock for {name}") return self def __exit__(self, exc_type, exc, tb): for name, lk in reversed(self._pairs): if DEBUG_ON: log.debug(f"_LockGroup: releasing write lock for {name}") lk.release_write() return False class _ReadLockGroup(contextlib.AbstractContextManager): """ Acquire multiple device **read** locks in deterministic order. Useful for multi-device snapshots or read-only, consistent views. """ def __init__(self, ordered_pairs: List[tuple[str, _RWLock]]): self._pairs = ordered_pairs def __enter__(self): for name, lk in self._pairs: if DEBUG_ON: log.debug(f"_ReadLockGroup: acquiring read lock for {name}") lk.acquire_read() if DEBUG_ON: log.debug(f"_ReadLockGroup: acquired read lock for {name}") return self def __exit__(self, exc_type, exc, tb): for name, lk in reversed(self._pairs): if DEBUG_ON: log.debug(f"_ReadLockGroup: releasing read lock for {name}") lk.release_read() return False class _WaitAndLock(contextlib.AbstractContextManager): """ Context manager returned by pool.wait(scope, lock=True). On enter: acquires writer locks over the scope in canonical order, which inherently waits for in-flight readers (tasks) to drain. On exit: releases locks. """ def __init__(self, pairs: List[tuple[str, _RWLock]]): self._pairs = pairs self._group = _LockGroup(pairs) def __enter__(self): return self._group.__enter__() def __exit__(self, exc_type, exc, tb): return self._group.__exit__(exc_type, exc, tb) class _WorkerWarmupState: """ Coordinate a single warmup invocation. Device-scoped states may be shared across workers for the same physical device. Thread-scoped states are owned by a single worker. """ def __init__(self, warmup_task: WarmupTask, warmup_ctx: WarmUpCtx): self._warmup_task = warmup_task self._warmup_ctx = warmup_ctx self._claim_lock = threading.Lock() self._started = False self._done = threading.Event() self._error: Optional[BaseException] = None def run(self, *, device: torch.device, rwlock: _RWLock) -> None: if self._done.is_set(): self._raise_if_failed() return should_run = False with self._claim_lock: if self._done.is_set(): pass elif not self._started: self._started = True should_run = True if should_run: try: with ctx(rwlock.reader(), _device_ctx(device)): self._warmup_task.fn(device, self._warmup_ctx) except BaseException as exc: with self._claim_lock: self._error = exc raise finally: self._done.set() else: self._done.wait() self._raise_if_failed() def _raise_if_failed(self) -> None: with self._claim_lock: error = self._error if error is not None: raise error # --------------------------- Worker Thread --------------------------- # Each worker is bound to a specific device and runs a single thread. Tasks are # executed under the device’s read lock; GC acquires the writer lock to keep # memory management steps from interleaving with tasks. class _DeviceWorker: """ Single worker thread bound to one device. Queue entries: (is_task: bool, fn, args, kwargs, future) - is_task=False is a sentinel to exit the thread loop. - Tasks run within a device-scoped reader lock to prevent interleaving with GC passes (which need a write lock). """ def __init__( self, device: torch.device, rwlock: _RWLock, on_task_finished: Callable[[str], None], on_worker_exit: Callable[[str, "_DeviceWorker"], None], name: Optional[str] = None, inference_mode: bool = False, cpu_core: Optional[int] = None, warmup_plan: Optional[Tuple[_WorkerWarmupState, ...]] = None, *, key_override: Optional[str] = None, ): self.device = device self.rwlock = rwlock self._on_task_finished = on_task_finished self._on_worker_exit = on_worker_exit self._pending_warmups = warmup_plan or () if key_override is not None: self.key = key_override else: self.key = f"{device.type}:{device.index}" if device.index is not None else device.type self.name = name or f"DPWorker-{self.key}" self._q: "queue.Queue[Tuple[bool, Callable[..., Any], tuple, dict, Future]]" = queue.Queue() self._stop = threading.Event() self._inference_mode = inference_mode self._target_cpu_core = cpu_core self._affinity_applied = False self._thread = threading.Thread(target=self._run, name=self.name, daemon=True) self._thread.start() if DEBUG_ON: log.debug(f"Spawned worker '{self.name}' for {self.key}") def submit(self, fn: Callable[..., Any], /, *args, **kwargs) -> Future: """ Enqueue a callable and return a Future that resolves with its result/exception. """ fut = Future() self._q.put((True, fn, args, kwargs, fut)) if DEBUG_ON: log.debug(f"{self.name}: task enqueued; qsize={self._q.qsize()}") return fut def stop(self): """ Request thread exit by setting a sentinel work item. The run loop exits ASAP after receiving it. """ self._stop.set() self._q.put((False, lambda: None, (), {}, Future())) if DEBUG_ON: log.debug(f"{self.name}: stop requested; sentinel queued") def join(self): """ Join the worker thread; for graceful shutdowns and tests. """ if DEBUG_ON: log.debug(f"{self.name}: joining thread") self._thread.join() def _apply_cpu_affinity(self) -> None: if self._affinity_applied or self._target_cpu_core is None: return if not hasattr(os, "sched_setaffinity"): log.warn( "Thread pinning unsupported on this platform; %s will use OS scheduling.", self.name, ) self._affinity_applied = True return try: os.sched_setaffinity(0, {int(self._target_cpu_core)}) self._affinity_applied = True #if DEBUG_ON: log.debug(f"{self.name}: pinned to CPU core {self._target_cpu_core}") except PermissionError as exc: log.warn( "Thread pinning permission denied for %s on core %s (%s); falling back to OS scheduling.", self.name, self._target_cpu_core, exc, ) self._affinity_applied = True except OSError as exc: log.warn( "Thread pinning failed for %s on core %s (%s); falling back to OS scheduling.", self.name, self._target_cpu_core, exc, ) self._affinity_applied = True def _run_warmup(self) -> None: pending_warmups = self._pending_warmups if not pending_warmups: return for warmup_state in pending_warmups: warmup_state.run(device=self.device, rwlock=self.rwlock) self._pending_warmups = () def _run(self): """ Main loop: pull tasks, set device context, execute, mark completion, and fulfill or fail the future. Completion is accounted BEFORE resolving the future to make stats() deterministic even under test interleavings. Workers default to inference mode for throughput but individual tasks may override via `inference_mode`. """ self._apply_cpu_affinity() _activate_thread_device(self.device) while not self._stop.is_set(): is_task, fn, args, kwargs, fut = self._q.get() try: if not is_task: if DEBUG_ON: log.debug(f"{self.name}: received sentinel; exiting") break if DEBUG_ON: log.debug(f"{self.name}: task begin; qsize={self._q.qsize()}") try: self._run_warmup() except BaseException as exc: self._abort_process(exc) return event = kwargs.pop("cuda_event", None) override_inference = _pop_public_kwarg( kwargs, "inference_mode", "_threadx_inference_mode" ) use_inference = self._inference_mode if override_inference is None else bool(override_inference) # Tasks take a **read** lock so janitor's write lock can't interleave with ctx(self.rwlock.reader(), _device_ctx(self.device)): inference_ctx = torch.inference_mode() if use_inference else contextlib.nullcontext() with inference_ctx: if event is not None and self.device.type == "cuda": event.synchronize() result = fn(*args, **kwargs) # Counters must be updated before resolving futures to prevent # tests reading stats mid-transition and seeing stale totals. self._on_task_finished(self.key) if not fut.cancelled(): fut.set_result(result) if DEBUG_ON: log.debug(f"{self.name}: task done") except BaseException as exc: # Even on exception we must decrement inflight and update totals. self._on_task_finished(self.key) if not fut.cancelled(): fut.set_exception(exc) if DEBUG_ON: log.debug(f"{self.name}: task exception: {exc!r}") self._abort_process(exc) finally: self._q.task_done() try: self._on_worker_exit(self.key, self) finally: if DEBUG_ON: log.debug(f"{self.name}: exited") def _abort_process(self, exc: BaseException) -> None: """Dump the stack trace and terminate the interpreter without cleanup.""" try: traceback.print_exception(type(exc), exc, exc.__traceback__, file=sys.stderr) sys.stderr.flush() except Exception: pass try: os._exit(1) except Exception: # Last resort if os._exit is unavailable for some reason. os.system("kill -9 %d" % os.getpid()) class _SyncWorker: """Fallback worker that executes tasks synchronously when threads are unsafe.""" def __init__( self, *, key: str, device: torch.device, rwlock: _RWLock, on_task_finished: Callable[[str], None], on_worker_exit: Callable[[str, "_SyncWorker"], None], inference_mode: bool = False, ) -> None: self.key = key self.device = device self.rwlock = rwlock self._on_task_finished = on_task_finished self._on_worker_exit = on_worker_exit self._inference_mode = inference_mode self.name = f"DPWorker-{key}#sync" def submit(self, fn: Callable[..., Any], /, *args, cuda_event=None, **kwargs) -> Future: fut = Future() try: event = cuda_event override_inference = _pop_public_kwarg( kwargs, "inference_mode", "_threadx_inference_mode" ) use_inference = self._inference_mode if override_inference is None else bool(override_inference) with ctx(self.rwlock.reader(), _device_ctx(self.device)): # with tctl.threadpool_limits(limits=1): inference_ctx = torch.inference_mode() if use_inference else contextlib.nullcontext() with inference_ctx: if event is not None and self.device.type == "cuda": event.synchronize() result = fn(*args, **kwargs) self._on_task_finished(self.key) if not fut.cancelled(): fut.set_result(result) except BaseException as exc: self._on_task_finished(self.key) if not fut.cancelled(): fut.set_exception(exc) return fut def stop(self) -> None: self._on_worker_exit(self.key, self) def join(self) -> None: return # --------------------------- Public Pool --------------------------- # - Builds workers per device with per-device RWLocks # - Tracks inflight counts (with condition vars) and completion counters # - Provides wait() with optional lock=True for exclusive operations # - Runs a janitor background thread that performs periodic empty_cache() under # exclusive writer locks, coalescing triggers via a short debounce window. class DeviceThreadPool: """ Multi-device thread pool with: - Eager discovery/creation of workers and locks for CUDA/XPU/MPS/CPU. - Configurable worker counts per device (default 1). - Correct per-thread device context. - submit()/do() for async/sync, with optional CUDA event synchronization. - Per-device RWLocks + global lock and family/all read-locks. - wait(scope, lock=False/True) to drain tasks (optionally with exclusive locks). - Per-device/global completed counters and in-flight counters. - Janitor: triggers empty-cache after N completions on accelerator devices (per-device lock). - GC diagnostics helpers (snapshots and ANSI tables). """ def __init__( self, devices: Optional[Iterable[DeviceLike]] = None, *, include_cuda: bool = True, include_xpu: bool = True, include_npu: bool = True, include_mps: bool = True, include_cpu: bool = True, inference_mode: bool = False, warmups: Optional[Dict[str, WarmupTask]] = None, empty_cache_every_n: int = 50, # <=0 disables janitor workers: Optional[Dict[str, int]] = None, # e.g. {'cpu':4, 'cuda:per':1, 'cuda:0':3} gc_debounce_seconds: float = 0.02, # absorb bursty triggers before GC gc_min_interval_seconds: float = 0.0, # throttle janitor passes pin_cpu_workers: bool = False, pin_accelerator_workers: bool = False, ): """ Args: devices: explicit list of devices. If None, auto-discover per include_* flags. workers: dict mapping worker-count policy: - 'cpu': N -> N workers total for CPU - 'mps': N -> N workers for MPS (single device) - 'cuda:per': N -> N workers per CUDA index - 'xpu:per': N -> N workers per XPU index - 'npu:per': N -> N workers per NPU index - 'cuda:': N -> override for specific CUDA index - 'xpu:': N -> override for specific XPU index - 'npu:': N -> override for specific NPU index - ':': N -> virtual pool sharing locks with parent device (CUDA/NPU parents must include an explicit index, e.g. 'alias:cuda:0') Unspecified devices default to 1 worker each. gc_debounce_seconds: short wait to coalesce multiple triggers. warmups: optional mapping from device family (e.g. 'cuda') to a `WarmupTask`. Each warmup callable receives `(device, ctx)`, where `ctx` is `WarmUpCtx.DEVICE` or `WarmUpCtx.THREAD`. A task scoped to `WarmUpCtx.THREAD_AND_DEVICE` runs once per physical device with `ctx=DEVICE`, then once per worker thread with `ctx=THREAD`. A special key 'default' applies when no family-specific warmup is found. gc_min_interval_seconds: minimum interval between GC passes. Values <= 0 disable throttling. pin_cpu_workers: bind CPU device workers to individual CPU cores when affinity APIs are available. Defaults to False so CPU tasks may float across cores unless explicitly opt-in. pin_accelerator_workers: bind CUDA/XPU/MPS workers to dedicated CPU cores when affinity APIs are available. Defaults to False so accelerator workers inherit the process CPU affinity mask. """ if devices is None: discovered: List[torch.device] = [] if include_cuda and torch.cuda.is_available(): for i in range(torch.cuda.device_count()): discovered.append(torch.device("cuda", i)) if include_xpu and hasattr(torch, "xpu") and torch.xpu.is_available(): for i in range(torch.xpu.device_count()): discovered.append(torch.device("xpu", i)) if include_npu and HAS_NPU: for i in range(torch.npu.device_count()): discovered.append(torch.device("npu", i)) if include_mps and _mps_available(): discovered.append(torch.device("mps")) if include_cpu: discovered.append(torch.device("cpu")) devices = discovered # Locks and device registry (keyed by "type[:index]" strings like 'cuda:0'). self._locks: Dict[str, _RWLock] = {} self._devices_by_key: Dict[str, torch.device] = {} # Worker groups and RR dispatch bookkeeping. self._worker_groups: Dict[str, List[_DeviceWorker]] = {} self._dispatch_rr: Dict[str, int] = {} self._dispatch_lock = threading.Lock() self._serial_workers: Dict[str, _DeviceWorker] = {} # Stats / GC / inflight control self._stats_lock = threading.Lock() self._per_device_done: Dict[str, int] = {} self._total_done: int = 0 self._empty_cache_every_n = int(empty_cache_every_n) self._gc_event = threading.Event() self._stop_event = threading.Event() self._janitor: Optional[threading.Thread] = None # Auto-GC disable tracking (allows latency-sensitive regions to pause janitor) self._auto_gc_disable_count = 0 self._auto_gc_disable_cv = threading.Condition() # In-flight (scheduled but not finished) counters + per-device CVs. # Each device has a condition variable to let wait() callers block # until inflight hits zero for that device scope. self._inflight: Dict[str, int] = {} self._inflight_cv: Dict[str, threading.Condition] = {} # GC dedupe/coalesce: debounce window to absorb bursty triggers; # per-device "done" watermark to skip redundant GC passes. self._gc_debounce_s = float(gc_debounce_seconds) self._gc_min_interval_s = max(0.0, float(gc_min_interval_seconds)) self._last_gc_done_per_device: Dict[str, int] = {} # Physical-device GC bookkeeping (per accelerator index). self._gc_done_physical: Dict[str, int] = {} self._last_gc_done_physical: Dict[str, int] = {} self._gc_pending_physical: Dict[str, int] = {} self._physical_children: Dict[str, Set[str]] = {} # Device-SMI handles are created lazily for GC logging, while their # synchronization state is an invariant of every pool instance. self._device_smi_lock = threading.Lock() self._device_smi_handles: Dict[str, Any] = {} self._device_smi_failures: Set[str] = set() self._inference_mode = bool(inference_mode) self._worker_warmups = self._normalize_worker_warmups(warmups) self._device_warmup_lock = threading.Lock() self._device_warmup_states: Dict[str, _WorkerWarmupState] = {} workers_cfg = workers or {} base_workers: Dict[str, int] = {} virtual_workers: Dict[str, Tuple[str, int]] = {} for raw_key, raw_count in workers_cfg.items(): alias_info = self._parse_virtual_worker_key(raw_key) try: count = int(raw_count) except (TypeError, ValueError) as exc: raise ValueError(f"Worker count for '{raw_key}' must be an integer (got {raw_count!r})") from exc if alias_info is None: base_workers[raw_key] = count else: _, parent_key = alias_info if parent_key.startswith("cuda") or parent_key.startswith("npu"): parts = parent_key.split(":") if len(parts) < 2 or not parts[1].isdigit(): raise ValueError( f"Virtual pool '{raw_key}' must target a concrete accelerator index (e.g. 'alias:cuda:0'); got '{parent_key}'" ) virtual_workers[raw_key] = (parent_key, count) self._virtual_to_parent: Dict[str, str] = { v_key: parent for v_key, (parent, _) in virtual_workers.items() } # Pre-compute CPU core assignments (when opted in) so GPU/XPU workers # do not collide with CPU workers. When affinity APIs are not available # we silently fall back to OS scheduling. affinity_plan = self._plan_worker_affinity( devices, base_workers, pin_cpu_workers=pin_cpu_workers, pin_accelerator_workers=pin_accelerator_workers, ) # Eagerly build workers, locks, inflight tracking, and counters. for d in devices: dev = _coerce_device(d) if dev.type not in ("cuda", "xpu", "npu", "mps", "cpu"): continue key = self._key(dev) if key in self._devices_by_key: continue self._devices_by_key[key] = dev self._locks[key] = _RWLock() self._per_device_done[key] = 0 self._inflight[key] = 0 self._inflight_cv[key] = threading.Condition() self._last_gc_done_per_device[key] = 0 self._physical_children[key] = {key} if dev.type in ("cuda", "xpu", "npu", "mps"): self._gc_done_physical[key] = 0 self._last_gc_done_physical[key] = 0 n_workers = self._resolve_workers_for_device(dev, base_workers) group: List[_DeviceWorker] = [] for wid in range(int(max(1, n_workers))): cpu_core = affinity_plan.get((key, wid)) worker = self._spawn_worker(dev, key, name=f"DPWorker-{key}#{wid}", cpu_core=cpu_core) group.append(worker) self._worker_groups[key] = group self._dispatch_rr[key] = 0 if group: self._serial_workers[key] = group[0] for v_key, (parent_key, limit) in virtual_workers.items(): if limit < 1: raise ValueError(f"Virtual pool '{v_key}' requires at least one worker (got {limit})") parent_dev = self._devices_by_key.get(parent_key) if parent_dev is None: raise ValueError(f"Virtual pool '{v_key}' references unknown parent '{parent_key}'") if parent_dev.type in ("cuda", "npu") and parent_dev.index is None: raise ValueError( f"Virtual pool '{v_key}' requires an indexed accelerator parent device (e.g. '{v_key}:cuda:0'); got '{parent_key}'" ) parent_group = self._worker_groups.get(parent_key, []) parent_budget = len(parent_group) if limit > parent_budget: raise ValueError( f"Virtual pool '{v_key}' requests {limit} workers but parent '{parent_key}' only has {parent_budget}" ) self._locks[v_key] = self._locks[parent_key] self._devices_by_key[v_key] = parent_dev self._per_device_done[v_key] = 0 self._inflight[v_key] = 0 self._inflight_cv[v_key] = threading.Condition() self._last_gc_done_per_device[v_key] = 0 self._physical_children.setdefault(parent_key, set()).add(v_key) alias_group: List[_DeviceWorker] = [] for wid in range(limit): worker = self._spawn_worker(parent_dev, v_key, name=f"DPWorker-{v_key}#{wid}") alias_group.append(worker) self._worker_groups[v_key] = alias_group self._dispatch_rr[v_key] = 0 if alias_group: self._serial_workers[v_key] = alias_group[0] # A canonical ordering for multi-device lock acquisitions. self._ordered_keys = sorted(self._locks.keys()) self._rebuild_family_keys() # GC diagnostics counters self._gc_passes = 0 self._last_gc_ts: Optional[float] = None self._gc_generation: int = 0 self._last_consumed_gc_generation: int = 0 # Start janitor if enabled and there exists at least one accelerator. if self._empty_cache_every_n > 0 and any( self._devices_by_key[k].type in ("cuda", "xpu", "npu", "mps") for k in self._ordered_keys ): self._janitor = threading.Thread( target=self._janitor_loop, name="DP-Janitor", daemon=True ) self._janitor.start() if DEBUG_ON: log.debug( f"DP-Janitor thread started (debounce={self._gc_debounce_s:.3f}s, " f"min_interval={self._gc_min_interval_s:.3f}s, threshold={self._empty_cache_every_n})" ) else: if DEBUG_ON: log.debug("DP-Janitor disabled (no accelerators or threshold <= 0)") # --------------- Worker management --------------- @staticmethod def _parse_virtual_worker_key(key: str) -> Optional[Tuple[str, str]]: if not isinstance(key, str) or ":" not in key: return None head, tail = key.split(":", 1) if head in {"cuda", "xpu", "npu", "mps", "cpu"}: return None if not tail: return None return head, tail def _rebuild_family_keys(self) -> None: fam_map: Dict[str, List[str]] = {} for key in sorted(self._locks.keys()): fam = key.split(":", 1)[0] fam_map.setdefault(fam, []).append(key) for alias_key, parent_key in self._virtual_to_parent.items(): parent_fam = parent_key.split(":", 1)[0] fam_map.setdefault(parent_fam, []) if alias_key not in fam_map[parent_fam]: fam_map[parent_fam].append(alias_key) for fam, keys in fam_map.items(): keys.sort() self._family_keys = fam_map def _resolve_device_key(self, device: DeviceLike | str) -> str: if isinstance(device, str): if device == "all": raise ValueError("'all' is not a valid concrete device specification") if device in self._locks: return device fam_keys = self._family_keys.get(device) if fam_keys is not None: if len(fam_keys) == 1: return fam_keys[0] raise ValueError( f"Device specification '{device}' is ambiguous; choose one of {fam_keys}" ) dev = _coerce_device(device) key = self._key(dev) if key not in self._locks: raise ValueError(f"Device not in pool: {device}") return key def _plan_worker_affinity( self, devices: Iterable[torch.device], worker_table: Dict[str, int], *, pin_cpu_workers: bool, pin_accelerator_workers: bool, ) -> Dict[Tuple[str, int], Optional[int]]: """Return per-worker CPU core targets to minimise contention.""" if not hasattr(os, "sched_getaffinity") or not hasattr(os, "sched_setaffinity"): return {} try: available_cores = sorted(os.sched_getaffinity(0)) except AttributeError: return {} except OSError as exc: log.warn( "Thread pinning disabled: unable to read CPU affinity mask (%s).", exc, ) return {} if not available_cores: return {} worker_specs: List[Tuple[str, str, int]] = [] for raw_dev in devices: dev = _coerce_device(raw_dev) if dev.type not in ("cuda", "xpu", "npu", "mps", "cpu"): continue if dev.type == "cpu" and not pin_cpu_workers: continue if dev.type in ("cuda", "xpu", "npu", "mps") and not pin_accelerator_workers: continue key = self._key(dev) n_workers = int(max(1, self._resolve_workers_for_device(dev, worker_table))) for wid in range(n_workers): worker_specs.append((dev.type, key, wid)) if not worker_specs: return {} def _priority(dev_type: str) -> int: if dev_type in ("cuda", "xpu", "npu"): return 0 if dev_type == "mps": return 1 return 2 # cpu last so it yields remaining cores worker_specs.sort(key=lambda spec: (_priority(spec[0]), spec[1], spec[2])) core_iter = iter(available_cores) plan: Dict[Tuple[str, int], Optional[int]] = {} exhaustion_logged = False for dev_type, key, wid in worker_specs: try: core = next(core_iter) except StopIteration: core = None if core is None: if not exhaustion_logged: log.warn( "Thread pinning: insufficient CPU cores for %d workers; remaining workers will use OS scheduling.", len(worker_specs) - len(plan), ) exhaustion_logged = True plan[(key, wid)] = None else: plan[(key, wid)] = core return plan @staticmethod def _normalize_worker_warmups( warmups: Optional[Dict[str, WarmupTask]], ) -> Optional[Dict[str, WarmupTask]]: if not warmups: return None normalized: Dict[str, WarmupTask] = {} for raw_key, raw_task in warmups.items(): if not isinstance(raw_task, WarmupTask): raise TypeError( "warmups entries must be WarmupTask instances; " f"got {type(raw_task).__name__!s} for key '{raw_key}'" ) normalized[str(raw_key).lower()] = raw_task return normalized def _resolve_worker_warmup( self, dev: torch.device, key: str, ) -> Tuple[_WorkerWarmupState, ...]: mapping = self._worker_warmups if not mapping: return () family = dev.type.lower() warmup_task = mapping.get(family) primary_key = key.split(":", 1)[0].lower() if warmup_task is None and primary_key in mapping: warmup_task = mapping[primary_key] if warmup_task is None: warmup_task = mapping.get("default") if warmup_task is None: return () states: List[_WorkerWarmupState] = [] physical_key = self._virtual_to_parent.get(key, key) for warmup_ctx in warmup_task.scope.invocation_order(): if warmup_ctx is WarmUpCtx.THREAD: states.append(_WorkerWarmupState(warmup_task, warmup_ctx)) continue with self._device_warmup_lock: state = self._device_warmup_states.get(physical_key) if state is None: state = _WorkerWarmupState(warmup_task, warmup_ctx) self._device_warmup_states[physical_key] = state states.append(state) return tuple(states) def _spawn_worker( self, dev: torch.device, key: str, name: Optional[str] = None, cpu_core: Optional[int] = None, ) -> _DeviceWorker: """ Create and start a worker bound to the provided device. """ warmup_plan = self._resolve_worker_warmup(dev, key) w = _DeviceWorker( device=dev, rwlock=self._locks[key], on_task_finished=self._on_task_finished, on_worker_exit=self._on_worker_exit, name=name, inference_mode=self._inference_mode, cpu_core=cpu_core, warmup_plan=warmup_plan, key_override=key, ) return w def _on_worker_exit(self, key: str, worker: _DeviceWorker) -> None: """ Clean up worker bookkeeping after a thread exits. """ with self._dispatch_lock: group = self._worker_groups.get(key, []) if worker in group: group.remove(worker) self._worker_groups[key] = group if group: self._dispatch_rr[key] %= len(group) else: self._dispatch_rr[key] = 0 self._refresh_serial_worker_locked(key) if DEBUG_ON: log.debug(f"Worker '{worker.name}' exited for {key}") # --------------- Public Work API --------------- def submit( self, device: DeviceLike, fn: Callable[..., Any], /, *args, cuda_event: Optional[torch.cuda.Event] = None, **kwargs, ) -> Future: """ Asynchronously schedule work on the given device; returns a Future. Optional (CUDA): pass `cuda_event=` to ensure work waits for the recorded event. """ key = self._resolve_device_key(device) dev = self._devices_by_key[key] worker = self._pick_worker(key) if cuda_event is not None and dev.type != "cuda": raise ValueError("cuda_event is only valid for CUDA devices") if DEBUG_ON: log.debug(f"submit: device={key} fn={getattr(fn, '__name__', repr(fn))}") # Mark in-flight before enqueue to avoid races with wait(). self._mark_scheduled(key) try: return worker.submit(capture_device_telemetry(fn), *args, cuda_event=cuda_event, **kwargs) except BaseException: # Roll back inflight if enqueue fails (rare) self._mark_finished(key) raise def submit_serial( self, device: DeviceLike, fn: Callable[..., Any], /, *args, cuda_event: Optional[torch.cuda.Event] = None, **kwargs, ) -> Future: """ Schedule work that must execute sequentially on a device. Tasks are enqueued onto a dedicated worker so they run in submission order. """ key = self._resolve_device_key(device) dev = self._devices_by_key[key] if cuda_event is not None and dev.type != "cuda": raise ValueError("cuda_event is only valid for CUDA devices") with self._dispatch_lock: group = self._worker_groups.get(key) if group is None: group = [] self._worker_groups[key] = group if key not in self._dispatch_rr: self._dispatch_rr[key] = 0 if not group: fresh = self._spawn_worker(dev, key, name=f"DPWorker-{key}#0") group.append(fresh) self._dispatch_rr[key] = 0 self._refresh_serial_worker_locked(key) worker = self._serial_workers.get(key) if worker is None: raise ValueError(f"No serial worker available for device '{key}'") if DEBUG_ON: log.debug(f"submit_serial: device={key} fn={getattr(fn, '__name__', repr(fn))}") self._mark_scheduled(key) try: return worker.submit(capture_device_telemetry(fn), *args, cuda_event=cuda_event, **kwargs) except BaseException: self._mark_finished(key) raise def do( self, device: DeviceLike, fn: Callable[..., Any], /, *args, cuda_event: Optional[torch.cuda.Event] = None, **kwargs, ) -> Any: """ Synchronously schedule work and block for the result. """ fut = self.submit(device, fn, *args, cuda_event=cuda_event, **kwargs) return fut.result() def shutdown(self, wait: bool = True): """ Gracefully stop all workers and janitor. IMPORTANT: We snapshot groups before stopping/joining to avoid mutating the lists while iterating (workers remove themselves on exit). """ self._stop_event.set() self._gc_event.set() # wake janitor if waiting # Take stable snapshots under the dispatch lock. with self._dispatch_lock: group_snapshots: Dict[str, List[_DeviceWorker]] = { key: list(group) for key, group in self._worker_groups.items() } # Stop janitor first so it won't grab locks while workers wind down. if self._janitor is not None and wait: if DEBUG_ON: log.debug("Joining DP-Janitor thread…") self._janitor.join() # Issue stop to every worker from the snapshots (no mutation hazards). for key, snapshot in group_snapshots.items(): for w in snapshot: w.stop() # Join everyone if requested. if wait: for key, snapshot in group_snapshots.items(): for w in snapshot: w.join() with self._device_smi_lock: for handle in list(self._device_smi_handles.values()): try: handle.close() except Exception as exc: if DEBUG_ON: log.debug("DeviceThreadPool: ignoring Device-SMI handle close failure during shutdown: %s", exc) self._device_smi_handles.clear() if DEBUG_ON: log.debug("DeviceThreadPool shutdown complete") @contextlib.contextmanager def no_auto_gc(self): """ Temporarily disable automatic empty-cache passes. Useful for latency-sensitive critical sections (e.g., forwarding) where janitor interference is undesirable. """ with self._auto_gc_disable_cv: self._auto_gc_disable_count += 1 try: yield finally: should_signal = False with self._auto_gc_disable_cv: if self._auto_gc_disable_count > 0: self._auto_gc_disable_count -= 1 if self._auto_gc_disable_count == 0: should_signal = True self._auto_gc_disable_cv.notify_all() if should_signal: # Wake janitor in case a trigger is pending. self._gc_event.set() # --------------- Public Lock API --------------- def device_lock(self, device: DeviceLike): """ Obtain an exclusive lock for a single device (blocks all its workers). """ key = self._resolve_device_key(device) lk = self._locks.get(key) if lk is None: raise ValueError(f"Unknown device for pool: {device}") return lk.writer() def read_lock(self, device: DeviceLike | str): """ Obtain a read/shared lock. Accepts: - Concrete device: torch.device('cuda:0'), 'cuda:1' - Family device: 'cuda', 'xpu', 'mps', 'cpu' - 'all' for every device in the pool Returns a context manager. """ # Family string shortcut (e.g., "cuda" or "all") if isinstance(device, str): if device == "all": pairs = [(k, self._locks[k]) for k in self._ordered_keys] return _ReadLockGroup(pairs) fam_keys = self._family_keys.get(device) if fam_keys is not None: pairs = [(k, self._locks[k]) for k in fam_keys] return _ReadLockGroup(pairs) if device in self._locks: return self._locks[device].reader() raise ValueError(f"Unknown device for pool: {device}") # torch.device / int / 'cuda:0' etc. dev = _coerce_device(device) key = self._key(dev) # Family device with index=None -> all devices of that type if dev.index is None: fam_keys = self._family_keys.get(key) if not fam_keys: raise ValueError(f"No devices of type '{dev.type}' in pool") pairs = [(k, self._locks[k]) for k in fam_keys] return _ReadLockGroup(pairs) # Concrete device lk = self._locks.get(key) if lk is None: raise ValueError(f"Unknown device for pool: {dev}") return lk.reader() def lock(self, devices: Optional[Iterable[DeviceLike]] = None): """ Exclusive lock across multiple devices (default: all pool devices). Acquires each device's write lock in canonical order to avoid deadlocks. """ if devices is None: pairs = [(k, self._locks[k]) for k in self._ordered_keys] else: keys = sorted(self._normalize_scope_to_keys(devices)) pairs = [(k, self._locks[k]) for k in keys] return _LockGroup(pairs) # --------------- Public Wait API --------------- # The wait() primitive blocks until inflight work drains for a scope. # With lock=True, it returns a context manager that first drains, then # acquires writer locks—handy for "drain-and-free" sequences. def wait( self, scope: Optional[Union[str, DeviceLike, Iterable[DeviceLike]]] = None, *, lock: bool = False, ) -> None | _WaitAndLock: """ Wait until in-flight tasks for `scope` drain to zero. scope: - None or 'all' -> all devices - 'cuda' | 'xpu' | 'mps' | 'cpu' -> all devices of that type - 'cuda:0' | 'xpu:1' -> specific device key - torch.device or iterable of the above lock: - False: block until drained. - True: return a context manager that drains then acquires writer locks. """ keys = self._resolve_scope_to_keys(scope) if lock: pairs = [(k, self._locks[k]) for k in sorted(keys)] class _WaitThenLock(contextlib.AbstractContextManager): """ Drain inflight for the given keys, then acquire writer locks in canonical order; release on exit. """ def __init__(self, outer: DeviceThreadPool, pairs_local: List[tuple[str, _RWLock]], keys_local: List[str]): self._outer = outer self._pairs = pairs_local self._keys = keys_local self._group = _LockGroup(pairs_local) def __enter__(self): if DEBUG_ON: log.debug(f"wait(lock=True) drain start: keys={self._keys}") for kk in self._keys: cv = self._outer._inflight_cv[kk] with cv: while self._outer._inflight[kk] > 0: if DEBUG_ON: log.debug(f"wait(lock=True) blocking on inflight[{kk}]={self._outer._inflight[kk]}") cv.wait() if DEBUG_ON: log.debug(f"wait(lock=True) acquire writer locks: keys={self._keys}") return self._group.__enter__() def __exit__(self, exc_type, exc, tb): if DEBUG_ON: log.debug(f"wait(lock=True) releasing writer locks: keys={self._keys}") return self._group.__exit__(exc_type, exc, tb) return _WaitThenLock(self, pairs, keys) # Simple drain (no locks) if DEBUG_ON: log.debug(f"wait(lock=False) drain start: keys={keys}") for k in keys: cv = self._inflight_cv[k] with cv: while self._inflight[k] > 0: if DEBUG_ON: log.debug(f"wait(lock=False) blocking on inflight[{k}]={self._inflight[k]}") cv.wait() if DEBUG_ON: log.debug(f"wait(lock=False) drain done: keys={keys}") return None # --------------- Public Stats API --------------- def stats(self) -> Dict[str, Any]: """ Return a snapshot of counters. Use under tests or ad-hoc diagnostics. """ with self._stats_lock: return { "per_device": dict(self._per_device_done), "total": int(self._total_done), "threshold": int(self._empty_cache_every_n), } def device_completed(self, device: DeviceLike) -> int: """ Convenience accessor for per-device completed count (atomic snapshot). """ key = self._resolve_device_key(device) with self._stats_lock: return int(self._per_device_done.get(key, 0)) def total_completed(self) -> int: """ Convenience accessor for global completed count (atomic snapshot). """ with self._stats_lock: return int(self._total_done) # --------------- Internals --------------- def _key(self, dev: torch.device) -> str: idx = "" if dev.index is None else f":{dev.index}" return f"{dev.type}{idx}" def _pick_worker(self, key: str) -> _DeviceWorker: """ Round-robin selection across available workers for a device key. If no workers exist (should not happen under normal init), we spawn one lazily for robustness. """ with self._dispatch_lock: group = self._worker_groups.get(key) if not group: dev = self._devices_by_key[key] w = self._spawn_worker(dev, key, name=f"DPWorker-{key}#0") group = [w] self._worker_groups[key] = group self._dispatch_rr[key] = 0 self._refresh_serial_worker_locked(key) return w n = len(group) idx = self._dispatch_rr[key] % n w = group[idx] self._dispatch_rr[key] = (idx + 1) % n self._refresh_serial_worker_locked(key) return w def _refresh_serial_worker_locked(self, key: str) -> None: """ Ensure the serial worker mapping references a live worker. Caller must hold self._dispatch_lock. """ group = self._worker_groups.get(key, []) if not group: self._serial_workers.pop(key, None) return current = self._serial_workers.get(key) if current in group: return self._serial_workers[key] = group[0] def _resolve_workers_for_device(self, dev: torch.device, table: Dict[str, int]) -> int: """ Resolve worker count from policy table: - exact key (e.g. 'cuda:0') overrides - family-per (e.g. 'cuda:per') applies to all indices - family singletons ('cpu', 'mps') apply to that single device """ key = self._key(dev) if key in table: return int(table[key]) fam_key = f"{dev.type}:per" if fam_key in table: return int(table[fam_key]) if dev.type in ("cpu", "mps") and dev.type in table: return int(table[dev.type]) return 1 def _normalize_scope_to_keys(self, scope: Iterable[DeviceLike]) -> List[str]: """ Normalize a scope specification into a sorted list of device keys. Accepts strings ('cuda', 'cuda:0', 'all'), ints (device indices), and torch.device objects. """ keys: List[str] = [] seen: set[str] = set() for s in scope: if isinstance(s, str): if s == "all": for key in self._ordered_keys: if key not in seen: keys.append(key) seen.add(key) continue fam_keys = self._family_keys.get(s) if fam_keys is not None: for key in fam_keys: if key not in seen: keys.append(key) seen.add(key) continue if s not in self._locks: raise ValueError(f"Unknown device key in scope: {s}") if s not in seen: keys.append(s) seen.add(s) else: dev = _coerce_device(s) k = self._key(dev) if dev.index is None: fam_keys = self._family_keys.get(k) if not fam_keys: raise ValueError(f"No devices of type '{dev.type}' in pool") for key in fam_keys: if key not in seen: keys.append(key) seen.add(key) else: if k not in self._locks: raise ValueError(f"Device not in pool: {dev}") if k not in seen: keys.append(k) seen.add(k) return keys def _resolve_scope_to_keys(self, scope: Optional[Union[str, DeviceLike, Iterable[DeviceLike]]] = None) -> List[str]: """ Helper for wait()/lock(): expand a scope into concrete device keys. """ if scope is None or (isinstance(scope, str) and scope == "all"): return list(self._ordered_keys) if isinstance(scope, (str, torch.device, int)): return self._normalize_scope_to_keys([scope]) return self._normalize_scope_to_keys(scope) def _physical_key(self, key: str) -> str: """ Map a logical worker/alias key back to its physical device key. """ return getattr(self, "_virtual_to_parent", {}).get(key, key) def _run_empty_cache_for_device(self, key: str, dev: torch.device) -> Optional[float]: """ Execute an empty_cache call for the given device. Returns execution time in seconds. """ start = time.time() try: # Tests and runtime hooks may monkeypatch empty_cache between janitor passes. torch_utils.resolve_empty_cache_callable.cache_clear() success = torch_empty_cache_any(device=dev, gc=False) except Exception as exc: if DEBUG_ON: log.debug("DP-Janitor: empty_cache failed for %s (%s)", key, exc) return None if not success and DEBUG_ON: log.debug("DP-Janitor: unsupported device type '%s' for key %s", dev.type, key) return None return time.time() - start @staticmethod def _format_gib_value(value: float) -> str: text = f"{value:.1f}" if text.endswith(".0"): text = text[:-2] return f"{text}G" def _device_smi_identifier(self, dev: torch.device) -> Optional[str]: if Device is None: return None if dev.type == "cuda": idx = dev.index if idx is None: return None prefix = "rocm" if getattr(torch.version, "hip", None) else "cuda" return f"{prefix}:{idx}" if dev.type == "xpu": idx = dev.index if idx is None: return None return f"xpu:{idx}" if dev.type == "npu": idx = dev.index if idx is None: return None return f"npu:{idx}" return None def _query_device_vram_gib(self, key: str) -> Optional[float]: if Device is None: return None dev = self._devices_by_key.get(key) if dev is None: return None identifier = self._device_smi_identifier(dev) if identifier is None: return None with self._device_smi_lock: if identifier in self._device_smi_failures: return None handle = self._device_smi_handles.get(identifier) if handle is None: try: handle = Device(identifier) except Exception: self._device_smi_failures.add(identifier) return None self._device_smi_handles[identifier] = handle try: metrics = handle.metrics(fast=True) except Exception: with self._device_smi_lock: self._device_smi_failures.add(identifier) stored = self._device_smi_handles.pop(identifier, None) if stored is not None: try: stored.close() except Exception: pass return None memory_used = getattr(metrics, "memory_used", None) if memory_used is None: return None return float(memory_used) / (1024 ** 3) def _format_vram_summary(self, physical_keys: Iterable[str]) -> str: readings: List[str] = [] for key in physical_keys: value = self._query_device_vram_gib(key) if value is None: continue readings.append(f"{key}={self._format_gib_value(value)}") return ", ".join(readings) if readings else "n/a" # ---- inflight & completion accounting ---- def _mark_scheduled(self, key: str) -> None: """ Increment inflight for a device key and emit a debug breadcrumb if enabled. """ cv = self._inflight_cv[key] with cv: self._inflight[key] = self._inflight.get(key, 0) + 1 if DEBUG_ON: log.debug(f"inflight[{key}] ++ -> {self._inflight[key]}") def _mark_finished(self, key: str) -> None: """ Decrement inflight for a device key, clamp at zero on underflow, and notify waiters when the device drains. """ cv = self._inflight_cv[key] with cv: new_val = self._inflight.get(key, 0) - 1 if new_val < 0: if DEBUG_ON: log.debug(f"WARNING: inflight[{key}] underflow ({new_val}); clamping to 0") new_val = 0 self._inflight[key] = new_val if DEBUG_ON: log.debug(f"inflight[{key}] -- -> {self._inflight[key]}") if self._inflight[key] == 0: cv.notify_all() def _on_task_finished(self, key: str) -> None: """ Called at the end of every task (success or failure). Updates counters and signals the janitor if the per-device threshold is reached. """ if not hasattr(self, "_gc_done_physical"): self._gc_done_physical = {} if not hasattr(self, "_gc_pending_physical"): self._gc_pending_physical = {} elif not isinstance(self._gc_pending_physical, dict): self._gc_pending_physical = dict.fromkeys(self._gc_pending_physical, 1) if not hasattr(self, "_last_gc_done_physical"): self._last_gc_done_physical = {} if not hasattr(self, "_physical_children"): self._physical_children = {} self._mark_finished(key) trigger_gc = False with self._stats_lock: self._per_device_done[key] = self._per_device_done.get(key, 0) + 1 self._total_done += 1 dev = self._devices_by_key.get(key) if ( dev is not None and self._empty_cache_every_n > 0 and dev.type in ("cuda", "xpu", "npu", "mps") ): physical_key = self._physical_key(key) current = self._gc_done_physical.get(physical_key, 0) + 1 self._gc_done_physical[physical_key] = current last_done = self._last_gc_done_physical.get(physical_key, 0) completed_chunks = max(0, (current - last_done) // self._empty_cache_every_n) pending_map = self._gc_pending_physical already_pending = pending_map.get(physical_key, 0) new_triggers = completed_chunks - already_pending if new_triggers > 0: pending_map[physical_key] = already_pending + new_triggers self._gc_generation += new_triggers trigger_gc = True if DEBUG_ON: log.debug( "GC trigger set by %s (physical=%s): per_physical_done=%d threshold=%d total_done=%d pending=%s", key, physical_key, current, self._empty_cache_every_n, self._total_done, sorted(k for k, count in pending_map.items() if count > 0), ) if trigger_gc: self._gc_event.set() # ---- ANSI table rendering for GC diagnostics ---- def _ansi_table(self, headers: List[str], rows: List[List[str]]) -> str: """ Render a simple ANSI/ASCII table with bold headers. Used only for human-readable diagnostics; not used in hot paths. """ widths = [len(h) for h in headers] for r in rows: for i, cell in enumerate(r): widths[i] = max(widths[i], len(cell)) def hrule(sep_left="+", sep_mid="+", sep_right="+", h="-"): parts = [sep_left] for i, w in enumerate(widths): parts.append(h * (w + 2)) parts.append(sep_mid if i < len(widths) - 1 else sep_right) return "".join(parts) def format_row(cols: List[str]): out = ["|"] for i, cell in enumerate(cols): out.append(" " + cell.ljust(widths[i]) + " ") out.append("|") return "".join(out) BOLD = "\x1b[1m" RESET = "\x1b[0m" top = hrule() mid = hrule(h="=") bot = hrule() lines = [top, format_row([BOLD + h + RESET for h in headers]), mid] for r in rows: lines.append(format_row(r)) lines.append(bot) return "\n".join(lines) def _collect_state_snapshot(self) -> Dict[str, Any]: """ Safely collect a snapshot of pool state for diagnostics and GC decisions. """ with self._stats_lock: per_done = dict(self._per_device_done) total_done = int(self._total_done) threshold = int(self._empty_cache_every_n) inflight: Dict[str, int] = {} for k, cv in self._inflight_cv.items(): with cv: inflight[k] = int(self._inflight[k]) workers = {k: len(self._worker_groups.get(k, [])) for k in self._devices_by_key.keys()} meta: Dict[str, Dict[str, str]] = {} for k, dev in self._devices_by_key.items(): idx = "" if dev.index is None else str(dev.index) meta[k] = {"type": dev.type, "index": idx} physical_children = getattr(self, "_physical_children", {}) per_done_physical: Dict[str, int] = {} for phys_key, members in physical_children.items(): per_done_physical[phys_key] = sum(per_done.get(member, 0) for member in members) raw_pending = getattr(self, "_gc_pending_physical", {}) if isinstance(raw_pending, dict): pending_gc = sorted(k for k, count in raw_pending.items() if count > 0) else: pending_gc = sorted(raw_pending) snap: Dict[str, Any] = { "devices": sorted(self._devices_by_key.keys()), "per_done": per_done, "per_done_physical": per_done_physical, "total_done": total_done, "threshold": threshold, "inflight": inflight, "workers": workers, "meta": meta, "total_inflight": sum(inflight.values()), "total_workers": sum(workers.values()), "gc_passes": int(self._gc_passes), "gc_generation": int(self._gc_generation), "gc_generation_consumed": int(self._last_consumed_gc_generation), "last_gc_ts": self._last_gc_ts, "now": time.time(), "pending_gc": pending_gc, } return snap def _render_gc_table(self, snap: Dict[str, Any]) -> str: """ Pretty-print a GC state table; used only when someone wants to log it. """ headers = [ "Device", "Type", "Index", "Workers", "Inflight", "Done", "Threshold", "NextGC", "Accel" ] rows: List[List[str]] = [] thr = snap["threshold"] for k in snap["devices"]: t = snap["meta"][k]["type"] idx = snap["meta"][k]["index"] w = snap["workers"].get(k, 0) infl = snap["inflight"].get(k, 0) done = snap["per_done"].get(k, 0) accel = "Y" if t in ("cuda", "xpu", "npu", "mps") else "N" if thr > 0 and t in ("cuda", "xpu", "npu", "mps"): rem = thr - (done % thr) if (done % thr) != 0 else 0 nextgc = "now" if rem == 0 and done > 0 else str(rem) else: nextgc = "-" rows.append([k, t, idx, str(w), str(infl), str(done), str(thr) if thr > 0 else "-", nextgc, accel]) table_main = self._ansi_table(headers, rows) totals_headers = ["Total Workers", "Total Inflight", "Total Done", "GC Passes", "Since Last GC (s)"] since = "-" if self._last_gc_ts is None else f"{time.time() - self._last_gc_ts:.3f}" totals_rows = [[ str(sum(len(v) for v in self._worker_groups.values())), str(snap["total_inflight"]), str(snap["total_done"]), str(snap["gc_passes"]), since, ]] table_totals = self._ansi_table(totals_headers, totals_rows) return table_main + "\n" + table_totals # ---- janitor (per-device empty-cache under exclusive writer lock) ---- # The janitor runs in the background. When a device completes N tasks, a trigger # is set. The janitor debounces triggers, takes a snapshot, and if at least one # accelerator device progressed by >= N tasks since the last pass, it iterates # devices, acquiring each device's writer lock before calling empty_cache(). def _synchronize_all(self): """ Optionally ensure devices are idle before empty_cache() to avoid races with outstanding kernels. Keeping this disabled by default for performance. """ # CUDA for key in self._ordered_keys: dev = self._devices_by_key[key] if dev.type != "cuda": continue with torch.cuda.device(dev.index): torch.cuda.synchronize() # XPU for key in self._ordered_keys: dev = self._devices_by_key[key] if dev.type != "xpu": continue if hasattr(torch, "xpu") and hasattr(torch.xpu, "synchronize"): with torch.xpu.device(dev.index): torch.xpu.synchronize() # NPU for key in self._ordered_keys: dev = self._devices_by_key[key] if dev.type != "npu": continue if HAS_NPU: with torch.npu.device(dev.index): torch.npu.synchronize() # MPS has_mps_device = any(self._devices_by_key[k].type == "mps" for k in self._ordered_keys) if has_mps_device and hasattr(torch, "mps") and hasattr(torch.mps, "synchronize"): torch.mps.synchronize() def _should_run_gc_from_snapshot(self, snap: Dict[str, Any]) -> bool: """ Decide whether to run GC by comparing per-device progress since last GC. This deduplicates bursty triggers that occurred before the previous GC ran. """ thr = snap["threshold"] if thr <= 0: return False pending = snap.get("pending_gc") or [] if pending: return True per_done_physical = snap.get("per_done_physical") or {} last_done_physical = getattr(self, "_last_gc_done_physical", {}) for phys_key, done_now in per_done_physical.items(): done_prev = last_done_physical.get(phys_key, 0) if done_now - done_prev >= thr: return True return False def _update_gc_watermarks(self, snap_after: Dict[str, Any]) -> None: """ Record 'done' counters as of a GC pass to require fresh progress before a subsequent pass is allowed. """ int(self._empty_cache_every_n) per_done_physical = snap_after.get("per_done_physical") or {} per_done = snap_after.get("per_done") or {} snap_after.get("meta") or {} processed = snap_after.get("_gc_processed_devices") if processed is None: processed_iter = per_done_physical.keys() else: processed_iter = processed for phys_key in processed_iter: done_phys = per_done_physical.get(phys_key) if done_phys is None: continue self._last_gc_done_physical[phys_key] = done_phys members = self._physical_children.get(phys_key, {phys_key}) for member in members: done_member = per_done.get(member) if done_member is None: continue self._last_gc_done_per_device[member] = done_member def _janitor_loop(self): """ Main janitor loop: - Waits on a trigger (with short timeout to honor shutdowns promptly). - Debounces additional triggers for a brief window. - Takes a snapshot and decides whether to run. - For each accelerator device, acquires its writer lock and calls empty_cache() using the LIVE attribute if callable, otherwise the HARD COPY captured at import time. """ while True: if DEBUG_ON: log.debug("DP-Janitor: waiting for trigger…") self._gc_event.wait() self._gc_event.clear() if self._stop_event.is_set(): if DEBUG_ON: log.debug("DP-Janitor: stop event set; exiting") break # Debounce window: absorb additional triggers before deciding. if self._gc_debounce_s > 0: t_end = time.time() + self._gc_debounce_s if DEBUG_ON: log.debug(f"DP-Janitor: debounce window start ({self._gc_debounce_s:.3f}s)") while time.time() < t_end: if self._stop_event.is_set(): if DEBUG_ON: log.debug("DP-Janitor: stop during debounce; exiting") return self._gc_event.wait(timeout=max(0.0, t_end - time.time())) self._gc_event.clear() if DEBUG_ON: log.debug("DP-Janitor: debounce window end") with self._auto_gc_disable_cv: while self._auto_gc_disable_count > 0 and not self._stop_event.is_set(): if DEBUG_ON: log.debug("DP-Janitor: auto-GC disabled; waiting…") self._auto_gc_disable_cv.wait() if self._stop_event.is_set(): if DEBUG_ON: log.debug("DP-Janitor: stop event set during auto-GC wait; exiting") break with self._stats_lock: current_generation = self._gc_generation last_generation = self._last_consumed_gc_generation last_gc_ts = self._last_gc_ts if current_generation == last_generation: if DEBUG_ON: log.debug("DP-Janitor: trigger generation already consumed; skipping") continue min_interval = self._gc_min_interval_s if min_interval > 0.0 and last_gc_ts is not None: elapsed = time.time() - last_gc_ts if elapsed < min_interval: wait_for = min_interval - elapsed if DEBUG_ON: log.debug( f"DP-Janitor: last pass {elapsed * 1000:.1f}ms ago; waiting {wait_for * 1000:.1f}ms to honor min interval" ) if self._stop_event.wait(timeout=wait_for): if DEBUG_ON: log.debug("DP-Janitor: stop event set during min-interval wait; exiting") break if self._stop_event.is_set(): if DEBUG_ON: log.debug("DP-Janitor: stop event observed after min-interval wait; exiting") break with self._stats_lock: current_generation = self._gc_generation last_generation = self._last_consumed_gc_generation if current_generation == last_generation: if DEBUG_ON: log.debug("DP-Janitor: no pending GC generation after min-interval wait; skipping") continue # Snapshot & decision try: pre = self._collect_state_snapshot() if DEBUG_ON: log.debug(f"DP-Janitor: pre-snapshot taken: total_done={pre['total_done']}, threshold={pre['threshold']}, inflight={pre['inflight']}") log.debug("GC trigger received; evaluating whether to run…") except Exception as e: # Fallback snapshot (unlikely path; logging should not crash janitor) try: log.warn(f"Failed to render GC pre-snapshot: {e!r}") except Exception: pass pre = { "devices": list(self._devices_by_key.keys()), "per_done": {k: self._per_device_done.get(k, 0) for k in self._devices_by_key.keys()}, "threshold": self._empty_cache_every_n, "meta": {k: {"type": self._devices_by_key[k].type} for k in self._devices_by_key.keys()}, "inflight": dict.fromkeys(self._devices_by_key.keys(), 0), "workers": {k: len(self._worker_groups.get(k, [])) for k in self._devices_by_key.keys()}, "total_inflight": 0, "total_workers": sum(len(v) for v in self._worker_groups.values()), "gc_passes": self._gc_passes, "last_gc_ts": self._last_gc_ts, "now": time.time(), "total_done": self._total_done, } if not self._should_run_gc_from_snapshot(pre): if DEBUG_ON: log.debug("DP-Janitor: skip GC (no device progressed by threshold since last pass)") with self._stats_lock: self._last_consumed_gc_generation = current_generation continue t0 = time.time() try: pre = self._collect_state_snapshot() if DEBUG_ON: log.debug( "DP-Janitor: pre-snapshot taken: total_done=%s threshold=%s inflight=%s pending=%s", pre["total_done"], pre["threshold"], pre["inflight"], pre.get("pending_gc"), ) log.debug("GC trigger received; evaluating whether to run…") pending_targets = [k for k in pre.get("pending_gc", []) if k in self._locks] except Exception as e: try: log.warn(f"Failed to render GC pre-snapshot: {e!r}") except Exception: pass raw_pending = getattr(self, "_gc_pending_physical", {}) if isinstance(raw_pending, dict): pending_targets = sorted(k for k, count in raw_pending.items() if count > 0) else: pending_targets = sorted(raw_pending) if not pending_targets: if DEBUG_ON: log.debug("DP-Janitor: no pending devices after snapshot; marking generation %d consumed", current_generation) with self._stats_lock: self._last_consumed_gc_generation = max(self._last_consumed_gc_generation, current_generation) continue # A single trigger runs a full accelerator sweep so memory pressure is # normalized across all active devices under the janitor's exclusive pass. sweep_targets: List[str] = [] seen_targets: Set[str] = set() for candidate in self._ordered_keys: physical = self._physical_key(candidate) if physical in seen_targets: continue dev = self._devices_by_key.get(physical) if dev is None or dev.type not in ("cuda", "xpu", "npu", "mps"): continue if physical not in self._locks: continue seen_targets.add(physical) sweep_targets.append(physical) if not sweep_targets: sweep_targets = pending_targets processed_devices: List[str] = [] skipped_devices: List[str] = [] per_device_durations: Dict[str, float] = {} for key in sweep_targets: dev = self._devices_by_key.get(key) if dev is None or dev.type not in ("cuda", "xpu", "npu", "mps"): skipped_devices.append(key) continue lk = self._locks.get(key) if lk is None: skipped_devices.append(key) continue if DEBUG_ON: log.debug("DP-Janitor: attempting writer lock for %s", key) with lk.writer(): if DEBUG_ON: log.debug("DP-Janitor: acquired writer lock for %s", key) duration = self._run_empty_cache_for_device(key, dev) if duration is not None: per_device_durations[key] = duration processed_devices.append(key) if not processed_devices and DEBUG_ON: log.debug("DP-Janitor: no eligible accelerator devices found in pending=%s", pending_targets) t1 = time.time() prev_gc_ts = self._last_gc_ts if processed_devices: self._gc_passes += 1 self._last_gc_ts = t1 gc_timestamp = datetime.fromtimestamp(t1, tz=timezone.utc).isoformat() if prev_gc_ts is None: since_last_gc = "since last GC: n/a" else: delta_s = t1 - prev_gc_ts since_last_gc = f"since last GC: {delta_s:.3f}s ({delta_s * 1000:.1f}ms)" post = None if processed_devices: try: post = self._collect_state_snapshot() post["_gc_processed_devices"] = processed_devices self._update_gc_watermarks(post) except Exception as e: try: log.warn(f"Failed to update GC watermarks: {e!r}") except Exception as log_exc: _best_effort_stderr_write( f"Failed to update GC watermarks: {e!r}; secondary logging failed: {log_exc!r}" ) with self._stats_lock: pending_map = self._gc_pending_physical if not isinstance(pending_map, dict): pending_map = dict.fromkeys(pending_map, 1) self._gc_pending_physical = pending_map for key in processed_devices: pending_map.pop(key, None) for key in skipped_devices: pending_map.pop(key, None) self._last_consumed_gc_generation = max(self._last_consumed_gc_generation, current_generation) if any(count > 0 for count in pending_map.values()): self._gc_event.set() if processed_devices: vram_summary = self._format_vram_summary(processed_devices) try: devices_clause = ", ".join(processed_devices) log.info( f"GC completed in {t1 - t0:.3f}s (pass #{self._gc_passes}) at {gc_timestamp}; devices={devices_clause}; VRAM {vram_summary}; {since_last_gc}." ) if DEBUG_ON: log.debug( "DP-Janitor: post-snapshot inflight=%s per_done=%s per_done_physical=%s durations=%s", post["inflight"], post["per_done"], post.get("per_done_physical"), per_device_durations, ) except Exception as e: try: log.warn(f"Failed to render GC post-snapshot: {e!r}") except Exception as log_exc: _best_effort_stderr_write( f"Failed to render GC post-snapshot: {e!r}; secondary logging failed: {log_exc!r}" ) def _empty_all_caches(self): torch_empty_cache_any(gc=False)