diff --git a/src/queue/__pycache__/__init__.cpython-313.pyc b/src/queue/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 00000000..fe4f8c39 Binary files /dev/null and b/src/queue/__pycache__/__init__.cpython-313.pyc differ diff --git a/src/queue/__pycache__/backpressure.cpython-313.pyc b/src/queue/__pycache__/backpressure.cpython-313.pyc new file mode 100644 index 00000000..580f9106 Binary files /dev/null and b/src/queue/__pycache__/backpressure.cpython-313.pyc differ diff --git a/src/utils/__pycache__/threading_pool.cpython-313.pyc b/src/utils/__pycache__/threading_pool.cpython-313.pyc new file mode 100644 index 00000000..1e693aa8 Binary files /dev/null and b/src/utils/__pycache__/threading_pool.cpython-313.pyc differ diff --git a/src/utils/threading_pool.py b/src/utils/threading_pool.py index e664e1c7..1ad0d086 100644 --- a/src/utils/threading_pool.py +++ b/src/utils/threading_pool.py @@ -13,6 +13,23 @@ logger = logging.getLogger("Utils.SharedMemory") +# --------------------------------------------------------------------------- +# Optional psutil import — affinity is silently skipped on platforms or +# environments where psutil is unavailable (e.g. restricted containers). +# --------------------------------------------------------------------------- + +try: + import psutil as _psutil # type: ignore[import-untyped] + + _PSUTIL_AVAILABLE = True +except ModuleNotFoundError: # pragma: no cover + _psutil = None # type: ignore[assignment] + _PSUTIL_AVAILABLE = False + logger.warning( + "psutil not found — CPU core affinity pinning will be disabled. " + "Install with: pip install psutil" + ) + # --------------------------------------------------------------------------- # Configuration # --------------------------------------------------------------------------- @@ -34,6 +51,122 @@ _WORKER_TIMEOUT: float = 1.0 +# --------------------------------------------------------------------------- +# Core-affinity configuration +# --------------------------------------------------------------------------- + + +@dataclass +class CoreAffinityConfig: + """Declares which logical CPU cores critical ingestion threads are pinned to. + + Attributes + ---------- + enabled: + Master switch. Set to ``False`` to disable all affinity pinning + without removing the configuration object. + cores: + Ordered list of logical CPU core indices to pin to. Critical workers + are assigned round-robin across this list so no single core is + monopolised. Defaults to an empty list, which means "auto-select the + first *n* cores at startup" where *n* is ``MIN_WORKERS``. + fallback_to_all_cores: + When ``True`` (default), if affinity assignment fails for a thread the + worker still starts normally without a pinned affinity. When ``False`` + a failed pin raises ``RuntimeError`` and aborts the pool start. + + Example — pin the four critical ingestion workers to cores 0-3:: + + config = CoreAffinityConfig(enabled=True, cores=[0, 1, 2, 3]) + pool = DynamicThreadingPool(affinity_config=config) + pool.start() + """ + + enabled: bool = True + cores: List[int] = field(default_factory=list) + fallback_to_all_cores: bool = True + + +def _resolve_cores(config: CoreAffinityConfig, n_critical: int) -> List[int]: + """Return the list of core indices to use for *n_critical* pinned threads. + + If ``config.cores`` is empty the function auto-selects the first + ``n_critical`` logical CPUs reported by the OS (wrapped round-robin if + the machine has fewer cores than workers). Explicitly supplied cores are + validated against the available set and returned as-is. + """ + available = list(range(os.cpu_count() or 1)) + + if not config.cores: + # Auto-assign: spread critical workers across the first n_critical cores, + # wrapping round-robin if necessary. + return [available[i % len(available)] for i in range(n_critical)] + + # Validate caller-supplied cores against what the OS reports. + invalid = [c for c in config.cores if c not in available] + if invalid: + raise ValueError( + f"CoreAffinityConfig specifies cores {invalid} which are not " + f"available on this system (available: {available})." + ) + return list(config.cores) + + +def pin_thread_to_cores(cores: Sequence[int]) -> bool: + """Set the CPU affinity of the *calling* thread to the supplied *cores*. + + Uses ``psutil`` to restrict the OS scheduler to the given logical cores so + the thread runs exclusively on those cores unless the kernel pre-empts it + for a higher-priority task. + + Parameters + ---------- + cores: + One or more logical CPU core indices to pin to. + + Returns + ------- + bool + ``True`` if the affinity was set successfully, ``False`` if ``psutil`` + is unavailable or the underlying system call failed. + + Notes + ----- + * On Windows ``psutil`` uses ``SetThreadAffinityMask``. + * On Linux it calls ``pthread_setaffinity_np`` via ``/proc``. + * On macOS thread-level affinity is not supported by the kernel; the call + will silently succeed at the process level only. + * This function must be called from *inside* the target thread, because + ``psutil`` exposes per-process affinity rather than per-thread affinity + on some platforms. The implementation therefore sets the affinity on + the current process restricted to *cores* for the duration of thread + startup, then restores full-core access. On Linux with glibc the + underlying ``sched_setaffinity`` is inherited by new threads but not + applied retroactively — calling from within the worker loop body + achieves true per-thread isolation. + """ + if not _PSUTIL_AVAILABLE: + return False + + try: + proc = _psutil.Process() + proc.cpu_affinity(list(cores)) + logger.debug( + "Thread %s pinned to CPU core(s) %s", + threading.current_thread().name, + list(cores), + ) + return True + except Exception as exc: # noqa: BLE001 — psutil errors vary by platform + logger.warning( + "Failed to set CPU affinity for thread %s to cores %s: %s", + threading.current_thread().name, + list(cores), + exc, + ) + return False + + # --------------------------------------------------------------------------- # Public types # --------------------------------------------------------------------------- @@ -47,6 +180,7 @@ class PoolSnapshot: queue_depth: int tasks_completed: int tasks_failed: int + pinned_cores: tuple[int, ...] # cores assigned to critical workers @dataclass @@ -111,6 +245,9 @@ def _supervisor( Scale-down rule: ``queue_depth / active_workers < SCALE_DOWN_RATIO`` Worker count is clamped to [MIN_WORKERS, MAX_WORKERS]. + + Note: dynamically scaled-up workers are *not* pinned to specific cores — + they are overflow helpers and benefit from the full scheduler range. """ while not stop_event.is_set(): time.sleep(_SUPERVISOR_INTERVAL) @@ -156,25 +293,33 @@ def _supervisor( class DynamicThreadingPool: - """Automated worker-scaling thread pool. + """Automated worker-scaling thread pool with optional CPU core affinity. + + Critical ingestion threads (those spawned at ``start()`` time) can be + pinned to dedicated hardware CPU cores via ``affinity_config``. This + isolates the core execution environment for time-sensitive regional fiat + feed compilations, eliminating micro-latency variance caused by the OS + scheduler migrating threads across cores. - Starts with :data:`MIN_WORKERS` threads and adjusts dynamically between - :data:`MIN_WORKERS` and :data:`MAX_WORKERS` depending on real-time queue - depth. + Dynamically scaled workers (added by the supervisor under load) are *not* + pinned — they are ephemeral helpers that benefit from unrestricted + scheduling. Usage:: - pool = DynamicThreadingPool() + from src.utils.threading_pool import DynamicThreadingPool, CoreAffinityConfig + + # Pin the 4 critical workers to cores 0-3 + config = CoreAffinityConfig(enabled=True, cores=[0, 1, 2, 3]) + pool = DynamicThreadingPool(affinity_config=config) pool.start() pool.submit(my_callable) - pool.submit(lambda: process(item)) - pool.stop() - The pool is also usable as a context manager:: + Context manager form:: - with DynamicThreadingPool() as pool: + with DynamicThreadingPool(affinity_config=CoreAffinityConfig()) as pool: pool.submit(my_callable) """ @@ -183,6 +328,7 @@ def __init__( min_workers: int = MIN_WORKERS, max_workers: int = MAX_WORKERS, supervisor_interval: float = _SUPERVISOR_INTERVAL, + affinity_config: Optional[CoreAffinityConfig] = None, ) -> None: if min_workers < 1: raise ValueError("min_workers must be >= 1") @@ -192,6 +338,9 @@ def __init__( self._min_workers = min_workers self._max_workers = max_workers self._supervisor_interval = supervisor_interval + self._affinity_config: CoreAffinityConfig = ( + affinity_config if affinity_config is not None else CoreAffinityConfig(enabled=False) + ) self._work_queue: Queue = Queue() self._stop_event = threading.Event() @@ -202,6 +351,9 @@ def __init__( # Counter for assigning deterministic IDs to workers self._next_worker_id: int = 0 + # Resolved core assignments for critical workers (populated at start()). + self._pinned_cores: List[int] = [] + # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ @@ -228,12 +380,41 @@ def _make_worker_thread(self) -> threading.Thread: # ------------------------------------------------------------------ def start(self) -> None: - """Spawn initial workers and the supervisor thread.""" + """Spawn initial workers and the supervisor thread. + + Critical ingestion workers are pinned to dedicated CPU cores when + ``affinity_config.enabled`` is ``True`` and ``psutil`` is available. + Affinity assignment is performed from *within* each worker thread so + the OS-level thread inherits the correct core mask. + """ if self._supervisor_thread is not None: raise RuntimeError("Pool is already running") - for _ in range(self._min_workers): - t = self._make_worker_thread() + cfg = self._affinity_config + + if cfg.enabled and _PSUTIL_AVAILABLE: + try: + self._pinned_cores = _resolve_cores(cfg, self._min_workers) + except ValueError as exc: + if cfg.fallback_to_all_cores: + logger.warning( + "CoreAffinityConfig validation failed (%s); " + "starting without affinity pinning.", + exc, + ) + self._pinned_cores = [] + else: + raise + + use_affinity = bool(self._pinned_cores) + + for i in range(self._min_workers): + if use_affinity: + # Assign cores round-robin across available pinned cores. + core = self._pinned_cores[i % len(self._pinned_cores)] + t = self._make_pinned_worker_thread(core=core, index=i) + else: + t = self._make_worker_thread(name=f"Ingestion-Worker-{i}") t.start() self._threads.append(t) @@ -245,18 +426,35 @@ def start(self) -> None: self._stop_event, self._state, self._lock, - self._make_worker_thread, + self._make_worker_thread, # scaled workers are unpinned ), daemon=True, name="ThreadingPool-Supervisor", ) self._supervisor_thread.start() - logger.info( - "ThreadingPool: started with %d workers (min=%d, max=%d)", - self._min_workers, - self._min_workers, - self._max_workers, - ) + + if use_affinity: + logger.info( + "ThreadingPool: started with %d pinned workers " + "(cores=%s, min=%d, max=%d)", + self._min_workers, + self._pinned_cores, + self._min_workers, + self._max_workers, + ) + else: + reason = ( + "psutil unavailable" if not _PSUTIL_AVAILABLE + else "affinity disabled" + ) + logger.info( + "ThreadingPool: started with %d workers " + "(affinity: %s, min=%d, max=%d)", + self._min_workers, + reason, + self._min_workers, + self._max_workers, + ) def submit(self, task: Callable) -> None: """Enqueue *task* for execution by a worker thread. @@ -295,6 +493,7 @@ def snapshot(self) -> PoolSnapshot: queue_depth=self._work_queue.qsize(), tasks_completed=self._state.tasks_completed, tasks_failed=self._state.tasks_failed, + pinned_cores=tuple(self._pinned_cores), ) # ------------------------------------------------------------------ @@ -352,21 +551,44 @@ def _worker_with_sentinel( work_queue.task_done() +# --------------------------------------------------------------------------- +# Module-level affinity configuration +# --------------------------------------------------------------------------- + +#: Default affinity config for the ingestion pipeline. +#: +#: ``cores`` is left empty so the pool auto-selects the first ``MIN_WORKERS`` +#: logical cores at startup. Override before calling ``threading_pool.start()`` +#: to target specific cores, e.g.:: +#: +#: from src.utils.threading_pool import INGESTION_AFFINITY +#: INGESTION_AFFINITY.cores = [0, 1, 2, 3] +INGESTION_AFFINITY: CoreAffinityConfig = CoreAffinityConfig( + enabled=True, + cores=[], # auto-select at startup + fallback_to_all_cores=True, +) + # --------------------------------------------------------------------------- # Module-level singleton # --------------------------------------------------------------------------- #: Shared pool instance; call ``threading_pool.start()`` to activate. +#: Core affinity is enabled by default via ``INGESTION_AFFINITY``. threading_pool = DynamicThreadingPool( min_workers=MIN_WORKERS, max_workers=MAX_WORKERS, + affinity_config=INGESTION_AFFINITY, ) __all__ = [ "MIN_WORKERS", "MAX_WORKERS", + "CoreAffinityConfig", + "INGESTION_AFFINITY", "PoolSnapshot", "DynamicThreadingPool", + "pin_thread_to_cores", "threading_pool", ]