"""Multi-process Eve SDK worker pool. Spawns N child processes, each hosting an independent ``EveWrapper`` instance with its own ``ctypes.CDLL`` handle. The main Gradio process communicates with workers via ``multiprocessing.Pipe`` (one bidirectional pipe per worker). Frame data is serialised as raw bytes (``ndarray.tobytes()``) rather than pickled ``ndarray`` objects for speed (~1-2 ms for a 1280x720x3 frame). Usage:: pool = EveWorkerPool(max_workers=4, ram_headroom_gb=2.0) worker = pool.acquire(session_hash="abc123") try: result = worker.send_inference(frame, features) finally: pool.release(worker) """ import atexit import logging import multiprocessing as mp import os import threading import time from collections.abc import Callable from dataclasses import dataclass from multiprocessing.connection import Connection import numpy as np from eve_messages import ( CalibrateNewUserCmd, CalibrateOkResponse, CalibrationResultMsg, EnableFaceIdCmd, ErrorResponse, FeatureFlags, GalleryRestoredResponse, GetProfileStatsCmd, GetTimingStatsCmd, HeartbeatResponse, InferenceCmd, InferenceResponse, OkResponse, ProcessVideoCmd, ProfileStatsResponse, ProgressResponse, ReadyResponse, RemoveAllUsersCmd, RemoveUsersOkResponse, RestoreGalleryCmd, RestoreGalleryOkResponse, SerializedFrame, ShutdownCmd, StartProfilingCmd, StopProfilingCmd, TimingStatsResponse, VideoProcessingDoneResponse, WorkerCmd, WorkerResponse, ) from log_utils import setup_logger logger = setup_logger("EveWorkerPool") # Cached per-process so the probe runs once per worker, not per video. _h264_encoder_cache: tuple[str, dict[str, str]] | None = None def _get_h264_encoder() -> tuple[str, dict[str, str]]: """Return (codec_name, codec_options) for H.264 encoding. Tries NVENC (GPU) first, falls back to libx264 (CPU). """ global _h264_encoder_cache if _h264_encoder_cache is not None: return _h264_encoder_cache import logging import av log = logging.getLogger("EveWorkerPool") try: ctx = av.CodecContext.create("h264_nvenc", "w") ctx.width = 64 ctx.height = 64 ctx.pix_fmt = "yuv420p" ctx.open() ctx.close() log.info("Using GPU encoder: h264_nvenc") _h264_encoder_cache = ( "h264_nvenc", {"profile": "baseline", "preset": "p1", "delay": "0"}, ) except Exception: log.info("GPU encoder unavailable, falling back to libx264") _h264_encoder_cache = ( "libx264", {"profile": "baseline", "level": "3.1", "preset": "ultrafast", "threads": "1"}, ) return _h264_encoder_cache def log_worker_activity( log: logging.Logger, action: str, feature: str, pool: "EveWorkerPool", worker_id: int | None = None, note: str = "", ) -> None: """Single-line worker activity log for tracking acquire/release/queue events.""" busy = pool.worker_count - pool.idle_count w = f" w{worker_id}" if worker_id is not None else "" n = f" ({note})" if note else "" log.info( f"[worker] {action}{w} for {feature}{n} " f"[busy={busy} idle={pool.idle_count} wait={pool.waiting_count} " f"sessions={pool.session_count}]" ) def compute_fps_sleep( elapsed: float, max_fps: float | None, min_fps: float | None, load: float, ) -> float: """Compute how long to sleep after one frame to stay within the FPS cap. Linearly interpolates the target cycle time between ``1/max_fps`` (at *load* 0) and ``1/min_fps`` (at *load* 1), then subtracts the time already spent. Used by both the live-inference bridge and the video-processing loop so that the throttle logic lives in one place. """ if max_fps is None: return 0.0 effective_min = min_fps if min_fps is not None else max_fps max_period = 1.0 / max_fps min_period = 1.0 / effective_min target_period = max_period + load * (min_period - max_period) return max(0.0, target_period - elapsed) # Use 'spawn' on all platforms so each child gets a fresh interpreter # (required on Windows; avoids CDLL sharing on Linux/fork). _mp_ctx = mp.get_context("spawn") DO_PROBE_WORKER = ( False # Set to False to skip the RAM probe and just use an estimate of the RAM used ) # --------------------------------------------------------------------------- # Worker process entry point (runs in the child) # --------------------------------------------------------------------------- def _eve_worker_main( conn: Connection, eve_bin_path: str, eve_lib_path: str, worker_id: int, max_jobs_per_worker: int | None = None, shared_load: "mp.Value | None" = None, ) -> None: """Top-level function executed inside each worker ``Process``. Communicates with the main process exclusively through *conn*. *shared_load* is a cross-process float (0.0–1.0) updated by the pool on acquire/release, used by the video processing loop to throttle. """ import ctypes import gc import platform import signal import sys from pathlib import Path signal.signal(signal.SIGINT, signal.SIG_IGN) # CPU affinity (for stress testing) affinity_env = os.environ.get(f"WORKER{worker_id}_AFFINITY") if affinity_env: cpu_id = int(affinity_env.removeprefix("CPU")) import psutil psutil.Process().cpu_affinity([cpu_id]) logger.info(f"Worker {worker_id} pinned to CPU {cpu_id}") # Ensure shared package is importable inside the child shared_dir = str(Path(__file__).resolve().parent) if shared_dir not in sys.path: sys.path.insert(0, shared_dir) from eve_wrapper import EveWrapper # noqa: E402 (deferred import) from frame_utils import load_media_frames_raw # noqa: E402 eve: EveWrapper | None = None try: eve = EveWrapper(eve_bin_path, eve_lib_path) conn.send(ReadyResponse(pid=os.getpid())) except Exception as exc: conn.send(ErrorResponse(error=str(exc))) return # Optional cProfile (enabled via ENABLE_PROFILER env var) profiler: "cProfile.Profile | None" = None profiling_enabled = os.environ.get("ENABLE_PROFILER", "").strip() not in ("", "0", "false") if profiling_enabled: import cProfile # Use process_time so I/O wait (pipe polling) doesn't dominate the profile profiler = cProfile.Profile(timer=time.process_time) job_count = 0 while True: try: if not conn.poll(timeout=2.0): # No command received — send heartbeat try: conn.send(HeartbeatResponse()) except (BrokenPipeError, OSError): break continue cmd: WorkerCmd = conn.recv() except (EOFError, OSError): break if isinstance(cmd, ShutdownCmd): conn.send(OkResponse()) break try: if isinstance(cmd, InferenceCmd): frame = np.frombuffer(cmd.frame_bytes, dtype=cmd.dtype).reshape(cmd.shape) features = cmd.features del cmd # free recv'd bytes early; numpy holds its own ref eve.enable_face_and_person_detection( faceEnabled=features.face_detection, personEnabled=features.person_detection, ) eve.enable_face_id(enabled=features.face_id) eve.enable_hand_gesture(enabled=features.hand_gesture) eve.enable_object_detection(config=features.mod_model_config) eve.enable_mirror(True) result = eve.inference(frame) del frame conn.send( InferenceResponse( frame_bytes=result.tobytes(), shape=result.shape, dtype=str(result.dtype), ) ) del result elif isinstance(cmd, CalibrateNewUserCmd): frames = _deserialize_frames(cmd.frames_data) result = eve.calibrate_new_user(frames) conn.send( CalibrateOkResponse( result=CalibrationResultMsg(result.success, result.user_id, result.message), ) ) elif isinstance(cmd, RemoveAllUsersCmd): ok = eve.remove_all_users() conn.send(RemoveUsersOkResponse(result=ok)) elif isinstance(cmd, RestoreGalleryCmd): frames_per_user = [_deserialize_frames(fd) for fd in cmd.frames_per_user_data] eve.remove_all_users() results = eve.restore_gallery(frames_per_user) conn.send( RestoreGalleryOkResponse( results=[ CalibrationResultMsg(r.success, r.user_id, r.message) for r in results ], ) ) elif isinstance(cmd, EnableFaceIdCmd): eve.enable_face_id(enabled=cmd.enabled) conn.send(OkResponse()) elif isinstance(cmd, ProcessVideoCmd): import av import cv2 from fractions import Fraction features = cmd.features # Gallery restore (if requested) — must happen before # applying user feature flags because restore_gallery() # internally enables face/person detection for calibration. gallery_results: list[CalibrationResultMsg] = [] if cmd.remove_all_users: eve.remove_all_users() if cmd.gallery_paths: frames_per_user = [load_media_frames_raw(p) for p in cmd.gallery_paths] raw_results = eve.restore_gallery(frames_per_user) del frames_per_user gallery_results = [ CalibrationResultMsg(r.success, r.user_id, r.message) for r in raw_results ] del raw_results conn.send(GalleryRestoredResponse(results=gallery_results)) # Reset needs to be after the restore_gallery() function since we activate # features of EVE inside, which can mess up the ideal user. eve.reset_pipeline() # Apply user's feature flags after gallery restore so they # aren't overridden by restore_gallery's internal SDK calls. eve.enable_face_and_person_detection( faceEnabled=features.face_detection, personEnabled=features.person_detection, ) eve.enable_face_id(enabled=features.face_id) eve.enable_hand_gesture(enabled=features.hand_gesture) eve.enable_object_detection(config=features.mod_model_config) eve.enable_mirror(False) # Video processing loop cap = cv2.VideoCapture(cmd.input_path) codec_name, codec_opts = _get_h264_encoder() container = av.open(cmd.output_path, mode="w", options={"movflags": "faststart"}) stream = container.add_stream(codec_name, rate=round(cmd.fps)) stream.time_base = Fraction(1, round(cmd.fps)) stream.width = cmd.width stream.height = cmd.height stream.pix_fmt = "yuv420p" stream.codec_context.options = codec_opts frame_count = 0 v_max_fps = cmd.max_fps v_min_fps = cmd.min_fps if cmd.min_fps is not None else v_max_fps try: while True: ret, frame = cap.read() if not ret: break t0 = time.monotonic() result = eve.inference(frame) vf = av.VideoFrame.from_ndarray(result, format="bgr24") vf.pts = frame_count for pkt in stream.encode(vf): container.mux(pkt) del pkt del result, vf, frame frame_count += 1 if frame_count % 10 == 0: conn.send( ProgressResponse( current=frame_count, total=cmd.total_frames, ) ) load = shared_load.value if shared_load is not None else 0.0 # FPS throttle (only if there's some load already) if load > 0.0: sleep = compute_fps_sleep( time.monotonic() - t0, v_max_fps, v_min_fps, load ) if sleep > 0: time.sleep(sleep) # Flush encoder for pkt in stream.encode(): container.mux(pkt) del pkt finally: cap.release() container.close() del ( cap, container, stream, ) job_count += 1 conn.send( VideoProcessingDoneResponse( frames_processed=frame_count, gallery_results=gallery_results, recycle=max_jobs_per_worker is not None and job_count >= max_jobs_per_worker, ) ) del cmd gc.collect() if platform.system() == "Linux": try: ctypes.CDLL("libc.so.6").malloc_trim(0) except Exception: pass if max_jobs_per_worker is not None and job_count >= max_jobs_per_worker: logger.info(f"Worker {worker_id} recycling after {job_count} jobs") break elif isinstance(cmd, StartProfilingCmd): if profiler is not None: profiler.enable() conn.send(OkResponse()) elif isinstance(cmd, StopProfilingCmd): if profiler is not None: profiler.disable() conn.send(OkResponse()) elif isinstance(cmd, GetProfileStatsCmd): if profiler is not None: import io import marshal import pstats profiler.disable() stream = io.StringIO() ps = pstats.Stats(profiler, stream=stream) conn.send(ProfileStatsResponse(stats_data=marshal.dumps(ps.stats))) else: conn.send(ProfileStatsResponse(stats_data=b"")) elif isinstance(cmd, GetTimingStatsCmd): conn.send(TimingStatsResponse(stats=eve.get_timing_stats(reset=cmd.reset))) else: conn.send(ErrorResponse(error=f"Unknown command: {cmd}")) except Exception as exc: logger.error(f"Worker {worker_id} error handling '{type(cmd).__name__}': {exc}") try: conn.send(ErrorResponse(error=str(exc))) except (BrokenPipeError, OSError): break # Clean shutdown if eve is not None: eve.shutdown() def _serialize_frames(frames: list[np.ndarray]) -> list[SerializedFrame]: """Convert a list of ndarrays to a picklable representation.""" return [SerializedFrame(data=f.tobytes(), shape=f.shape, dtype=str(f.dtype)) for f in frames] def _deserialize_frames(data: list[SerializedFrame]) -> list[np.ndarray]: """Reconstruct ndarrays from their serialized form.""" return [np.frombuffer(d.data, dtype=d.dtype).reshape(d.shape) for d in data] # --------------------------------------------------------------------------- # EveWorker — main-process handle for one child process # --------------------------------------------------------------------------- class EveWorker: """Main-process proxy for a single worker process. Not thread-safe by itself — the ``EveWorkerPool`` ensures only one thread accesses a worker at a time. """ def __init__( self, process: mp.Process, conn: Connection, worker_id: int, ): self.process = process self._conn = conn self.worker_id = worker_id self.status: str = "idle" # idle | busy | dead self.session_hash: str | None = None self.busy_since: float | None = None self.last_heartbeat: float = time.monotonic() self.is_live_stream: bool = False self.pending_recycle: bool = False # -- command helpers --------------------------------------------------- def _recv_response(self) -> WorkerResponse: """Receive the next non-heartbeat response from the worker.""" while True: try: resp: WorkerResponse = self._conn.recv() except (EOFError, OSError): raise RuntimeError(f"Worker {self.worker_id} pipe closed") except Exception as exc: self.pending_recycle = True raise RuntimeError( f"Worker {self.worker_id} pipe corrupt: {exc}" ) from exc if isinstance(resp, HeartbeatResponse): self.last_heartbeat = time.monotonic() continue return resp def _expect_ok(self, context: str) -> WorkerResponse: """Receive a response and raise on error.""" resp = self._recv_response() if isinstance(resp, ErrorResponse): raise RuntimeError(f"Worker {self.worker_id} {context} error: {resp.error}") return resp def send_inference(self, frame: np.ndarray, features: FeatureFlags) -> np.ndarray: """Send a frame for inference and return the annotated result.""" self._conn.send( InferenceCmd( frame_bytes=frame.tobytes(), shape=frame.shape, dtype=str(frame.dtype), features=features, ) ) resp = self._expect_ok("inference") if not isinstance(resp, InferenceResponse): raise RuntimeError(f"Worker {self.worker_id} inference: unexpected {resp}") return np.frombuffer(resp.frame_bytes, dtype=resp.dtype).reshape(resp.shape) def send_calibrate_new_user(self, frames: list[np.ndarray]) -> CalibrationResultMsg: self._conn.send(CalibrateNewUserCmd(frames_data=_serialize_frames(frames))) resp = self._expect_ok("calibrate") if not isinstance(resp, CalibrateOkResponse): raise RuntimeError(f"Worker {self.worker_id} calibrate: unexpected {resp}") return resp.result def send_remove_all_users(self) -> bool: self._conn.send(RemoveAllUsersCmd()) resp = self._expect_ok("remove") if not isinstance(resp, RemoveUsersOkResponse): raise RuntimeError(f"Worker {self.worker_id} remove: unexpected {resp}") return resp.result def send_restore_gallery( self, frames_per_user: list[list[np.ndarray]] ) -> list[CalibrationResultMsg]: self._conn.send( RestoreGalleryCmd( frames_per_user_data=[_serialize_frames(fu) for fu in frames_per_user], ) ) resp = self._expect_ok("restore") if not isinstance(resp, RestoreGalleryOkResponse): raise RuntimeError(f"Worker {self.worker_id} restore: unexpected {resp}") return resp.results def send_enable_face_id(self, enabled: bool) -> None: self._conn.send(EnableFaceIdCmd(enabled=enabled)) self._expect_ok("enable_face_id") def send_process_video( self, input_path: str, output_path: str, features: FeatureFlags, gallery_paths: list[str], remove_all_users: bool, fps: float, width: int, height: int, total_frames: int, progress: Callable[..., None] | None = None, max_fps: float | None = None, min_fps: float | None = None, ) -> tuple[int, list[CalibrationResultMsg]]: """Run the full video processing loop inside the worker process. The worker reads the input video, runs Eve inference on each frame, encodes the result with PyAV, and writes the output video. Only small progress messages travel over the pipe — no frame data. Args: input_path: Path to the input video file. output_path: Path where the output video will be written. features: Feature toggles for the Eve SDK. gallery_paths: Media paths for Face ID gallery restore (may be empty). remove_all_users: Whether to clear the Face ID gallery first. fps: Output video frame rate. width: Output video width in pixels. height: Output video height in pixels. total_frames: Total frame count (for progress reporting). progress: Optional ``gr.Progress``-compatible callback. max_fps: FPS cap when idle (``None`` = unlimited). min_fps: FPS cap under full load (defaults to ``max_fps``). Returns: Tuple of (frames_processed, gallery_results). Raises: RuntimeError: If the worker reports an error. """ self._conn.send( ProcessVideoCmd( input_path=input_path, output_path=output_path, features=features, gallery_paths=gallery_paths, remove_all_users=remove_all_users, fps=fps, width=width, height=height, total_frames=total_frames, max_fps=max_fps, min_fps=min_fps, ) ) gallery_results: list[CalibrationResultMsg] = [] while True: resp = self._recv_response() if isinstance(resp, GalleryRestoredResponse): gallery_results = resp.results elif isinstance(resp, ProgressResponse): if progress is not None: progress( (resp.current, resp.total), desc=f"Processing frame {resp.current}/{resp.total}", ) elif isinstance(resp, VideoProcessingDoneResponse): if resp.recycle: self.pending_recycle = True return resp.frames_processed, gallery_results elif isinstance(resp, ErrorResponse): raise RuntimeError(f"Worker {self.worker_id} process_video error: {resp.error}") else: raise RuntimeError(f"Worker {self.worker_id} unexpected response: {resp}") def send_start_profiling(self) -> None: self._conn.send(StartProfilingCmd()) self._expect_ok("start_profiling") def send_stop_profiling(self) -> None: self._conn.send(StopProfilingCmd()) self._expect_ok("stop_profiling") def send_get_profile_stats(self) -> bytes: """Request profiling stats from the worker. Returns marshalled pstats data.""" self._conn.send(GetProfileStatsCmd()) resp = self._recv_response() if isinstance(resp, ErrorResponse): raise RuntimeError(f"Worker {self.worker_id} get_profile_stats error: {resp.error}") if not isinstance(resp, ProfileStatsResponse): raise RuntimeError(f"Worker {self.worker_id} get_profile_stats: unexpected {resp}") return resp.stats_data def send_get_timing_stats(self, reset: bool = True) -> dict[str, tuple[int, float]]: """Request per-SDK-call timing stats from the worker.""" self._conn.send(GetTimingStatsCmd(reset=reset)) resp = self._recv_response() if isinstance(resp, ErrorResponse): raise RuntimeError(f"Worker {self.worker_id} get_timing_stats error: {resp.error}") if not isinstance(resp, TimingStatsResponse): raise RuntimeError(f"Worker {self.worker_id} get_timing_stats: unexpected {resp}") return resp.stats def send_shutdown(self) -> None: try: self._conn.send(ShutdownCmd()) self._conn.recv() except (EOFError, OSError, BrokenPipeError): pass def drain_heartbeats(self) -> None: """Consume any pending heartbeat messages on the pipe.""" while self._conn.poll(timeout=0): try: msg = self._conn.recv() if isinstance(msg, HeartbeatResponse): self.last_heartbeat = time.monotonic() except (EOFError, OSError): break except Exception: # Pipe has non-pickle data (e.g. native library output) — the # framing is now unreliable so schedule this worker for recycling. logger.warning( "Worker %d: corrupt data on pipe, scheduling recycle", self.worker_id, ) self.pending_recycle = True break # --------------------------------------------------------------------------- # EveWorkerPool — manages all workers # --------------------------------------------------------------------------- @dataclass class _PoolConfig: max_workers: int = 8 max_ram_gb: float = 32.0 ram_headroom_gb: float = 2.0 rss_safety_factor: float = 2.5 stuck_timeout_s: float = 300.0 # 5 min for video processing health_check_interval_s: float = 5.0 max_jobs_per_worker: int | None = 50 class EveWorkerPool: """Pool of Eve SDK worker processes. Measures available RAM at startup, spawns as many workers as fit (with safety margins), and provides acquire / release semantics. Args: max_workers: Upper bound on the number of workers. max_ram_gb: Cap on available RAM (GB) considered for worker sizing. Useful on shared machines where reported available RAM is much larger than what the demo should consume. ram_headroom_gb: Free RAM (GB) to reserve for main process / OS. rss_safety_factor: Multiplier applied to the measured init RSS to estimate peak runtime memory per worker. eve_bin_path: Forwarded to ``EveWrapper``. eve_lib_path: Forwarded to ``EveWrapper``. """ def __init__( self, max_workers: int = 8, max_ram_gb: float = 32.0, ram_headroom_gb: float = 2.0, rss_safety_factor: float = 2.5, eve_bin_path: str = "", eve_lib_path: str = "", max_jobs_per_worker: int | None = 50, ): self._cfg = _PoolConfig( max_workers=max_workers, max_ram_gb=max_ram_gb, ram_headroom_gb=ram_headroom_gb, rss_safety_factor=rss_safety_factor, max_jobs_per_worker=max_jobs_per_worker, ) self._eve_bin_path = eve_bin_path self._eve_lib_path = eve_lib_path self._lock = threading.Condition(threading.Lock()) self._workers: list[EveWorker] = [] self._shutting_down = False self._next_id = 0 self._waiting_count = 0 self.session_count = 0 # Shared load signal readable by worker processes (0.0–1.0). # lock=False is safe: single writer (main process), readers get # a slightly stale value at worst. self._shared_load: mp.Value = _mp_ctx.Value("d", 0.0, lock=False) # Spawn workers based on RAM capacity count = self._compute_worker_count() self._spawn_workers_parallel(count) logger.info(f"EveWorkerPool started with {len(self._workers)} workers") # Start health-check daemon self._health_thread = threading.Thread(target=self._health_monitor, daemon=True) self._health_thread.start() # Hourly memory report from memory_monitor import start_memory_reporter self._mem_thread = start_memory_reporter( get_workers=lambda: list(self._workers), lock=self._lock, shutdown_flag=lambda: self._shutting_down, max_ram_gb=self._cfg.max_ram_gb, ) atexit.register(self.shutdown_all) # -- public API -------------------------------------------------------- @property def worker_count(self) -> int: """Number of live workers (for setting ``concurrency_limit``).""" with self._lock: return sum(1 for w in self._workers if w.status != "dead") @property def idle_count(self) -> int: """Number of idle workers.""" with self._lock: return sum(1 for w in self._workers if w.status == "idle") @property def waiting_count(self) -> int: """Number of callers blocked in ``acquire()`` waiting for a worker.""" with self._lock: return self._waiting_count def try_acquire(self, session_hash: str) -> EveWorker | None: """Non-blocking acquire — returns an idle worker or ``None``. Unlike :meth:`acquire`, this never blocks or increments the waiting count. Used by :class:`LiveStreamManager` so that WebRTC frame handlers can return immediately when no worker is available. """ with self._lock: return self._try_acquire(session_hash) def _try_acquire(self, session_hash: str) -> EveWorker | None: """Try to grab an idle worker (must hold ``self._lock``). Returns the worker if successful, ``None`` otherwise. """ # Prefer session-affiliated idle worker (Face ID gallery affinity) for w in self._workers: if w.status == "idle" and not w.pending_recycle and w.session_hash == session_hash: w.status = "busy" w.busy_since = time.monotonic() w.drain_heartbeats() self._update_shared_load() return w # Any idle worker for w in self._workers: if w.status == "idle" and not w.pending_recycle: w.status = "busy" w.session_hash = session_hash w.busy_since = time.monotonic() w.drain_heartbeats() self._update_shared_load() return w return None def acquire( self, session_hash: str, timeout: float = 300.0, progress: Callable[..., None] | None = None, eta_fn: Callable[[int], float | None] | None = None, ) -> EveWorker: """Reserve a worker for *session_hash*, blocking until one is free. Prefers a worker already affiliated with the session (Face ID gallery affinity). Falls back to any idle worker. Blocks up to *timeout* seconds if all workers are busy. Args: session_hash: Gradio session hash for worker affinity. timeout: Max seconds to wait for a free worker. progress: Optional ``gr.Progress``-compatible callback invoked while waiting so the UI can show queue position. eta_fn: Optional callable accepting a 1-based queue position and returning estimated seconds until a worker frees up. The return value is shown in the progress description. Raises: RuntimeError: If no worker becomes available within *timeout*. """ deadline = time.monotonic() + timeout # Fast path — try without incrementing waiting count with self._lock: worker = self._try_acquire(session_hash) if worker is not None: return worker self._waiting_count += 1 # Slow path — block until a worker is released try: while True: # Send progress update OUTSIDE the lock to avoid blocking release() if progress is not None: with self._lock: pos = self._waiting_count eta = eta_fn(pos) if eta_fn is not None else None if eta is not None: eta_int = int(eta) eta_mins = max(0.5, round(eta / 30) * 0.5) progress( (0, eta_int), desc=f"⏳ In queue ~{eta_mins:g} minutes", ) else: progress((0, 1), desc=f"⏳ In queue") with self._lock: worker = self._try_acquire(session_hash) if worker is not None: return worker remaining = deadline - time.monotonic() if remaining <= 0: raise RuntimeError( "No idle Eve workers available — timed out after " f"{timeout:.0f}s. Try again shortly." ) self._lock.wait(timeout=min(remaining, 1.0)) finally: with self._lock: self._waiting_count -= 1 def release(self, worker: EveWorker) -> None: """Return a worker to the idle pool and wake any waiting acquirers.""" with self._lock: worker.status = "idle" worker.busy_since = None worker.is_live_stream = False self._update_shared_load() self._lock.notify_all() def get_live_workers(self) -> list[EveWorker]: """Return a snapshot of non-dead workers (thread-safe).""" with self._lock: return [w for w in self._workers if w.status != "dead"] def shutdown_all(self) -> None: """Gracefully shut down all workers.""" self._shutting_down = True with self._lock: for w in self._workers: if w.status != "dead": try: w.send_shutdown() except Exception: pass w.process.join(timeout=5) if w.process.is_alive(): w.process.kill() w.process.join(timeout=2) w.status = "dead" def _update_shared_load(self) -> None: """Refresh the shared load signal (must hold ``self._lock``).""" total = sum(1 for w in self._workers if w.status != "dead") if total <= 1: self._shared_load.value = 0.0 else: idle = sum(1 for w in self._workers if w.status == "idle") self._shared_load.value = 1.0 - idle / total # -- internals --------------------------------------------------------- def _compute_worker_count(self) -> int: """Determine how many workers fit in available RAM.""" try: import psutil except ImportError: logger.warning("psutil not installed — defaulting to 1 worker") return 1 raw_available_mb = psutil.virtual_memory().available / (1024 * 1024) cap_mb = self._cfg.max_ram_gb * 1024 available_mb = min(raw_available_mb, cap_mb) headroom_mb = self._cfg.ram_headroom_gb * 1024 # Spawn a probe worker to measure init RSS if DO_PROBE_WORKER: probe = self._spawn_worker(probe=True) try: import psutil as _ps probe_rss_mb = _ps.Process(probe.process.pid).memory_info().rss / (1024 * 1024) except Exception: logger.warning("Could not measure worker RSS — defaulting to 1 worker") return 1 else: # Estimated 500MB RSS for an idle worker without running inference (based on observed values) # It's around 200MB idle, and there's the video being loaded in memory, depends on the size of the video. probe_rss_mb = 250 estimated_peak_mb = probe_rss_mb * self._cfg.rss_safety_factor usable_mb = available_mb - headroom_mb # +1 because the probe worker already consumed memory max_by_ram = max(1, int(usable_mb / estimated_peak_mb) + 1) if DO_PROBE_WORKER: actual = min(self._cfg.max_workers - 1, max_by_ram) else: actual = min(self._cfg.max_workers, max_by_ram) logger.info( f"RAM estimation: init_rss={probe_rss_mb:.0f}MB, " f"estimated_peak={estimated_peak_mb:.0f}MB, " f"raw_available={raw_available_mb:.0f}MB, " f"available={available_mb:.0f}MB (cap={cap_mb:.0f}MB), " f"headroom={headroom_mb:.0f}MB, " f"max_by_ram={max_by_ram}, max_workers={self._cfg.max_workers}, " f"actual={actual}" ) return actual def _launch_worker(self) -> tuple[int, mp.Process, Connection]: """Start a worker process without waiting for readiness. Returns: Tuple of (worker_id, process, parent_conn). """ wid = self._next_id self._next_id += 1 parent_conn, child_conn = _mp_ctx.Pipe() proc = _mp_ctx.Process( target=_eve_worker_main, args=( child_conn, self._eve_bin_path, self._eve_lib_path, wid, self._cfg.max_jobs_per_worker, self._shared_load, ), daemon=True, ) proc.start() return wid, proc, parent_conn def _wait_for_ready(self, wid: int, proc: mp.Process, parent_conn: Connection) -> EveWorker: """Block until a launched worker sends its "ready" message. Raises: RuntimeError: If the worker doesn't become ready within 120 s. """ if not parent_conn.poll(timeout=120): proc.kill() proc.join(timeout=5) raise RuntimeError(f"Worker {wid} did not start within 120 s") msg: WorkerResponse = parent_conn.recv() if not isinstance(msg, ReadyResponse): proc.kill() proc.join(timeout=5) error = msg.error if isinstance(msg, ErrorResponse) else str(msg) raise RuntimeError(f"Worker {wid} failed to initialise: {error}") worker = EveWorker(proc, parent_conn, wid) with self._lock: self._workers.append(worker) self._lock.notify_all() logger.info(f"Worker {wid} (pid={proc.pid}) ready") return worker def _spawn_workers_parallel(self, count: int) -> None: """Launch *count* workers in parallel and wait for all to be ready.""" pending = [self._launch_worker() for _ in range(count)] for wid, proc, conn in pending: try: self._wait_for_ready(wid, proc, conn) except RuntimeError: logger.error(f"Worker {wid} failed to start — skipping") def _spawn_worker(self, probe: bool = False) -> EveWorker: """Spawn a single worker and wait for it to become ready. Used for the RAM probe and for respawning crashed workers. """ wid, proc, conn = self._launch_worker() return self._wait_for_ready(wid, proc, conn) def _respawn_worker(self, dead_worker: EveWorker) -> None: """Replace a dead worker with a fresh one.""" if self._shutting_down: return logger.warning( f"Respawning worker {dead_worker.worker_id} " f"(was pid={dead_worker.process.pid})" ) dead_worker.process.join(timeout=0) with self._lock: self._workers.remove(dead_worker) try: self._spawn_worker() except RuntimeError as exc: logger.error(f"Failed to respawn worker: {exc}") def _health_monitor(self) -> None: """Background daemon that detects dead / stuck workers.""" while not self._shutting_down: time.sleep(self._cfg.health_check_interval_s) if self._shutting_down: break with self._lock: workers_snapshot = list(self._workers) for w in workers_snapshot: if w.status == "dead": continue # Detect crashed / zombie processes if not w.process.is_alive(): if w.pending_recycle: logger.info(f"Worker {w.worker_id} (pid={w.process.pid}) recycled cleanly") else: logger.error( f"Worker {w.worker_id} (pid={w.process.pid}) died unexpectedly" ) w.status = "dead" self._respawn_worker(w) continue # Detect stuck workers (skip live streams — they're long by design) if w.status == "busy" and not w.is_live_stream and w.busy_since is not None: elapsed = time.monotonic() - w.busy_since if elapsed > self._cfg.stuck_timeout_s * 2: logger.error(f"Worker {w.worker_id} stuck for {elapsed:.0f}s — killing") w.process.kill() w.process.join(timeout=5) w.status = "dead" self._respawn_worker(w) elif elapsed > self._cfg.stuck_timeout_s: logger.warning( f"Worker {w.worker_id} busy for {elapsed:.0f}s " f"(threshold={self._cfg.stuck_timeout_s}s)" )