hao-ai-lab-FastVideo/apps/dreamverse/server/gpu_pool.py

# pyright: reportMissingTypeArgument=false, reportArgumentType=false, reportOptionalSubscript=false, reportAttributeAccessIssue=false, reportOptionalMemberAccess=false, reportUndefinedVariable=false
# ruff: noqa: UP038, SIM105, F821
# mypy: ignore-errors
import asyncio
import multiprocessing as mp
import os
import subprocess
import time
import traceback
from dataclasses import dataclass
from enum import Enum
from multiprocessing import Process, Queue

from config import (
    DEFAULT_MODEL_ID,
    MODEL_REGISTRY,
    STARTUP_WARMUP_ENABLED,
    STARTUP_WARMUP_PROMPT,
    STARTUP_WARMUP_TIMEOUT_SECONDS,
)
from av_streaming import (
    SHARED_STREAM_BUFFER_BYTES,
    USE_SHARED_STREAM_BUFFER,
    StreamChunk,
    StreamComplete,
    StreamEvent,
    StreamInit,
    generate_stream_id,
    stream_fmp4,
)
from worker_ipc import (
    CommandPayload,
    InitAck,
    JoinAck,
    LeaveAck,
    MediaChunk,
    MediaComplete,
    MediaInit,
    ReloadAck,
    ReloadModelPayload,
    ShutdownAck,
    StepComplete,
    UserStepPayload,
    WarmupComplete,
    WarmupPayload,
    WorkerError,
    WorkerEvent,
)


def _parse_requested_gpu_limit() -> int | None:
    raw_value = os.getenv("FASTVIDEO_GPU_COUNT", "").strip().lower()
    if not raw_value:
        return 1
    if raw_value == "all":
        return None
    try:
        requested = int(raw_value)
    except ValueError as exc:
        raise RuntimeError("Invalid FASTVIDEO_GPU_COUNT. Use a positive integer or 'all'.") from exc
    if requested <= 0:
        raise RuntimeError("Invalid FASTVIDEO_GPU_COUNT. Use a positive integer or 'all'.")
    return requested


def _limit_gpu_ids(gpu_ids: list[int]) -> list[int]:
    requested_limit = _parse_requested_gpu_limit()
    print(f"[INFO] Using GPU limit={requested_limit}")
    if requested_limit is None:
        return gpu_ids
    return gpu_ids[:requested_limit] or gpu_ids


class CommandType(Enum):
    """Commands sent from main process to GPU worker."""
    INIT = "init"
    WARMUP = "warmup"
    SHUTDOWN = "shutdown"
    USER_JOIN = "user_join"
    USER_STEP = "user_step"
    USER_LEAVE = "user_leave"
    RELOAD_MODEL = "reload_model"


@dataclass
class Command:
    """Command sent to GPU worker subprocess.

    Commands that carry data (USER_STEP, WARMUP, RELOAD_MODEL)
    populate ``payload`` with a typed payload from ``worker_ipc``.
    Commands that don't (INIT, SHUTDOWN, USER_JOIN, USER_LEAVE)
    leave ``payload`` as ``None``.
    """
    type: CommandType
    payload: CommandPayload | None = None
    user_id: str | None = None


def _stream_event_to_worker_event(
    event: StreamEvent,
    user_id: str,
    segment_idx: int,
) -> WorkerEvent:
    """Translate an av_streaming event into a typed worker event.

    ``StreamEvent`` (ffmpeg output layer) carries no routing info;
    this adds ``user_id`` + ``segment_idx`` so the pool's
    ``_response_reader`` can dispatch to the right per-user queue.
    """
    match event:
        case StreamInit(stream_id=sid, mime=m, uses_shared_buffer=u):
            return MediaInit(
                user_id=user_id,
                segment_idx=segment_idx,
                stream_id=sid,
                mime=m,
                uses_shared_buffer=u,
            )
        case StreamChunk(
            stream_id=sid,
            chunk=c,
            chunk_offset=co,
            chunk_length=cl,
            uses_shared_buffer=u,
        ):
            return MediaChunk(
                user_id=user_id,
                segment_idx=segment_idx,
                stream_id=sid,
                chunk=c,
                chunk_offset=co,
                chunk_length=cl,
                uses_shared_buffer=u,
            )
        case StreamComplete(stream_id=sid, chunks=n):
            return MediaComplete(
                user_id=user_id,
                segment_idx=segment_idx,
                stream_id=sid,
                chunks=n,
            )
        case _:
            raise ValueError(f"unknown stream event: {type(event).__name__}")


def gpu_worker_process(
    gpu_id: int,
    cuda_device: str,
    command_queue: Queue,
    response_queue: Queue,
    shared_stream_buffer=None,
    shared_stream_buffer_bytes: int = 0,
):
    """Worker process that runs on a single GPU.

    CUDA_VISIBLE_DEVICES must be set BEFORE importing VideoGenerationWorker
    (which transitively touches CUDA).  Delegates model lifecycle and
    generation to VideoGenerationWorker; AV muxing to av_streaming.
    """
    os.environ["CUDA_VISIBLE_DEVICES"] = cuda_device
    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"

    from video_generation import VideoGenerationWorker

    worker = VideoGenerationWorker(gpu_id)

    def event_loop(first_cmd: Command = None):
        """Blocking event loop for LTX2; dispatches user commands."""
        print(f"[GPU {gpu_id}] Entering event loop")

        def handle_command(cmd: Command):
            if cmd.type == CommandType.USER_JOIN:
                print(f"[GPU {gpu_id}] User {cmd.user_id[:8]} joined")
                worker.clear_conditioning()
                response_queue.put(JoinAck(user_id=cmd.user_id))

            elif cmd.type == CommandType.USER_STEP:
                try:
                    assert isinstance(cmd.payload, UserStepPayload), (f"USER_STEP requires UserStepPayload, "
                                                                      f"got {type(cmd.payload).__name__}")
                    payload = cmd.payload
                    segment_idx = payload.segment_idx
                    step_result = worker.generate_step(
                        payload.prompt,
                        segment_idx,
                        image_path=payload.image_path,
                        reset_conditioning=payload.reset_conditioning,
                    )
                    head_trim_frames = step_result.head_trim_frames
                    head_trim_audio_frames = step_result.head_trim_audio_frames
                    if head_trim_frames > 0 or head_trim_audio_frames > 0:
                        print(f"[GPU {gpu_id}] Segment {segment_idx}: "
                              f"trimming video={head_trim_frames} "
                              f"audio={head_trim_audio_frames} "
                              f"overlap frames from AV output")
                    audio_shape = getattr(step_result.audio, "shape", None)
                    print(f"[GPU {gpu_id}] AV attempt segment "
                          f"{segment_idx}: "
                          f"audio_present={step_result.audio is not None}, "
                          f"audio_shape={audio_shape}, "
                          f"audio_sample_rate={step_result.audio_sample_rate}")
                    stream_id = generate_stream_id(segment_idx)

                    def _publish(event: StreamEvent) -> None:
                        response_queue.put(_stream_event_to_worker_event(event, cmd.user_id, segment_idx))

                    av_ok, av_error = stream_fmp4(
                        frames=step_result.frames,
                        audio=step_result.audio,
                        audio_sample_rate=step_result.audio_sample_rate,
                        stream_id=stream_id,
                        timings=step_result.timings,
                        head_trim_frames=head_trim_frames,
                        head_trim_audio_frames=head_trim_audio_frames,
                        shared_buffer=shared_stream_buffer,
                        shared_buffer_bytes=shared_stream_buffer_bytes,
                        publish=_publish,
                        log_prefix=f"[GPU {gpu_id}]",
                    )
                    if not av_ok:
                        raise RuntimeError(av_error or "worker av_fmp4 stream failed")
                    step_result.timings["ipc_put_start_ns"] = time.time_ns()
                    response_queue.put(
                        StepComplete(
                            user_id=cmd.user_id,
                            segment_idx=segment_idx,
                            timings=step_result.timings,
                        ))
                except Exception as e:
                    print(f"[GPU {gpu_id}] Step error: {e}")
                    traceback.print_exc()
                    response_queue.put(WorkerError(
                        user_id=cmd.user_id,
                        message=str(e),
                    ))

            elif cmd.type == CommandType.USER_LEAVE:
                print(f"[GPU {gpu_id}] User {cmd.user_id[:8]} left")
                worker.clear_conditioning()
                response_queue.put(LeaveAck(user_id=cmd.user_id))

            elif cmd.type == CommandType.SHUTDOWN:
                return False  # Signal to exit

            elif cmd.type == CommandType.WARMUP:
                try:
                    assert isinstance(cmd.payload, WarmupPayload), (f"WARMUP requires WarmupPayload, "
                                                                    f"got {type(cmd.payload).__name__}")
                    timings = worker.warmup(cmd.payload.prompt)
                    response_queue.put(WarmupComplete(
                        user_id=cmd.user_id,
                        timings=timings,
                    ))
                except Exception as e:
                    print(f"[GPU {gpu_id}] Warmup error: {e}")
                    traceback.print_exc()
                    response_queue.put(WorkerError(
                        user_id=cmd.user_id,
                        message=str(e),
                    ))

            elif cmd.type == CommandType.RELOAD_MODEL:
                try:
                    assert isinstance(cmd.payload, ReloadModelPayload), (f"RELOAD_MODEL requires ReloadModelPayload, "
                                                                         f"got {type(cmd.payload).__name__}")
                    worker.initialize(cmd.payload.model_config)
                    response_queue.put(ReloadAck(user_id=cmd.user_id))
                except Exception as e:
                    print(f"[GPU {gpu_id}] Reload error: {e}")
                    traceback.print_exc()
                    response_queue.put(WorkerError(
                        user_id=cmd.user_id,
                        message=str(e),
                    ))

            return True  # Continue loop

        if first_cmd is not None:
            if first_cmd.type == CommandType.SHUTDOWN:
                worker.shutdown()
                response_queue.put(ShutdownAck())
                return
            handle_command(first_cmd)

        import queue as queue_module
        while True:
            try:
                cmd = command_queue.get(timeout=1.0)
            except queue_module.Empty:
                continue

            if cmd.type == CommandType.SHUTDOWN:
                print(f"[GPU {gpu_id}] Event loop shutting down")
                worker.shutdown()
                response_queue.put(ShutdownAck())
                return

            if not handle_command(cmd):
                return

    print(f"[GPU {gpu_id}] Worker process starting...")

    try:
        while True:
            cmd: Command = command_queue.get()

            if cmd.type == CommandType.SHUTDOWN:
                print(f"[GPU {gpu_id}] Shutting down...")
                worker.shutdown()
                response_queue.put(ShutdownAck())
                break

            elif cmd.type == CommandType.INIT:
                try:
                    worker.initialize()
                    response_queue.put(InitAck(success=True))
                except Exception as e:
                    print(f"[GPU {gpu_id}] Init error: {e}")
                    traceback.print_exc()
                    response_queue.put(InitAck(success=False, error=str(e)))

            elif cmd.type == CommandType.WARMUP:
                try:
                    assert isinstance(cmd.payload, WarmupPayload), (f"WARMUP requires WarmupPayload, "
                                                                    f"got {type(cmd.payload).__name__}")
                    timings = worker.warmup(cmd.payload.prompt)
                    response_queue.put(WarmupComplete(
                        user_id=cmd.user_id,
                        timings=timings,
                    ))
                except Exception as e:
                    print(f"[GPU {gpu_id}] Warmup error: {e}")
                    traceback.print_exc()
                    response_queue.put(WorkerError(
                        user_id=cmd.user_id,
                        message=str(e),
                    ))

            elif cmd.type == CommandType.RELOAD_MODEL:
                try:
                    assert isinstance(cmd.payload, ReloadModelPayload), (f"RELOAD_MODEL requires ReloadModelPayload, "
                                                                         f"got {type(cmd.payload).__name__}")
                    worker.initialize(cmd.payload.model_config)
                    response_queue.put(ReloadAck(user_id=cmd.user_id))
                except Exception as e:
                    print(f"[GPU {gpu_id}] Reload error: {e}")
                    traceback.print_exc()
                    response_queue.put(WorkerError(
                        user_id=cmd.user_id,
                        message=str(e),
                    ))

            elif cmd.type in (CommandType.USER_JOIN, CommandType.USER_STEP, CommandType.USER_LEAVE):
                event_loop(first_cmd=cmd)
                break

    except Exception as e:
        print(f"[GPU {gpu_id}] Worker crashed: {e}")
        traceback.print_exc()

    print(f"[GPU {gpu_id}] Worker process exiting")


class GPUSlot:
    """Manages a single GPU worker subprocess."""

    def __init__(self, gpu_id: int, cuda_device: str):
        self.gpu_id = gpu_id
        self.cuda_device = cuda_device
        self.process: Process | None = None
        self.command_queue: Queue | None = None
        self.response_queue: Queue | None = None
        self.ready: bool = False
        self.warmup_enabled: bool = STARTUP_WARMUP_ENABLED
        self.warmup_success: bool = False
        self.warmup_error: str | None = None
        self.warmup_timings: dict[str, float] = {}
        self._lock = asyncio.Lock()

        # Client state
        self.connected_users: set[str] = set()
        self._pending_futures: dict[str, asyncio.Future] = {}
        self._stream_queues: dict[str, asyncio.Queue] = {}
        self._response_reader_task: asyncio.Task | None = None
        self._active: bool = False
        self._reader_lock: asyncio.Lock | None = None
        self.current_model_id: str = DEFAULT_MODEL_ID
        self.shared_stream_buffer = None
        self.shared_stream_buffer_size = SHARED_STREAM_BUFFER_BYTES

    @property
    def client_count(self) -> int:
        return len(self.connected_users)

    @property
    def is_available(self) -> bool:
        """A GPU is available if it has no active users."""
        alive = self.ready and self.process is not None and self.process.is_alive()
        if not alive:
            return False
        return len(self.connected_users) == 0

    @property
    def is_empty(self) -> bool:
        return len(self.connected_users) == 0

    async def start(self):
        """Start the GPU worker subprocess."""
        self.ready = False
        self.warmup_success = False
        self.warmup_error = None
        self.warmup_timings = {}

        ctx = mp.get_context("spawn")
        self.command_queue = ctx.Queue()
        self.response_queue = ctx.Queue()
        if USE_SHARED_STREAM_BUFFER and self.shared_stream_buffer_size > 0:
            # Fixed shared byte buffer to avoid per-chunk IPC payload copies.
            self.shared_stream_buffer = mp.RawArray("B", self.shared_stream_buffer_size)

        self.process = ctx.Process(
            target=gpu_worker_process,
            args=(
                self.gpu_id,
                self.cuda_device,
                self.command_queue,
                self.response_queue,
                self.shared_stream_buffer,
                self.shared_stream_buffer_size,
            ),
            daemon=False,
        )
        loop = asyncio.get_event_loop()
        await loop.run_in_executor(None, self.process.start)

        # Send init command and wait for response
        init_response = await self._send_command(Command(CommandType.INIT), timeout=600.0)
        if not isinstance(init_response, InitAck) or not init_response.success:
            error_msg = (init_response.error if isinstance(init_response, InitAck) else
                         f"unexpected init response: {type(init_response).__name__}")
            raise RuntimeError(f"GPU {self.gpu_id} failed to initialize: {error_msg}")

        if self.warmup_enabled:
            try:
                warmup_response = await self._send_command(
                    Command(
                        CommandType.WARMUP,
                        payload=WarmupPayload(prompt=STARTUP_WARMUP_PROMPT),
                        user_id=f"__warmup_gpu_{self.gpu_id}__",
                    ),
                    timeout=float(STARTUP_WARMUP_TIMEOUT_SECONDS),
                )
            except Exception as exc:
                self.warmup_error = str(exc)
                raise RuntimeError(f"GPU {self.gpu_id} warmup failed: {self.warmup_error}") from exc
            match warmup_response:
                case WarmupComplete(timings=timings):
                    self.warmup_timings = {
                        key: float(value)
                        for key, value in timings.items() if isinstance(value, (int, float))
                    }
                    self.warmup_success = True
                case WorkerError(message=msg):
                    self.warmup_error = msg or "Warmup failed."
                    raise RuntimeError(f"GPU {self.gpu_id} warmup failed: {self.warmup_error}")
                case _:
                    self.warmup_error = (f"unexpected warmup response: "
                                         f"{type(warmup_response).__name__}")
                    raise RuntimeError(f"GPU {self.gpu_id} warmup failed: {self.warmup_error}")
        else:
            print(f"[GPU {self.gpu_id}] Startup warmup disabled by "
                  "FASTVIDEO_ENABLE_STARTUP_WARMUP")

        self.ready = True

    async def _send_command(self, cmd: Command, timeout: float = 300.0) -> WorkerEvent:
        """Send a command and wait for the response or worker death.

        Whichever arrives first wins.  A worker dying (exit, signal, OOM,
        segfault) flips the kernel-level sentinel fd readable, which
        asyncio notices via add_reader — typically within ~10 ms.  The
        queue timeout still bounds hangs where the worker stays alive but
        never replies.
        """
        loop = asyncio.get_running_loop()
        process = self.process
        await loop.run_in_executor(None, self.command_queue.put, cmd)

        response_fut = loop.run_in_executor(None, lambda: self.response_queue.get(timeout=timeout))

        death_fut: asyncio.Future | None = None
        sentinel_fd = process.sentinel if process is not None else None

        if sentinel_fd is not None:
            death_fut = loop.create_future()

            def _on_death() -> None:
                try:
                    loop.remove_reader(sentinel_fd)
                except (ValueError, OSError):
                    pass
                if death_fut is not None and not death_fut.done():
                    death_fut.set_result(None)

            loop.add_reader(sentinel_fd, _on_death)

        waiters = [response_fut] + ([death_fut] if death_fut is not None else [])
        try:
            done, _ = await asyncio.wait(waiters, return_when=asyncio.FIRST_COMPLETED)
        finally:
            if (sentinel_fd is not None and death_fut is not None and not death_fut.done()):
                try:
                    loop.remove_reader(sentinel_fd)
                except (ValueError, OSError):
                    pass
                death_fut.cancel()

        if (death_fut is not None and death_fut in done and response_fut not in done):
            try:
                return self.response_queue.get_nowait()
            except Exception:
                pass
            self.ready = False
            pid = process.pid if process is not None else "?"
            # Reap the process so .exitcode is populated.  Sentinel
            # readability means the kernel has already exited the
            # process; join is non-blocking in practice.
            if process is not None:
                try:
                    process.join(timeout=1)
                except Exception:
                    pass
            exitcode = process.exitcode if process is not None else None
            raise RuntimeError(f"GPU {self.gpu_id} worker died during command "
                               f"(pid={pid}, exitcode={exitcode})")

        return response_fut.result()

    async def _send_command_tagged(self, cmd: Command, timeout: float = 300.0) -> WorkerEvent:
        """Send a tagged command and wait for the matching response.

        Uses the response reader background task to route responses.
        """
        loop = asyncio.get_event_loop()
        future = loop.create_future()

        # Register this request's pending future
        self._pending_futures[cmd.user_id] = future

        # Send command
        await loop.run_in_executor(None, self.command_queue.put, cmd)

        # Ensure response reader is running
        await self._ensure_response_reader()

        # Wait for the response
        try:
            response = await asyncio.wait_for(future, timeout=timeout)
        except asyncio.TimeoutError:
            if (cmd.user_id in self._pending_futures and self._pending_futures[cmd.user_id] is future):
                self._pending_futures.pop(cmd.user_id, None)
            raise

        return response

    async def _ensure_response_reader(self):
        """Start the response reader background task if not running."""
        if self._reader_lock is None:
            self._reader_lock = asyncio.Lock()

        async with self._reader_lock:
            if (self._response_reader_task is None or self._response_reader_task.done()):
                self._response_reader_task = asyncio.create_task(self._response_reader())

    async def _response_reader(self):
        """Background task that reads response queue and routes to per-user futures."""
        loop = asyncio.get_event_loop()

        while self._active:
            try:

                def get_response_nonblocking():
                    try:
                        return self.response_queue.get(timeout=0.5)
                    except Exception:
                        return None

                event = await loop.run_in_executor(None, get_response_nonblocking)

                if event is None:
                    continue

                # AV streaming events → route to the user's stream queue.
                if isinstance(event, (MediaInit, MediaChunk, MediaComplete)):
                    stream_queue = self._stream_queues.get(event.user_id)
                    if stream_queue is not None:
                        await stream_queue.put(event)
                    else:
                        print(f"[GPU {self.gpu_id}] Unmatched stream event for user "
                              f"{event.user_id[:8]}")
                    continue

                # System-level acks shouldn't reach the tagged router; they
                # belong to the untagged `_send_command` path.
                if isinstance(event, (InitAck, ShutdownAck)):
                    print(f"[GPU {self.gpu_id}] System event leaked into "
                          f"tagged reader: {type(event).__name__}")
                    continue

                # Late-mutation of timings for observability.  Only events
                # that actually carry timings get this annotation.
                if isinstance(event, (StepComplete, WarmupComplete)):
                    event.timings["ipc_get_done_ns"] = time.time_ns()

                user_id = event.user_id
                if user_id and user_id in self._pending_futures:
                    future = self._pending_futures.pop(user_id)
                    if not future.done():
                        future.set_result(event)
                else:
                    print(f"[GPU {self.gpu_id}] Unmatched response for user "
                          f"{user_id[:8] if user_id else 'None'}")
            except asyncio.CancelledError:
                break
            except Exception as e:
                print(f"[GPU {self.gpu_id}] Response reader error: {e}")
                await asyncio.sleep(0.01)

    def register_stream_queue(self, user_id: str) -> asyncio.Queue:
        """Register stream-event queue for a specific user."""
        queue = asyncio.Queue()
        self._stream_queues[user_id] = queue
        return queue

    def unregister_stream_queue(self, user_id: str) -> None:
        """Remove stream-event queue for a specific user."""
        self._stream_queues.pop(user_id, None)

    async def join_user(self, user_id: str, model_id: str = None) -> JoinAck:
        """Add a user to this GPU."""
        if model_id is None:
            model_id = DEFAULT_MODEL_ID

        # Reload model if a different one is requested
        if model_id != self.current_model_id and model_id in MODEL_REGISTRY:
            print(f"[GPU {self.gpu_id}] Model switch: "
                  f"{self.current_model_id} -> {model_id}")

            for uid, future in list(self._pending_futures.items()):
                if not future.done():
                    future.set_exception(RuntimeError("Model changed, session reset"))
            self._pending_futures.clear()
            self._stream_queues.clear()
            self.connected_users.clear()

            model_config = MODEL_REGISTRY[model_id]
            reload_response = await self._send_command(Command(CommandType.RELOAD_MODEL,
                                                               payload=ReloadModelPayload(model_config=model_config),
                                                               user_id="__reload__"),
                                                       timeout=600.0)
            match reload_response:
                case ReloadAck():
                    pass
                case WorkerError(message=msg):
                    raise RuntimeError(f"Model reload failed: {msg}")
                case _:
                    raise RuntimeError(f"Unexpected reload response: "
                                       f"{type(reload_response).__name__}")

            self.current_model_id = model_id
            print(f"[GPU {self.gpu_id}] Model reloaded: {model_id}")

        self._active = True
        self.connected_users.add(user_id)
        try:
            response = await self._send_command_tagged(Command(CommandType.USER_JOIN, user_id=user_id), timeout=600.0)
            match response:
                case JoinAck() as ack:
                    return ack
                case WorkerError(message=msg):
                    self.connected_users.discard(user_id)
                    raise RuntimeError(f"User join failed for {user_id[:8]}: {msg}")
                case _:
                    self.connected_users.discard(user_id)
                    raise RuntimeError(f"Unexpected join response for {user_id[:8]}: "
                                       f"{type(response).__name__}")
        except Exception:
            self.connected_users.discard(user_id)
            raise

    async def user_step(
        self,
        user_id: str,
        prompt: str,
        segment_idx: int = 1,
        image_path: str | None = None,
        reset_conditioning: bool = False,
    ) -> dict[str, float]:
        """Execute a generation step for a specific user.

        Returns the timings dict.  Frames/audio are no longer part of
        this return — they stream asynchronously via the AV media
        events (MediaInit/MediaChunk/MediaComplete) which the caller
        consumes through ``register_stream_queue``.
        """
        payload = UserStepPayload(
            prompt=prompt,
            segment_idx=segment_idx,
            image_path=image_path,
            reset_conditioning=bool(reset_conditioning),
        )
        response = await self._send_command_tagged(Command(CommandType.USER_STEP, payload=payload, user_id=user_id),
                                                   timeout=1800.0)
        match response:
            case StepComplete(timings=timings):
                return timings
            case WorkerError(message=msg):
                raise RuntimeError(f"User step failed for {user_id[:8]}: {msg}")
            case _:
                raise RuntimeError(f"Unexpected step response for {user_id[:8]}: "
                                   f"{type(response).__name__}")

    async def leave_user(self, user_id: str) -> None:
        """Remove a user from this GPU."""
        try:
            await self._send_command_tagged(Command(CommandType.USER_LEAVE, user_id=user_id), timeout=30.0)
        except Exception as e:
            print(f"[GPU {self.gpu_id}] Leave user error: {e}")
        finally:
            self.connected_users.discard(user_id)
            self._pending_futures.pop(user_id, None)
            self._stream_queues.pop(user_id, None)

    async def shutdown(self):
        """Shutdown the worker subprocess."""
        self._active = False
        if self._response_reader_task and not self._response_reader_task.done():
            self._response_reader_task.cancel()
            try:
                await self._response_reader_task
            except asyncio.CancelledError:
                pass

        if self.process is not None:
            if self.process.is_alive():
                try:
                    await self._send_command(Command(CommandType.SHUTDOWN), timeout=30.0)
                except Exception:
                    pass
                self.process.terminate()
                self.process.join(timeout=5)
                if self.process.is_alive():
                    self.process.kill()
                    self.process.join(timeout=5)
            else:
                # Process already died (sentinel path).  Reap it so the
                # OS releases the PID.
                try:
                    self.process.join(timeout=1)
                except Exception:
                    pass

        for q in (self.command_queue, self.response_queue):
            if q is not None:
                try:
                    q.close()
                except Exception:
                    pass
                try:
                    q.join_thread()
                except Exception:
                    pass


class GPUPool:
    """Manages multiple GPU worker subprocesses."""

    def __init__(self, gpu_ids: list[int]):
        self.gpu_ids = gpu_ids
        self.slots: dict[int, GPUSlot] = {gpu_id: GPUSlot(gpu_id, str(gpu_id)) for gpu_id in gpu_ids}
        self.waiting_list: list[tuple[str, asyncio.Event, WebSocket]] = []
        self.client_gpu_map: dict[str, int] = {}
        self._pool_lock = asyncio.Lock()

    async def initialize(self):
        """Initialize all GPU workers and wait for them to be ready.

        Any per-GPU failure aborts startup so uvicorn refuses to serve
        traffic with no functional GPUs.
        """
        print(f"Initializing GPU pool with {len(self.gpu_ids)} GPUs: {self.gpu_ids}")
        await asyncio.gather(*(self._init_gpu(gpu_id) for gpu_id in self.gpu_ids))

    async def _init_gpu(self, gpu_id: int):
        """Initialize a single GPU and assign any waiting clients."""
        try:
            await self.slots[gpu_id].start()
        except Exception as e:
            print(f"GPU {gpu_id} failed to initialize: {e}")
            await self.slots[gpu_id].shutdown()
            raise

        print(f"GPU pool: {gpu_id} ready "
              f"({sum(1 for s in self.slots.values() if s.ready)}/{len(self.gpu_ids)})")

        # Check if anyone is waiting for a GPU
        async with self._pool_lock:
            slot = self.slots[gpu_id]
            if slot.is_available and self.waiting_list:
                waiting_client_id, ready_event, _ = self.waiting_list.pop(0)
                self.client_gpu_map[waiting_client_id] = gpu_id
                print(f"Client {waiting_client_id[:8]} assigned GPU {gpu_id} from queue")
                ready_event.set()

                await self._send_queue_updates()

    async def acquire(self, client_id: str, websocket=None) -> tuple[int, GPUSlot]:
        """Acquire a GPU slot for a client."""
        async with self._pool_lock:
            for gpu_id, slot in self.slots.items():
                if slot.is_available:
                    self.client_gpu_map[client_id] = gpu_id
                    print(f"Client {client_id[:8]} acquired GPU {gpu_id}")
                    return gpu_id, slot

        # No slot available, wait in queue
        print(f"Client {client_id[:8]} waiting in queue "
              f"(all {len(self.gpu_ids)} GPUs at capacity)")
        ready_event = asyncio.Event()
        async with self._pool_lock:
            self.waiting_list.append((client_id, ready_event, websocket))
            await self._send_queue_updates()

        try:
            await ready_event.wait()
        except asyncio.CancelledError:
            # Client disconnected while queued; remove stale queue entry.
            async with self._pool_lock:
                self.waiting_list = [
                    item for item in self.waiting_list if not (item[0] == client_id and item[1] is ready_event)
                ]
                self.client_gpu_map.pop(client_id, None)
                await self._send_queue_updates()
            raise

        gpu_id = self.client_gpu_map.get(client_id)
        if gpu_id is None:
            raise RuntimeError(f"Client {client_id} was signaled but has no GPU assigned")

        return gpu_id, self.slots[gpu_id]

    async def release(self, client_id: str):
        """Release a client from its GPU slot."""
        async with self._pool_lock:
            # Remove stale queue entries if the client disconnected while waiting.
            prev_wait_len = len(self.waiting_list)
            self.waiting_list = [item for item in self.waiting_list if item[0] != client_id]
            removed_from_queue = len(self.waiting_list) != prev_wait_len

            gpu_id = self.client_gpu_map.pop(client_id, None)
            if gpu_id is None:
                if removed_from_queue:
                    await self._send_queue_updates()
                return

            slot = self.slots[gpu_id]
            print(f"Client {client_id[:8]} released GPU {gpu_id}")

            try:
                await slot.leave_user(client_id)
            except Exception as e:
                print(f"[GPU {gpu_id}] Leave user failed: {e}")

            # Assign to next waiting client if GPU has capacity
            if slot.is_available and self.waiting_list:
                waiting_client_id, ready_event, _ = self.waiting_list.pop(0)
                self.client_gpu_map[waiting_client_id] = gpu_id
                print(f"Client {waiting_client_id[:8]} assigned GPU {gpu_id} from queue")
                ready_event.set()

                await self._send_queue_updates()

    async def _send_queue_updates(self):
        """Send updated queue positions to all waiting clients."""
        for i, (cid, _, ws) in enumerate(self.waiting_list):
            if ws is not None:
                try:
                    await ws.send_json({
                        "type": "queue_status",
                        "position": i + 1,
                        "total_gpus": len(self.gpu_ids),
                        "available_gpus": 0,
                    })
                except Exception:
                    pass

    async def shutdown(self):
        """Shutdown all GPU workers."""
        print("Shutting down GPU pool...")
        tasks = [slot.shutdown() for slot in self.slots.values()]
        await asyncio.gather(*tasks, return_exceptions=True)
        print("GPU pool shutdown complete")

    def get_status(self) -> dict:
        """Get the current status of the GPU pool."""
        warmed_count = sum(1 for slot in self.slots.values() if slot.warmup_success)
        warmup_failures = sum(1 for slot in self.slots.values()
                              if slot.warmup_enabled and slot.warmup_error is not None)
        return {
            "total_gpus": len(self.gpu_ids),
            "available_gpus": sum(1 for slot in self.slots.values() if slot.is_available),
            "queue_size": len(self.waiting_list),
            "warmup_enabled": STARTUP_WARMUP_ENABLED,
            "warmup_successful_gpus": warmed_count,
            "warmup_failed_gpus": warmup_failures,
            "gpu_status": {
                gpu_id: {
                    "ready": slot.ready,
                    "available": slot.is_available,
                    "client_count": slot.client_count,
                    "current_model_id": slot.current_model_id,
                    "process_alive": (slot.process.is_alive() if slot.process else False),
                    "warmup_enabled": slot.warmup_enabled,
                    "warmup_success": slot.warmup_success,
                    "warmup_error": slot.warmup_error,
                    "warmup_timings": slot.warmup_timings,
                }
                for gpu_id, slot in self.slots.items()
            }
        }


def get_available_gpus() -> list[int]:
    """Get list of available GPU IDs from environment or auto-detect."""
    cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")
    if cuda_visible:
        visible_gpu_ids = [int(x.strip()) for x in cuda_visible.split(",") if x.strip()]
        return _limit_gpu_ids(visible_gpu_ids)

    # Auto-detect available GPUs
    try:
        result = subprocess.run(["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"],
                                capture_output=True,
                                text=True)
        if result.returncode == 0:
            detected_gpu_ids = [int(x.strip()) for x in result.stdout.strip().split("\n") if x.strip()]
            print(f"Auto-detected GPU IDs: {detected_gpu_ids}")
            return _limit_gpu_ids(detected_gpu_ids)
    except Exception:
        pass

    return _limit_gpu_ids([0])