shared/eve_worker_pool.py

"""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)"
                        )