[fix] dreamverse-monorepo: carry generic-merged files product-local
Source: Dreamverse repo at will/integrate-public-fastvideo @ ec8ef92.\n\nCarry the generic-merged backend modules as Dreamverse product-local code because the public streaming APIs are not drop-in compatible with the product runtime yet. Revert the Phase 2 public compatibility shims so FastVideo's public streaming surface returns to the pre-migration state.\n\nCo-authored-by: Junda (David) Su <90978028+Davids048@users.noreply.github.com>\nCo-authored-by: Matthew Noto <99706358+RandNMR73@users.noreply.github.com>\nCo-authored-by: XOR-op <17672363+XOR-op@users.noreply.github.com>\nCo-authored-by: Zhang Peiyuan <42993249+jzhang38@users.noreply.github.com>
This commit is contained in:
SolitaryThinker
@@ -0,0 +1,403 @@
|
||||
# pyright: reportMissingTypeArgument=false, reportArgumentType=false, reportOptionalSubscript=false, reportOptionalMemberAccess=false, reportConstantRedefinition=false, reportCallIssue=false
|
||||
# ruff: noqa: UP007, SIM108, SIM105
|
||||
# mypy: ignore-errors
|
||||
"""ffmpeg fMP4 muxing with chunk-level event emission.
|
||||
|
||||
Self-contained: spawns ffmpeg as a subprocess, pipes raw frames into
|
||||
its stdin, reads fragmented-MP4 chunks from stdout, and publishes each
|
||||
chunk as a ``StreamEvent`` via the caller-supplied ``publish``
|
||||
callback. Knows nothing about multiprocessing queues, the GPU pool,
|
||||
or individual users — the caller decides what "publish" means.
|
||||
"""
|
||||
import fcntl
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
import wave
|
||||
from dataclasses import dataclass
|
||||
from typing import Union
|
||||
from collections.abc import Callable
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
FFMPEG_BIN = shutil.which(os.getenv("FASTVIDEO_FFMPEG_BIN", "ffmpeg"))
|
||||
AV_MEDIA_MIME = os.getenv(
|
||||
"STREAM_MIME_TYPE",
|
||||
'video/mp4; codecs="avc1.42E01E,mp4a.40.2"',
|
||||
)
|
||||
AV_CHUNK_SIZE_BYTES = 1048576
|
||||
TARGET_FPS = 24
|
||||
AV_FRAGMENT_DURATION_US = int(os.getenv("FASTVIDEO_FRAG_US", "250000"))
|
||||
X264_GOP_FRAMES = int(os.getenv("FASTVIDEO_X264_GOP", "12"))
|
||||
X264_PROFILE = os.getenv("FASTVIDEO_X264_PROFILE", "baseline").strip().lower()
|
||||
if X264_PROFILE not in {"baseline", "main", "high", "main10", "high10"}:
|
||||
print(f"[WARN] Unsupported FASTVIDEO_X264_PROFILE={X264_PROFILE}; using baseline")
|
||||
X264_PROFILE = "baseline"
|
||||
USE_SHARED_STREAM_BUFFER = (os.getenv("FASTVIDEO_USE_SHARED_STREAM_BUFFER", "1").strip().lower()
|
||||
not in {"0", "false", "no"})
|
||||
SHARED_STREAM_BUFFER_BYTES = int(os.getenv("FASTVIDEO_SHARED_STREAM_BUFFER_BYTES", str(256 * 1024 * 1024)))
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamInit:
|
||||
"""First event emitted — tells the consumer the stream is starting."""
|
||||
stream_id: str
|
||||
mime: str
|
||||
uses_shared_buffer: bool
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamChunk:
|
||||
"""One fMP4 chunk. Either ``chunk`` (raw bytes) or
|
||||
``chunk_offset``+``chunk_length`` (read from the shared buffer)
|
||||
will be populated, never both."""
|
||||
stream_id: str
|
||||
chunk: bytes | None = None
|
||||
chunk_offset: int | None = None
|
||||
chunk_length: int | None = None
|
||||
uses_shared_buffer: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamComplete:
|
||||
"""Final event emitted — muxing finished successfully."""
|
||||
stream_id: str
|
||||
chunks: int
|
||||
|
||||
|
||||
StreamEvent = Union[StreamInit, StreamChunk, StreamComplete]
|
||||
|
||||
|
||||
def generate_stream_id(segment_idx: int) -> str:
|
||||
"""Convenience: build a stream id of the form ``seg007-abcd1234``."""
|
||||
return f"seg{segment_idx:03d}-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
|
||||
def _normalize_audio_tensor(audio: object) -> tuple[np.ndarray, int] | None:
|
||||
"""Convert audio tensor/array into int16 ndarray [samples, channels]."""
|
||||
if audio is None:
|
||||
return None
|
||||
|
||||
if torch.is_tensor(audio):
|
||||
audio_np = audio.detach().cpu().float().numpy()
|
||||
else:
|
||||
audio_np = np.asarray(audio, dtype=np.float32)
|
||||
|
||||
if audio_np.ndim == 1:
|
||||
audio_np = audio_np[:, None]
|
||||
elif audio_np.ndim == 2:
|
||||
if audio_np.shape[0] <= 8 and audio_np.shape[1] > audio_np.shape[0]:
|
||||
audio_np = audio_np.T
|
||||
else:
|
||||
return None
|
||||
|
||||
audio_np = np.clip(audio_np, -1.0, 1.0)
|
||||
audio_int16 = (audio_np * 32767.0).astype(np.int16)
|
||||
num_channels = audio_int16.shape[1]
|
||||
return audio_int16, num_channels
|
||||
|
||||
|
||||
def _write_audio_wav(
|
||||
audio_int16: np.ndarray,
|
||||
num_channels: int,
|
||||
sample_rate: int,
|
||||
) -> str:
|
||||
"""Write normalized int16 audio to a temporary WAV file."""
|
||||
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
|
||||
wav_path = f.name
|
||||
with wave.open(wav_path, "wb") as wav_file:
|
||||
wav_file.setnchannels(num_channels)
|
||||
wav_file.setsampwidth(2)
|
||||
wav_file.setframerate(sample_rate)
|
||||
wav_file.writeframes(audio_int16.tobytes())
|
||||
return wav_path
|
||||
|
||||
|
||||
def stream_fmp4(
|
||||
*,
|
||||
frames: list[np.ndarray],
|
||||
audio: object,
|
||||
audio_sample_rate: int | None,
|
||||
stream_id: str,
|
||||
timings: dict,
|
||||
head_trim_frames: int = 0,
|
||||
head_trim_audio_frames: int | None = None,
|
||||
shared_buffer=None,
|
||||
shared_buffer_bytes: int = 0,
|
||||
publish: Callable[[StreamEvent], None],
|
||||
log_prefix: str = "",
|
||||
) -> tuple[bool, str | None]:
|
||||
"""Encode frames+audio with ffmpeg, publish each fMP4 chunk as an event.
|
||||
|
||||
Args:
|
||||
frames: RGB24 video frames as HxWx3 uint8 arrays.
|
||||
audio: 1D/2D tensor or ndarray, float values in [-1, 1].
|
||||
audio_sample_rate: sample rate of ``audio``.
|
||||
stream_id: caller-supplied identifier carried on every event.
|
||||
timings: dict mutated in place with ffmpeg/stream timing metrics.
|
||||
head_trim_frames: video frames to drop from the start
|
||||
(conditioning overlap).
|
||||
head_trim_audio_frames: video-frame-equivalent audio to drop.
|
||||
Defaults to ``head_trim_frames``.
|
||||
shared_buffer: optional ``mp.RawArray``-compatible object; when
|
||||
provided, chunks are written into it and emitted by offset
|
||||
rather than by bytes (avoids IPC copies).
|
||||
shared_buffer_bytes: size of ``shared_buffer`` in bytes.
|
||||
publish: callback invoked once per stream event.
|
||||
log_prefix: prepended to warning prints (e.g. ``"[GPU 0]"``).
|
||||
|
||||
Returns:
|
||||
``(True, None)`` on success, ``(False, error_message)`` on
|
||||
failure. On mid-stream failure, a ``StreamInit`` may have
|
||||
already been published — the caller is responsible for
|
||||
handling that.
|
||||
"""
|
||||
if not frames:
|
||||
return False, "no frames returned"
|
||||
if audio is None:
|
||||
return False, "audio is None"
|
||||
if audio_sample_rate is None:
|
||||
return False, "audio_sample_rate is None"
|
||||
if FFMPEG_BIN is None:
|
||||
return False, "ffmpeg not found"
|
||||
|
||||
if head_trim_audio_frames is None:
|
||||
head_trim_audio_frames = head_trim_frames
|
||||
|
||||
normalized_audio = _normalize_audio_tensor(audio)
|
||||
if normalized_audio is None:
|
||||
shape_hint = getattr(audio, "shape", None)
|
||||
return False, f"unsupported audio shape={shape_hint}"
|
||||
audio_int16, num_channels = normalized_audio
|
||||
|
||||
if head_trim_frames < 0:
|
||||
return False, (f"head_trim_frames must be >= 0, "
|
||||
f"got {head_trim_frames}")
|
||||
if head_trim_frames >= len(frames):
|
||||
return False, (f"head_trim_frames={head_trim_frames} removes "
|
||||
f"all {len(frames)} frames in segment")
|
||||
|
||||
out_frames = (frames[head_trim_frames:] if head_trim_frames > 0 else frames)
|
||||
sample_rate = int(audio_sample_rate)
|
||||
if head_trim_audio_frames > 0:
|
||||
trim_start_samples = int(round((head_trim_audio_frames / float(TARGET_FPS)) * sample_rate))
|
||||
if trim_start_samples >= audio_int16.shape[0]:
|
||||
return False, ("audio too short after overlap trim: "
|
||||
f"trim_start_samples={trim_start_samples}"
|
||||
f", audio_samples={audio_int16.shape[0]}")
|
||||
keep_samples = int(round((len(out_frames) / float(TARGET_FPS)) * sample_rate))
|
||||
trim_end_samples = min(
|
||||
audio_int16.shape[0],
|
||||
trim_start_samples + keep_samples,
|
||||
)
|
||||
if trim_end_samples <= trim_start_samples:
|
||||
return False, ("invalid audio trim range: "
|
||||
f"start={trim_start_samples}, "
|
||||
f"end={trim_end_samples}")
|
||||
audio_int16 = audio_int16[trim_start_samples:trim_end_samples]
|
||||
|
||||
height = int(out_frames[0].shape[0])
|
||||
width = int(out_frames[0].shape[1])
|
||||
codec = os.getenv("FASTVIDEO_VIDEO_CODEC", "libx264")
|
||||
wav_path = _write_audio_wav(audio_int16, num_channels, sample_rate)
|
||||
|
||||
cmd = [
|
||||
FFMPEG_BIN,
|
||||
"-hide_banner",
|
||||
"-loglevel",
|
||||
"error",
|
||||
"-y",
|
||||
"-f",
|
||||
"rawvideo",
|
||||
"-pix_fmt",
|
||||
"rgb24",
|
||||
"-s:v",
|
||||
f"{width}x{height}",
|
||||
"-r",
|
||||
str(TARGET_FPS),
|
||||
"-i",
|
||||
"pipe:0",
|
||||
"-i",
|
||||
wav_path,
|
||||
"-c:v",
|
||||
codec,
|
||||
]
|
||||
if codec.endswith("_nvenc"):
|
||||
cmd += [
|
||||
"-preset",
|
||||
os.getenv("FASTVIDEO_NVENC_PRESET", "p1"),
|
||||
"-tune",
|
||||
os.getenv("FASTVIDEO_NVENC_TUNE", "ull"),
|
||||
"-rc",
|
||||
os.getenv("FASTVIDEO_NVENC_RC", "constqp"),
|
||||
"-qp",
|
||||
os.getenv("FASTVIDEO_NVENC_QP", "28"),
|
||||
"-bf",
|
||||
os.getenv("FASTVIDEO_NVENC_BF", "0"),
|
||||
]
|
||||
else:
|
||||
cmd += [
|
||||
"-preset",
|
||||
os.getenv("FASTVIDEO_X264_PRESET", "ultrafast"),
|
||||
"-tune",
|
||||
"zerolatency",
|
||||
"-profile:v",
|
||||
X264_PROFILE,
|
||||
# Emit frequent keyframes so fragments are independently playable.
|
||||
"-g",
|
||||
str(X264_GOP_FRAMES),
|
||||
"-keyint_min",
|
||||
str(X264_GOP_FRAMES),
|
||||
"-x264-params",
|
||||
"scenecut=0",
|
||||
]
|
||||
cmd += [
|
||||
"-c:a",
|
||||
"aac",
|
||||
"-pix_fmt",
|
||||
os.getenv("FASTVIDEO_OUTPUT_PIX_FMT", "yuv420p"),
|
||||
"-shortest",
|
||||
"-movflags",
|
||||
"+frag_keyframe+empty_moov+default_base_moof",
|
||||
"-frag_duration",
|
||||
str(AV_FRAGMENT_DURATION_US),
|
||||
"-flush_packets",
|
||||
"1",
|
||||
"-muxdelay",
|
||||
"0",
|
||||
"-muxpreload",
|
||||
"0",
|
||||
"-f",
|
||||
"mp4",
|
||||
"pipe:1",
|
||||
]
|
||||
|
||||
proc: subprocess.Popen | None = None
|
||||
stderr_chunks: list[bytes] = []
|
||||
writer_error: list[Exception | None] = [None]
|
||||
t_stream_start = time.perf_counter()
|
||||
use_shared_buffer = (USE_SHARED_STREAM_BUFFER and shared_buffer is not None and shared_buffer_bytes > 0)
|
||||
shared_write_offset = 0
|
||||
shared_buffer_fallback = False
|
||||
shared_np = (np.frombuffer(
|
||||
shared_buffer,
|
||||
dtype=np.uint8,
|
||||
count=shared_buffer_bytes,
|
||||
) if use_shared_buffer else None)
|
||||
|
||||
try:
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
bufsize=0,
|
||||
)
|
||||
assert proc.stdin is not None
|
||||
assert proc.stdout is not None
|
||||
assert proc.stderr is not None
|
||||
fcntl.fcntl(proc.stdin.fileno(), fcntl.F_SETPIPE_SZ, 1048576)
|
||||
fcntl.fcntl(proc.stdout.fileno(), fcntl.F_SETPIPE_SZ, 1048576)
|
||||
|
||||
def _write_frames():
|
||||
try:
|
||||
for frame in out_frames:
|
||||
proc.stdin.write(np.ascontiguousarray(frame).tobytes())
|
||||
proc.stdin.close()
|
||||
except Exception as exc:
|
||||
writer_error[0] = exc
|
||||
try:
|
||||
proc.stdin.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _read_stderr():
|
||||
try:
|
||||
while True:
|
||||
data = proc.stderr.read(4096)
|
||||
if not data:
|
||||
break
|
||||
stderr_chunks.append(data)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
writer_thread = threading.Thread(target=_write_frames, daemon=True)
|
||||
stderr_thread = threading.Thread(target=_read_stderr, daemon=True)
|
||||
writer_thread.start()
|
||||
stderr_thread.start()
|
||||
|
||||
publish(StreamInit(
|
||||
stream_id=stream_id,
|
||||
mime=AV_MEDIA_MIME,
|
||||
uses_shared_buffer=use_shared_buffer,
|
||||
))
|
||||
|
||||
chunk_count = 0
|
||||
total_bytes = 0
|
||||
while True:
|
||||
chunk = proc.stdout.read(AV_CHUNK_SIZE_BYTES)
|
||||
if not chunk:
|
||||
break
|
||||
chunk_count += 1
|
||||
total_bytes += len(chunk)
|
||||
if use_shared_buffer and not shared_buffer_fallback:
|
||||
chunk_len = len(chunk)
|
||||
write_end = shared_write_offset + chunk_len
|
||||
if write_end <= shared_buffer_bytes:
|
||||
shared_np[shared_write_offset:write_end] = np.frombuffer(chunk, dtype=np.uint8)
|
||||
publish(
|
||||
StreamChunk(
|
||||
stream_id=stream_id,
|
||||
chunk_offset=shared_write_offset,
|
||||
chunk_length=chunk_len,
|
||||
uses_shared_buffer=True,
|
||||
))
|
||||
shared_write_offset = write_end
|
||||
continue
|
||||
shared_buffer_fallback = True
|
||||
print(f"{log_prefix} Shared stream buffer exhausted at "
|
||||
f"{shared_write_offset / (1024 * 1024):.1f}MB; "
|
||||
"falling back to queue chunk bytes")
|
||||
publish(StreamChunk(
|
||||
stream_id=stream_id,
|
||||
chunk=chunk,
|
||||
))
|
||||
|
||||
writer_thread.join(timeout=5.0)
|
||||
rc = proc.wait()
|
||||
stderr_thread.join(timeout=1.0)
|
||||
|
||||
if rc != 0:
|
||||
stderr_tail = b"".join(stderr_chunks).decode(errors="ignore")[-1200:]
|
||||
return False, f"ffmpeg av_fmp4 stream failed (rc={rc}): {stderr_tail}"
|
||||
if writer_error[0] is not None:
|
||||
return False, f"ffmpeg frame writer failed: {writer_error[0]}"
|
||||
|
||||
timings["av_encode_stream_ms"] = (time.perf_counter() - t_stream_start) * 1000
|
||||
timings["av_stream_bytes"] = total_bytes
|
||||
timings["av_trim_head_frames"] = head_trim_frames
|
||||
timings["av_trim_head_audio_frames"] = head_trim_audio_frames
|
||||
timings["av_frames_encoded"] = len(out_frames)
|
||||
timings["av_shared_buffer_used"] = (bool(use_shared_buffer and not shared_buffer_fallback))
|
||||
publish(StreamComplete(
|
||||
stream_id=stream_id,
|
||||
chunks=chunk_count,
|
||||
))
|
||||
return True, None
|
||||
except Exception as exc:
|
||||
return False, str(exc)
|
||||
finally:
|
||||
if proc is not None and proc.poll() is None:
|
||||
try:
|
||||
proc.kill()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
os.remove(wav_path)
|
||||
except OSError:
|
||||
pass
|
||||
@@ -0,0 +1,952 @@
|
||||
# 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])
|
||||
@@ -9,9 +9,8 @@ from pathlib import Path
|
||||
from fastapi import FastAPI, WebSocket
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from fastvideo.entrypoints.streaming.gpu_pool import GpuPool as GPUPool
|
||||
from fastvideo.entrypoints.streaming.gpu_pool import get_available_gpus
|
||||
from fastvideo.entrypoints.streaming.session_logger import SessionEventLogger
|
||||
from gpu_pool import GPUPool, get_available_gpus
|
||||
from session_logger import SessionEventLogger
|
||||
|
||||
from config import (
|
||||
DEVTOOLS_ENABLED,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -11,8 +11,8 @@ from __future__ import annotations
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastvideo.entrypoints.streaming.gpu_pool import GpuPool as GPUPool
|
||||
from fastvideo.entrypoints.streaming.session_logger import SessionEventLogger
|
||||
from gpu_pool import GPUPool
|
||||
from session_logger import SessionEventLogger
|
||||
from prompt_enhancer import PromptEnhancer
|
||||
from prompt_safety import PromptSafetyFilter
|
||||
|
||||
|
||||
@@ -0,0 +1,14 @@
|
||||
# pyright: reportMissingTypeArgument=false
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
def cli() -> None:
|
||||
try:
|
||||
from main import cli as main_cli
|
||||
except ModuleNotFoundError as exc:
|
||||
if exc.name in {"fastvideo", "torch", "safetensors"}:
|
||||
raise SystemExit("dreamverse-server requires the Python backend extra. "
|
||||
"Run `uv sync --extra server` first.") from exc
|
||||
raise
|
||||
|
||||
main_cli()
|
||||
@@ -25,9 +25,9 @@ import uuid
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from fastapi import WebSocket, WebSocketDisconnect
|
||||
from fastvideo.entrypoints.streaming.gpu_pool import GPUSlot
|
||||
from fastvideo.entrypoints.streaming.protocol import MediaChunk, MediaComplete, MediaInit
|
||||
from fastvideo.entrypoints.streaming.session_init_image import cleanup_session_init_image, persist_session_init_image
|
||||
from gpu_pool import GPUSlot
|
||||
from session_init_image import cleanup_session_init_image, persist_session_init_image
|
||||
from worker_ipc import MediaChunk, MediaComplete, MediaInit
|
||||
|
||||
from config import (
|
||||
DEFAULT_MODEL_ID,
|
||||
@@ -47,8 +47,8 @@ from utils import (
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastvideo.entrypoints.streaming.gpu_pool import GpuPool as GPUPool
|
||||
from fastvideo.entrypoints.streaming.session_logger import SessionEventLogger
|
||||
from gpu_pool import GPUPool
|
||||
from session_logger import SessionEventLogger
|
||||
from prompt_enhancer import PromptEnhancer
|
||||
from prompt_safety import PromptSafetyFilter
|
||||
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
# pyright: reportMissingTypeArgument=false, reportArgumentType=false
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
import re
|
||||
import shutil
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image
|
||||
|
||||
MAX_SESSION_INIT_IMAGE_BYTES = 15 * 1024 * 1024
|
||||
SUPPORTED_SESSION_INIT_IMAGE_MIME_TYPES = {
|
||||
"image/jpeg",
|
||||
"image/png",
|
||||
"image/webp",
|
||||
}
|
||||
|
||||
_DATA_URL_RE = re.compile(r"^data:(?P<mime>[-\w.+/]+);base64,(?P<data>[A-Za-z0-9+/=\s]+)$")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SessionInitImage:
|
||||
file_path: Path
|
||||
temp_dir: Path
|
||||
display_name: str
|
||||
mime_type: str
|
||||
|
||||
|
||||
def _sanitize_display_name(raw_name: object) -> str:
|
||||
text = str(raw_name or "").strip()
|
||||
if not text:
|
||||
return "uploaded-image"
|
||||
return Path(text).name or "uploaded-image"
|
||||
|
||||
|
||||
def persist_session_init_image(
|
||||
payload: object,
|
||||
*,
|
||||
temp_root: Path | None = None,
|
||||
) -> SessionInitImage | None:
|
||||
if payload is None:
|
||||
return None
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError("initial_image must be an object.")
|
||||
|
||||
data_url = str(payload.get("data_url") or "").strip()
|
||||
if not data_url:
|
||||
return None
|
||||
|
||||
data_url_match = _DATA_URL_RE.match(data_url)
|
||||
if data_url_match is None:
|
||||
raise ValueError("initial_image.data_url must be a base64 data URL.")
|
||||
|
||||
data_url_mime_type = data_url_match.group("mime").strip().lower()
|
||||
mime_type = str(payload.get("mime_type") or data_url_mime_type).strip().lower()
|
||||
if mime_type != data_url_mime_type:
|
||||
raise ValueError("initial_image.mime_type must match the data URL mime type.")
|
||||
if mime_type not in SUPPORTED_SESSION_INIT_IMAGE_MIME_TYPES:
|
||||
raise ValueError("initial_image must be a PNG, JPEG, or WebP image.")
|
||||
|
||||
encoded_bytes = re.sub(r"\s+", "", data_url_match.group("data"))
|
||||
try:
|
||||
raw_bytes = base64.b64decode(encoded_bytes, validate=True)
|
||||
except Exception as exc:
|
||||
raise ValueError("initial_image.data_url is not valid base64 data.") from exc
|
||||
|
||||
if len(raw_bytes) == 0:
|
||||
raise ValueError("initial_image.data_url did not contain image bytes.")
|
||||
if len(raw_bytes) > MAX_SESSION_INIT_IMAGE_BYTES:
|
||||
raise ValueError("initial_image must be 15 MB or smaller.")
|
||||
|
||||
try:
|
||||
with Image.open(io.BytesIO(raw_bytes)) as image:
|
||||
image.load()
|
||||
normalized = image.convert("RGBA" if "A" in image.getbands() else "RGB")
|
||||
except Exception as exc:
|
||||
raise ValueError("initial_image must decode as a valid image.") from exc
|
||||
|
||||
temp_dir = Path(
|
||||
tempfile.mkdtemp(
|
||||
prefix="ltx2_session_init_",
|
||||
dir=str(temp_root) if temp_root is not None else None,
|
||||
))
|
||||
file_path = temp_dir / "initial_frame.png"
|
||||
normalized.save(file_path, format="PNG")
|
||||
|
||||
return SessionInitImage(
|
||||
file_path=file_path,
|
||||
temp_dir=temp_dir,
|
||||
display_name=_sanitize_display_name(payload.get("name")),
|
||||
mime_type=mime_type,
|
||||
)
|
||||
|
||||
|
||||
def cleanup_session_init_image(session_image: SessionInitImage | None) -> None:
|
||||
if session_image is None:
|
||||
return
|
||||
shutil.rmtree(session_image.temp_dir, ignore_errors=True)
|
||||
@@ -0,0 +1,46 @@
|
||||
# pyright: reportMissingTypeArgument=false
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime, timezone
|
||||
import json
|
||||
import socket
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
def _utc_now_iso() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
class SessionEventLogger:
|
||||
|
||||
def __init__(self, root_dir: Path):
|
||||
self.hostname = socket.gethostname()
|
||||
timestamp = datetime.now(timezone.utc).strftime("%y%m%d_%H%M%S")
|
||||
self.directory = root_dir / self.hostname
|
||||
self.path = self.directory / f"{timestamp}.jsonl"
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
self.directory.mkdir(parents=True, exist_ok=True)
|
||||
self.path.touch(exist_ok=False)
|
||||
|
||||
async def write_event(
|
||||
self,
|
||||
*,
|
||||
event: str,
|
||||
client_id: str,
|
||||
payload: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
entry = {
|
||||
"ts": _utc_now_iso(),
|
||||
"event": event,
|
||||
"hostname": self.hostname,
|
||||
"client_id": client_id,
|
||||
}
|
||||
if payload:
|
||||
entry.update(payload)
|
||||
|
||||
async with self._lock:
|
||||
with self.path.open("a", encoding="utf-8") as fp:
|
||||
fp.write(json.dumps(entry, ensure_ascii=False) + "\n")
|
||||
@@ -5,17 +5,16 @@ import sys
|
||||
from pathlib import Path
|
||||
import types
|
||||
|
||||
import fastvideo.entrypoints.streaming.mock_server as mock_server
|
||||
|
||||
|
||||
SERVER_DIR = Path(__file__).resolve().parent
|
||||
if str(SERVER_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(SERVER_DIR))
|
||||
|
||||
import main as server_main # noqa: E402
|
||||
import mock_server # noqa: E402
|
||||
|
||||
|
||||
def _run_cli(module, monkeypatch, argv: list[str], entrypoint: str = "cli") -> list[dict[str, object]]:
|
||||
def _run_cli(module, monkeypatch, argv: list[str]) -> list[dict[str, object]]:
|
||||
calls: list[dict[str, object]] = []
|
||||
uvicorn_stub = types.ModuleType("uvicorn")
|
||||
|
||||
@@ -32,7 +31,7 @@ def _run_cli(module, monkeypatch, argv: list[str], entrypoint: str = "cli") -> l
|
||||
monkeypatch.setitem(sys.modules, "uvicorn", uvicorn_stub)
|
||||
monkeypatch.setattr(sys, "argv", argv)
|
||||
|
||||
getattr(module, entrypoint)()
|
||||
module.cli()
|
||||
return calls
|
||||
|
||||
|
||||
@@ -65,52 +64,37 @@ def test_server_cli_allows_explicit_host_and_port(monkeypatch):
|
||||
|
||||
|
||||
def test_mock_server_cli_defaults_to_local_web_port(monkeypatch):
|
||||
mock_app = object()
|
||||
build_calls: list[dict[str, object]] = []
|
||||
|
||||
def build_mock_app(*, sleep_ms: float = 0.0):
|
||||
build_calls.append({"sleep_ms": sleep_ms})
|
||||
return mock_app
|
||||
|
||||
monkeypatch.setattr(mock_server, "build_mock_app", build_mock_app)
|
||||
calls = _run_cli(
|
||||
mock_server,
|
||||
monkeypatch,
|
||||
["dreamverse-mock-server"],
|
||||
entrypoint="main",
|
||||
)
|
||||
|
||||
assert build_calls == [{"sleep_ms": 0.0}]
|
||||
assert calls == [
|
||||
{
|
||||
"app": mock_app,
|
||||
"host": "127.0.0.1",
|
||||
"port": 8000,
|
||||
"app": mock_server.app,
|
||||
"host": "0.0.0.0",
|
||||
"port": 8009,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_mock_server_cli_updates_latency(monkeypatch):
|
||||
mock_app = object()
|
||||
build_calls: list[dict[str, object]] = []
|
||||
old_latency_ms = mock_server.LATENCY_MS
|
||||
try:
|
||||
calls = _run_cli(
|
||||
mock_server,
|
||||
monkeypatch,
|
||||
["dreamverse-mock-server", "--latency", "321", "--port", "8111"],
|
||||
)
|
||||
|
||||
def build_mock_app(*, sleep_ms: float = 0.0):
|
||||
build_calls.append({"sleep_ms": sleep_ms})
|
||||
return mock_app
|
||||
|
||||
monkeypatch.setattr(mock_server, "build_mock_app", build_mock_app)
|
||||
calls = _run_cli(
|
||||
mock_server,
|
||||
monkeypatch,
|
||||
["dreamverse-mock-server", "--sleep-ms", "321", "--port", "8111"],
|
||||
entrypoint="main",
|
||||
)
|
||||
|
||||
assert build_calls == [{"sleep_ms": 321.0}]
|
||||
assert calls == [
|
||||
{
|
||||
"app": mock_app,
|
||||
"host": "127.0.0.1",
|
||||
"port": 8111,
|
||||
}
|
||||
]
|
||||
assert calls == [
|
||||
{
|
||||
"app": mock_server.app,
|
||||
"host": "0.0.0.0",
|
||||
"port": 8111,
|
||||
}
|
||||
]
|
||||
assert mock_server.LATENCY_MS == 321
|
||||
finally:
|
||||
mock_server.LATENCY_MS = old_latency_ms
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
# pyright: reportAttributeAccessIssue=false
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import multiprocessing as mp
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
SERVER_DIR = Path(__file__).resolve().parent
|
||||
if str(SERVER_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(SERVER_DIR))
|
||||
|
||||
import gpu_pool
|
||||
|
||||
|
||||
def _child_consume_and_exit(cmd_q, resp_q):
|
||||
"""Top-level so the spawn context can pickle it.
|
||||
|
||||
Signals startup by putting "READY" on resp_q (so the parent can
|
||||
wait out spawn-import latency separately from the actual test
|
||||
assertion), then consumes one command from cmd_q and exits without
|
||||
putting anything else on resp_q. Simulates a worker that dies
|
||||
mid-command — e.g. SIGQUIT'd after a fatal pipeline error.
|
||||
"""
|
||||
resp_q.put("READY")
|
||||
cmd_q.get()
|
||||
|
||||
|
||||
def test_get_available_gpus_defaults_to_single_detected_gpu(monkeypatch):
|
||||
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
monkeypatch.delenv("FASTVIDEO_GPU_COUNT", raising=False)
|
||||
|
||||
def fake_run(*args, **kwargs):
|
||||
del args, kwargs
|
||||
return SimpleNamespace(returncode=0, stdout="0\n1\n2\n")
|
||||
|
||||
monkeypatch.setattr(gpu_pool.subprocess, "run", fake_run)
|
||||
|
||||
assert gpu_pool.get_available_gpus() == [0]
|
||||
|
||||
|
||||
def test_get_available_gpus_respects_explicit_gpu_count(monkeypatch):
|
||||
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
monkeypatch.setenv("FASTVIDEO_GPU_COUNT", "2")
|
||||
|
||||
def fake_run(*args, **kwargs):
|
||||
del args, kwargs
|
||||
return SimpleNamespace(returncode=0, stdout="0\n1\n2\n")
|
||||
|
||||
monkeypatch.setattr(gpu_pool.subprocess, "run", fake_run)
|
||||
|
||||
assert gpu_pool.get_available_gpus() == [0, 1]
|
||||
|
||||
|
||||
def test_get_available_gpus_can_use_all_visible_devices(monkeypatch):
|
||||
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,5,7")
|
||||
monkeypatch.setenv("FASTVIDEO_GPU_COUNT", "all")
|
||||
|
||||
assert gpu_pool.get_available_gpus() == [3, 5, 7]
|
||||
|
||||
|
||||
def test_get_available_gpus_defaults_to_first_visible_device(monkeypatch):
|
||||
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,5,7")
|
||||
monkeypatch.delenv("FASTVIDEO_GPU_COUNT", raising=False)
|
||||
|
||||
assert gpu_pool.get_available_gpus() == [3]
|
||||
|
||||
|
||||
def test_get_available_gpus_rejects_invalid_gpu_count(monkeypatch):
|
||||
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
monkeypatch.setenv("FASTVIDEO_GPU_COUNT", "zero")
|
||||
|
||||
with pytest.raises(RuntimeError, match="Invalid FASTVIDEO_GPU_COUNT"):
|
||||
gpu_pool.get_available_gpus()
|
||||
|
||||
|
||||
def test_send_command_raises_on_worker_death():
|
||||
"""A worker that consumes a command and exits without replying must
|
||||
surface as RuntimeError via sentinel detection, not after the long
|
||||
queue timeout.
|
||||
|
||||
The child signals startup with a "READY" message; we drain that
|
||||
first so the wait_for budget bounds only the sentinel-detection
|
||||
time, not spawn-import latency.
|
||||
"""
|
||||
ctx = mp.get_context("spawn")
|
||||
cmd_q = ctx.Queue()
|
||||
resp_q = ctx.Queue()
|
||||
|
||||
proc = ctx.Process(
|
||||
target=_child_consume_and_exit, args=(cmd_q, resp_q)
|
||||
)
|
||||
proc.start()
|
||||
|
||||
# Wait for the spawn child to fully boot. Allow generous time —
|
||||
# this isn't what we're measuring.
|
||||
ready = resp_q.get(timeout=30.0)
|
||||
assert ready == "READY"
|
||||
|
||||
async def runner():
|
||||
slot = gpu_pool.GPUSlot(gpu_id=0, cuda_device="0")
|
||||
slot.process = proc
|
||||
slot.command_queue = cmd_q
|
||||
slot.response_queue = resp_q
|
||||
|
||||
# Now the child is blocked in cmd_q.get(). Sending the cmd
|
||||
# makes it exit ~immediately; sentinel must fire well before
|
||||
# the queue timeout.
|
||||
await asyncio.wait_for(
|
||||
slot._send_command(
|
||||
gpu_pool.Command(gpu_pool.CommandType.SHUTDOWN),
|
||||
timeout=10.0,
|
||||
),
|
||||
timeout=5.0,
|
||||
)
|
||||
|
||||
try:
|
||||
with pytest.raises(RuntimeError, match="worker died"):
|
||||
asyncio.run(runner())
|
||||
finally:
|
||||
proc.join(timeout=5)
|
||||
cmd_q.close()
|
||||
resp_q.close()
|
||||
@@ -0,0 +1,361 @@
|
||||
# pyright: reportArgumentType=false
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi import WebSocketDisconnect
|
||||
|
||||
|
||||
SERVER_DIR = Path(__file__).resolve().parent
|
||||
if str(SERVER_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(SERVER_DIR))
|
||||
|
||||
os.environ.setdefault("CEREBRAS_API_KEY", "dummy")
|
||||
os.environ.setdefault("GROQ_API_KEY", "dummy")
|
||||
|
||||
import mock_server
|
||||
|
||||
|
||||
class _FakeWebSocket:
|
||||
def __init__(self, messages: list[tuple[float, dict[str, object]]]):
|
||||
self._messages = messages
|
||||
self._index = 0
|
||||
self.sent_json: list[dict[str, object]] = []
|
||||
self.sent_bytes: list[bytes] = []
|
||||
|
||||
async def accept(self) -> None:
|
||||
return None
|
||||
|
||||
async def send_json(self, payload: dict[str, object]) -> None:
|
||||
self.sent_json.append(payload)
|
||||
|
||||
async def send_bytes(self, payload: bytes) -> None:
|
||||
self.sent_bytes.append(payload)
|
||||
|
||||
async def close(self, code: int = 1000, reason: str | None = None) -> None:
|
||||
del code, reason
|
||||
|
||||
async def receive_json(self) -> dict[str, object]:
|
||||
if self._index >= len(self._messages):
|
||||
raise WebSocketDisconnect()
|
||||
delay_s, payload = self._messages[self._index]
|
||||
self._index += 1
|
||||
if delay_s > 0:
|
||||
await asyncio.sleep(delay_s)
|
||||
return payload
|
||||
|
||||
|
||||
def test_mock_server_matches_current_single5s_protocol():
|
||||
old_segment_bytes = mock_server.MOCK_SEGMENT_BYTES
|
||||
old_latency_ms = mock_server.LATENCY_MS
|
||||
try:
|
||||
mock_server.MOCK_SEGMENT_BYTES = b"mock-fmp4-bytes"
|
||||
mock_server.LATENCY_MS = 1
|
||||
|
||||
ws = _FakeWebSocket(
|
||||
[
|
||||
(
|
||||
0.0,
|
||||
{
|
||||
"type": "session_init_v2",
|
||||
"preset_id": "simple_prompt_1",
|
||||
"curated_prompts": ["selected prompt"],
|
||||
"single_clip_mode": True,
|
||||
"enhancement_enabled": False,
|
||||
"auto_extension_enabled": False,
|
||||
"loop_generation_enabled": False,
|
||||
},
|
||||
),
|
||||
(
|
||||
0.01,
|
||||
{
|
||||
"type": "simple_generate",
|
||||
"preset_id": "simple_custom_prompt",
|
||||
"prompt_id": "simple_custom_prompt",
|
||||
"prompt": "custom prompt",
|
||||
"enhancement_enabled": True,
|
||||
"initial_image": None,
|
||||
},
|
||||
),
|
||||
(0.20, {"type": "leave"}),
|
||||
]
|
||||
)
|
||||
|
||||
asyncio.run(mock_server.websocket_endpoint(ws))
|
||||
|
||||
message_types = [payload["type"] for payload in ws.sent_json]
|
||||
assert message_types[0] == "queue_status"
|
||||
assert "gpu_assigned" in message_types
|
||||
assert message_types.count("ltx2_stream_start") == 2
|
||||
assert "prompt_received" in message_types
|
||||
assert "prompt_enhancing" in message_types
|
||||
assert "prompt_ready" in message_types
|
||||
assert "media_init" in message_types
|
||||
assert "media_segment_complete" in message_types
|
||||
assert message_types.count("ltx2_stream_complete") == 2
|
||||
assert "prompt_sources_blocked" not in message_types
|
||||
|
||||
segment_start_events = [
|
||||
payload
|
||||
for payload in ws.sent_json
|
||||
if payload["type"] == "ltx2_segment_start"
|
||||
]
|
||||
gpu_assigned_event = next(
|
||||
payload for payload in ws.sent_json if payload["type"] == "gpu_assigned"
|
||||
)
|
||||
assert gpu_assigned_event["session_timeout"] == mock_server.SESSION_TIMEOUT_SECONDS
|
||||
assert [payload["segment_idx"] for payload in segment_start_events] == [1, 1]
|
||||
assert segment_start_events[0]["prompt"] == "selected prompt"
|
||||
assert segment_start_events[1]["prompt"] == "custom prompt"
|
||||
|
||||
step_complete_events = [
|
||||
payload
|
||||
for payload in ws.sent_json
|
||||
if payload["type"] == "step_complete"
|
||||
]
|
||||
assert len(step_complete_events) == 2
|
||||
assert step_complete_events[0]["latency_ms"] == {
|
||||
"total": 121.0,
|
||||
"worker_e2e": 1.0,
|
||||
"main_user_step": 121.0,
|
||||
"overhead": 120.0,
|
||||
}
|
||||
|
||||
assert ws.sent_bytes
|
||||
assert ws.sent_bytes[0] == b"mock-fmp4-bytes"
|
||||
finally:
|
||||
mock_server.MOCK_SEGMENT_BYTES = old_segment_bytes
|
||||
mock_server.LATENCY_MS = old_latency_ms
|
||||
|
||||
|
||||
def test_mock_server_regular_cap_waits_for_rewrite_rollout():
|
||||
old_segment_bytes = mock_server.MOCK_SEGMENT_BYTES
|
||||
old_latency_ms = mock_server.LATENCY_MS
|
||||
old_generation_segment_cap = mock_server.GENERATION_SEGMENT_CAP
|
||||
try:
|
||||
mock_server.MOCK_SEGMENT_BYTES = b"mock-fmp4-bytes"
|
||||
mock_server.LATENCY_MS = 1
|
||||
mock_server.GENERATION_SEGMENT_CAP = 1
|
||||
|
||||
ws = _FakeWebSocket(
|
||||
[
|
||||
(
|
||||
0.0,
|
||||
{
|
||||
"type": "session_init_v2",
|
||||
"preset_id": "test_preset",
|
||||
"curated_prompts": ["segment one"],
|
||||
"enhancement_enabled": True,
|
||||
"auto_extension_enabled": False,
|
||||
"loop_generation_enabled": False,
|
||||
},
|
||||
),
|
||||
(
|
||||
0.02,
|
||||
{
|
||||
"type": "rewrite_seed_prompts",
|
||||
"rewrite_instruction": "start a new rollout",
|
||||
},
|
||||
),
|
||||
(0.20, {"type": "leave"}),
|
||||
]
|
||||
)
|
||||
|
||||
asyncio.run(mock_server.websocket_endpoint(ws))
|
||||
|
||||
message_types = [payload["type"] for payload in ws.sent_json]
|
||||
assert message_types.count("ltx2_stream_start") == 2
|
||||
assert message_types.count("ltx2_stream_complete") == 2
|
||||
assert "generation_cap_reached" not in message_types
|
||||
assert "prompt_sources_blocked" not in message_types
|
||||
|
||||
segment_start_events = [
|
||||
payload
|
||||
for payload in ws.sent_json
|
||||
if payload["type"] == "ltx2_segment_start"
|
||||
]
|
||||
assert [payload["segment_idx"] for payload in segment_start_events] == [1, 1]
|
||||
assert segment_start_events[0]["prompt"] == "segment one"
|
||||
assert segment_start_events[1]["prompt"] == "segment one [start a new rollout]"
|
||||
finally:
|
||||
mock_server.MOCK_SEGMENT_BYTES = old_segment_bytes
|
||||
mock_server.LATENCY_MS = old_latency_ms
|
||||
mock_server.GENERATION_SEGMENT_CAP = old_generation_segment_cap
|
||||
|
||||
|
||||
def test_mock_server_rewrite_during_active_segment_restarts_from_first_rewritten_prompt():
|
||||
old_segment_bytes = mock_server.MOCK_SEGMENT_BYTES
|
||||
old_latency_ms = mock_server.LATENCY_MS
|
||||
try:
|
||||
mock_server.MOCK_SEGMENT_BYTES = b"mock-fmp4-bytes"
|
||||
mock_server.LATENCY_MS = 100
|
||||
|
||||
ws = _FakeWebSocket(
|
||||
[
|
||||
(
|
||||
0.0,
|
||||
{
|
||||
"type": "session_init_v2",
|
||||
"preset_id": "test_preset",
|
||||
"curated_prompts": ["segment one", "segment two"],
|
||||
"enhancement_enabled": True,
|
||||
"auto_extension_enabled": False,
|
||||
"loop_generation_enabled": False,
|
||||
},
|
||||
),
|
||||
(
|
||||
0.02,
|
||||
{
|
||||
"type": "rewrite_seed_prompts",
|
||||
"rewrite_instruction": "restart from rewrite",
|
||||
},
|
||||
),
|
||||
(0.40, {"type": "leave"}),
|
||||
]
|
||||
)
|
||||
|
||||
asyncio.run(mock_server.websocket_endpoint(ws))
|
||||
|
||||
segment_start_events = [
|
||||
payload
|
||||
for payload in ws.sent_json
|
||||
if payload["type"] == "ltx2_segment_start"
|
||||
]
|
||||
assert [payload["prompt"] for payload in segment_start_events[:2]] == [
|
||||
"segment one",
|
||||
"segment one [restart from rewrite]",
|
||||
]
|
||||
assert all(
|
||||
payload["prompt"] != "segment two"
|
||||
for payload in segment_start_events[1:]
|
||||
)
|
||||
reset_events = [
|
||||
payload
|
||||
for payload in ws.sent_json
|
||||
if payload.get("type") == "seed_prompts_reset_applied"
|
||||
]
|
||||
assert any(
|
||||
payload.get("reason") == "rewrite_during_generation"
|
||||
for payload in reset_events
|
||||
)
|
||||
finally:
|
||||
mock_server.MOCK_SEGMENT_BYTES = old_segment_bytes
|
||||
mock_server.LATENCY_MS = old_latency_ms
|
||||
|
||||
|
||||
def test_mock_server_supports_initial_custom_rollout_prompt():
|
||||
old_segment_bytes = mock_server.MOCK_SEGMENT_BYTES
|
||||
old_latency_ms = mock_server.LATENCY_MS
|
||||
try:
|
||||
mock_server.MOCK_SEGMENT_BYTES = b"mock-fmp4-bytes"
|
||||
mock_server.LATENCY_MS = 1
|
||||
|
||||
ws = _FakeWebSocket(
|
||||
[
|
||||
(
|
||||
0.0,
|
||||
{
|
||||
"type": "session_init_v2",
|
||||
"preset_id": "custom_editable",
|
||||
"preset_label": "Custom rollout",
|
||||
"curated_prompts": [],
|
||||
"initial_rollout_prompt": "A moonbase corridor thriller with flooding",
|
||||
"enhancement_enabled": True,
|
||||
"auto_extension_enabled": False,
|
||||
"loop_generation_enabled": False,
|
||||
},
|
||||
),
|
||||
(0.20, {"type": "leave"}),
|
||||
]
|
||||
)
|
||||
|
||||
asyncio.run(mock_server.websocket_endpoint(ws))
|
||||
|
||||
message_types = [payload["type"] for payload in ws.sent_json]
|
||||
assert "rewrite_seed_prompts_started" in message_types
|
||||
assert "seed_prompts_updated" in message_types
|
||||
assert "rewrite_seed_prompts_complete" in message_types
|
||||
assert "seed_prompts_reset_applied" in message_types
|
||||
assert "ltx2_stream_start" in message_types
|
||||
assert "prompt_sources_blocked" not in message_types
|
||||
|
||||
segment_start_events = [
|
||||
payload
|
||||
for payload in ws.sent_json
|
||||
if payload["type"] == "ltx2_segment_start"
|
||||
]
|
||||
assert segment_start_events
|
||||
assert segment_start_events[0]["prompt"] == (
|
||||
"A moonbase corridor thriller with flooding [segment 1]"
|
||||
)
|
||||
finally:
|
||||
mock_server.MOCK_SEGMENT_BYTES = old_segment_bytes
|
||||
mock_server.LATENCY_MS = old_latency_ms
|
||||
|
||||
|
||||
def test_mock_server_can_start_new_project_without_reconnecting():
|
||||
old_segment_bytes = mock_server.MOCK_SEGMENT_BYTES
|
||||
old_latency_ms = mock_server.LATENCY_MS
|
||||
try:
|
||||
mock_server.MOCK_SEGMENT_BYTES = b"mock-fmp4-bytes"
|
||||
mock_server.LATENCY_MS = 40
|
||||
|
||||
ws = _FakeWebSocket(
|
||||
[
|
||||
(
|
||||
0.0,
|
||||
{
|
||||
"type": "session_init_v2",
|
||||
"preset_id": "test_preset",
|
||||
"curated_prompts": ["segment one"],
|
||||
"enhancement_enabled": True,
|
||||
"auto_extension_enabled": False,
|
||||
"loop_generation_enabled": False,
|
||||
},
|
||||
),
|
||||
(0.02, {"type": "end_project_keep_session"}),
|
||||
(
|
||||
0.20,
|
||||
{
|
||||
"type": "project_init_v1",
|
||||
"preset_id": "test_preset_2",
|
||||
"preset_label": "Test Preset 2",
|
||||
"curated_prompts": ["segment two"],
|
||||
"enhancement_enabled": True,
|
||||
"auto_extension_enabled": False,
|
||||
"loop_generation_enabled": False,
|
||||
},
|
||||
),
|
||||
(0.40, {"type": "leave"}),
|
||||
]
|
||||
)
|
||||
|
||||
asyncio.run(mock_server.websocket_endpoint(ws))
|
||||
|
||||
message_types = [payload["type"] for payload in ws.sent_json]
|
||||
assert message_types.count("gpu_assigned") == 1
|
||||
assert message_types.count("ltx2_stream_start") == 2
|
||||
assert "project_idle" in message_types
|
||||
|
||||
project_idle_index = message_types.index("project_idle")
|
||||
stream_start_indexes = [
|
||||
index for index, message_type in enumerate(message_types)
|
||||
if message_type == "ltx2_stream_start"
|
||||
]
|
||||
assert stream_start_indexes[0] < project_idle_index < stream_start_indexes[1]
|
||||
|
||||
segment_start_events = [
|
||||
payload
|
||||
for payload in ws.sent_json
|
||||
if payload["type"] == "ltx2_segment_start"
|
||||
]
|
||||
assert [payload["prompt"] for payload in segment_start_events[:2]] == [
|
||||
"segment one",
|
||||
"segment two",
|
||||
]
|
||||
finally:
|
||||
mock_server.MOCK_SEGMENT_BYTES = old_segment_bytes
|
||||
mock_server.LATENCY_MS = old_latency_ms
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,75 @@
|
||||
# pyright: reportMissingTypeArgument=false
|
||||
import base64
|
||||
import io
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image
|
||||
import pytest
|
||||
|
||||
from session_init_image import (
|
||||
MAX_SESSION_INIT_IMAGE_BYTES,
|
||||
cleanup_session_init_image,
|
||||
persist_session_init_image,
|
||||
)
|
||||
|
||||
|
||||
def make_data_url(format_name: str = "PNG") -> str:
|
||||
image = Image.new("RGB", (2, 2), color=(16, 32, 64))
|
||||
buffer = io.BytesIO()
|
||||
image.save(buffer, format=format_name)
|
||||
encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
|
||||
mime_type = "image/png" if format_name == "PNG" else "image/jpeg"
|
||||
return f"data:{mime_type};base64,{encoded}"
|
||||
|
||||
|
||||
def test_persist_session_init_image_saves_normalized_file(tmp_path: Path):
|
||||
session_image = persist_session_init_image(
|
||||
{
|
||||
"name": "frame.png",
|
||||
"mime_type": "image/png",
|
||||
"data_url": make_data_url("PNG"),
|
||||
},
|
||||
temp_root=tmp_path,
|
||||
)
|
||||
|
||||
assert session_image is not None
|
||||
assert session_image.display_name == "frame.png"
|
||||
assert session_image.file_path.is_file()
|
||||
|
||||
cleanup_session_init_image(session_image)
|
||||
assert not session_image.temp_dir.exists()
|
||||
|
||||
|
||||
def test_persist_session_init_image_returns_none_when_missing_data():
|
||||
assert persist_session_init_image(None) is None
|
||||
assert persist_session_init_image({"name": "frame.png", "data_url": ""}) is None
|
||||
|
||||
|
||||
def test_persist_session_init_image_rejects_unsupported_mime():
|
||||
with pytest.raises(ValueError, match="PNG, JPEG, or WebP"):
|
||||
persist_session_init_image(
|
||||
{
|
||||
"name": "frame.gif",
|
||||
"mime_type": "image/gif",
|
||||
"data_url": "data:image/gif;base64,R0lGODlhAQABAAAAACw=",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_persist_session_init_image_rejects_large_payload(monkeypatch):
|
||||
data_url = make_data_url("PNG")
|
||||
oversized = "a" * (MAX_SESSION_INIT_IMAGE_BYTES + 1)
|
||||
|
||||
def fake_b64decode(value: str, validate: bool = True):
|
||||
return oversized.encode("ascii")
|
||||
|
||||
monkeypatch.setattr(base64, "b64decode", fake_b64decode)
|
||||
|
||||
with pytest.raises(ValueError, match="15 MB or smaller"):
|
||||
persist_session_init_image(
|
||||
{
|
||||
"name": "frame.png",
|
||||
"mime_type": "image/png",
|
||||
"data_url": data_url,
|
||||
}
|
||||
)
|
||||
@@ -13,8 +13,6 @@ from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
from fastapi import WebSocketDisconnect
|
||||
from fastvideo.entrypoints.streaming.protocol import MediaChunk, MediaComplete, MediaInit
|
||||
from fastvideo.entrypoints.streaming.session_logger import SessionEventLogger
|
||||
from PIL import Image
|
||||
import pytest
|
||||
|
||||
@@ -28,6 +26,8 @@ os.environ.setdefault("GROQ_API_KEY", "dummy")
|
||||
|
||||
import main as server_main # noqa: E402
|
||||
import runtime # noqa: E402
|
||||
from session_logger import SessionEventLogger # noqa: E402
|
||||
from worker_ipc import MediaChunk, MediaComplete, MediaInit # noqa: E402
|
||||
from session import controller as session_controller # noqa: E402
|
||||
|
||||
pytestmark = pytest.mark.gpu
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
# pyright: reportMissingTypeArgument=false
|
||||
"""Typed worker IPC events: replaces the fat ``Response`` dataclass.
|
||||
|
||||
Each message kind from GPU worker → main process is its own small
|
||||
dataclass; ``WorkerEvent`` is the union of them. Consumers dispatch
|
||||
via ``match``/``case`` or ``isinstance`` — mirroring the existing
|
||||
``StreamEvent`` pattern in ``av_streaming.py``.
|
||||
|
||||
Invalid states are unrepresentable: a ``MediaChunk`` simply has no
|
||||
``frames`` field, a ``StepComplete`` has no ``chunk_offset`` field,
|
||||
and a ``JoinAck`` can't accidentally default to ``kind="step_result"``
|
||||
because there is no ``kind`` string.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
# ---- User-scoped events (carry user_id) ------------------------------------
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StepComplete:
|
||||
"""Generation step finished. Frames/audio were already sent via ``MediaChunk``."""
|
||||
user_id: str
|
||||
segment_idx: int
|
||||
timings: dict[str, float]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WorkerError:
|
||||
"""Any worker failure. ``user_id`` is None for system-scoped failures."""
|
||||
user_id: str | None
|
||||
message: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class JoinAck:
|
||||
user_id: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LeaveAck:
|
||||
user_id: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ReloadAck:
|
||||
user_id: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WarmupComplete:
|
||||
user_id: str | None
|
||||
timings: dict[str, float]
|
||||
|
||||
|
||||
# ---- Streaming events (carry user_id + segment_idx for routing) ------------
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MediaInit:
|
||||
user_id: str
|
||||
segment_idx: int
|
||||
stream_id: str
|
||||
mime: str
|
||||
uses_shared_buffer: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MediaChunk:
|
||||
"""One fMP4 chunk.
|
||||
|
||||
Either ``chunk`` (raw bytes) is set, or both ``chunk_offset`` and
|
||||
``chunk_length`` are set (read from the shared buffer). The
|
||||
invariant is enforced in ``__post_init__``.
|
||||
"""
|
||||
user_id: str
|
||||
segment_idx: int
|
||||
stream_id: str
|
||||
chunk: bytes | None = None
|
||||
chunk_offset: int | None = None
|
||||
chunk_length: int | None = None
|
||||
uses_shared_buffer: bool = False
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
has_bytes = self.chunk is not None
|
||||
has_offset = (self.chunk_offset is not None and self.chunk_length is not None)
|
||||
if has_bytes == has_offset:
|
||||
raise ValueError("MediaChunk must carry either chunk bytes or "
|
||||
"(chunk_offset + chunk_length), not both or neither")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MediaComplete:
|
||||
user_id: str
|
||||
segment_idx: int
|
||||
stream_id: str
|
||||
chunks: int
|
||||
|
||||
|
||||
# ---- System events (no user_id) --------------------------------------------
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class InitAck:
|
||||
success: bool
|
||||
error: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ShutdownAck:
|
||||
pass
|
||||
|
||||
|
||||
WorkerEvent = (StepComplete
|
||||
| WorkerError
|
||||
| JoinAck
|
||||
| LeaveAck
|
||||
| ReloadAck
|
||||
| WarmupComplete
|
||||
| MediaInit
|
||||
| MediaChunk
|
||||
| MediaComplete
|
||||
| InitAck
|
||||
| ShutdownAck)
|
||||
|
||||
# ---- Command payloads (main process → worker) ------------------------------
|
||||
#
|
||||
# The envelope (``Command`` + ``CommandType``) lives in ``gpu_pool.py``
|
||||
# alongside the enum; only the typed payloads live here so they sit
|
||||
# next to the response-side types. Commands without a payload
|
||||
# (``INIT``, ``SHUTDOWN``, ``USER_JOIN``, ``USER_LEAVE``) leave
|
||||
# ``Command.payload`` as ``None``.
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class UserStepPayload:
|
||||
prompt: str
|
||||
segment_idx: int
|
||||
image_path: str | None
|
||||
reset_conditioning: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WarmupPayload:
|
||||
prompt: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ReloadModelPayload:
|
||||
# ``model_config`` stays a dict because ``MODEL_REGISTRY`` values
|
||||
# are dicts shaped by the external config module; typing that dict
|
||||
# is a separate refactor.
|
||||
model_config: dict
|
||||
|
||||
|
||||
CommandPayload = UserStepPayload | WarmupPayload | ReloadModelPayload
|
||||
@@ -1,5 +1,4 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# pyright: reportMissingTypeArgument=false
|
||||
"""GPU pool manager for the streaming server.
|
||||
|
||||
Replaces the single-generator path in PR 7.5 with a typed pool
|
||||
@@ -27,9 +26,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import queue
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
@@ -74,13 +71,6 @@ class PoolAssignment:
|
||||
pinned_at: float = field(default_factory=time.monotonic)
|
||||
|
||||
|
||||
# Backwards-compatible name for the Dreamverse product controller during
|
||||
# the monorepo migration. The controller only uses this name for runtime
|
||||
# annotations; Phase 4 replaces the product-local controller path with the
|
||||
# generic streaming runtime.
|
||||
GPUSlot = PoolAssignment
|
||||
|
||||
|
||||
class GpuPool(ABC):
|
||||
"""Abstract GPU pool.
|
||||
|
||||
@@ -540,32 +530,7 @@ def _default_worker_factory(
|
||||
)
|
||||
|
||||
|
||||
def get_available_gpus() -> list[int]:
|
||||
"""Return CUDA-visible GPU ids, falling back to ``nvidia-smi`` then ``[0]``."""
|
||||
|
||||
cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")
|
||||
if cuda_visible:
|
||||
return [int(item.strip()) for item in cuda_visible.split(",") if item.strip()]
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
except OSError:
|
||||
return [0]
|
||||
|
||||
if result.returncode != 0:
|
||||
return [0]
|
||||
gpu_ids = [int(line.strip()) for line in result.stdout.splitlines() if line.strip()]
|
||||
return gpu_ids or [0]
|
||||
|
||||
|
||||
__all__ = [
|
||||
"get_available_gpus",
|
||||
"GPUSlot",
|
||||
"GpuPool",
|
||||
"InProcessGpuPool",
|
||||
"PoolAcquireTimeout",
|
||||
|
||||
@@ -11,7 +11,6 @@ both places in the same PR.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Annotated, Any, Literal, Union
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
@@ -143,8 +142,6 @@ class StepComplete(BaseModel):
|
||||
class MediaInit(BaseModel):
|
||||
"""Descriptor for the fMP4 initialization segment that follows."""
|
||||
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
type: Literal["media_init"] = "media_init"
|
||||
segment_idx: int
|
||||
mime: str = "video/mp4; codecs=\"avc1.64001f, mp4a.40.2\""
|
||||
@@ -152,36 +149,6 @@ class MediaInit(BaseModel):
|
||||
mode: Literal["av_fmp4"] = "av_fmp4"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MediaChunk:
|
||||
"""One fMP4 chunk emitted by the legacy Dreamverse worker bridge."""
|
||||
|
||||
user_id: str
|
||||
segment_idx: int
|
||||
stream_id: str
|
||||
chunk: bytes | None = None
|
||||
chunk_offset: int | None = None
|
||||
chunk_length: int | None = None
|
||||
uses_shared_buffer: bool = False
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
has_bytes = self.chunk is not None
|
||||
has_offset = self.chunk_offset is not None and self.chunk_length is not None
|
||||
if has_bytes == has_offset:
|
||||
raise ValueError("MediaChunk must carry either chunk bytes or "
|
||||
"(chunk_offset + chunk_length), not both or neither")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MediaComplete:
|
||||
"""Legacy Dreamverse fMP4 segment-complete event."""
|
||||
|
||||
user_id: str
|
||||
segment_idx: int
|
||||
stream_id: str
|
||||
chunks: int
|
||||
|
||||
|
||||
class MediaSegmentComplete(BaseModel):
|
||||
type: Literal["media_segment_complete"] = "media_segment_complete"
|
||||
segment_idx: int
|
||||
@@ -272,8 +239,6 @@ __all__ = [
|
||||
"Ltx2StreamStart",
|
||||
"LoopGenerationUpdated",
|
||||
"MediaInit",
|
||||
"MediaChunk",
|
||||
"MediaComplete",
|
||||
"MediaSegmentComplete",
|
||||
"QueueStatus",
|
||||
"SeedPromptsUpdated",
|
||||
|
||||
@@ -6,9 +6,6 @@ import base64
|
||||
import binascii
|
||||
import contextlib
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
import shutil
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
@@ -21,7 +18,6 @@ _ACCEPTED_MIMES = {
|
||||
}
|
||||
|
||||
_MAX_IMAGE_BYTES = 32 * 1024 * 1024 # 32 MiB cap
|
||||
_DATA_URL_RE = re.compile(r"^data:(?P<mime>[-\w.+/]+);base64,(?P<data>[A-Za-z0-9+/=\s]+)$")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -36,24 +32,6 @@ class SessionInitImage:
|
||||
display_name: str
|
||||
mime: str
|
||||
|
||||
@property
|
||||
def file_path(self) -> Path:
|
||||
"""Compatibility alias used by the migrated Dreamverse controller."""
|
||||
|
||||
return Path(self.path)
|
||||
|
||||
@property
|
||||
def temp_dir(self) -> Path:
|
||||
"""Compatibility alias for cleanup of the temporary image directory."""
|
||||
|
||||
return Path(self.path).parent
|
||||
|
||||
@property
|
||||
def mime_type(self) -> str:
|
||||
"""Compatibility alias for the UI payload's MIME field."""
|
||||
|
||||
return self.mime
|
||||
|
||||
|
||||
def persist_session_init_image(
|
||||
payload: Any,
|
||||
@@ -79,19 +57,11 @@ def persist_session_init_image(
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError("session init image must be an object")
|
||||
|
||||
data_url = payload.get("data_url")
|
||||
if isinstance(data_url, str) and data_url.strip():
|
||||
match = _DATA_URL_RE.match(data_url.strip())
|
||||
if match is None:
|
||||
raise ValueError("session init image data URL is not valid base64")
|
||||
mime = payload.get("mime_type") or match.group("mime").strip().lower()
|
||||
data_b64 = match.group("data")
|
||||
else:
|
||||
mime = payload.get("mime")
|
||||
data_b64 = payload.get("data")
|
||||
mime = payload.get("mime")
|
||||
if mime not in _ACCEPTED_MIMES:
|
||||
raise ValueError(f"session init image mime {mime!r} is not one of "
|
||||
f"{sorted(_ACCEPTED_MIMES)}")
|
||||
data_b64 = payload.get("data")
|
||||
if not isinstance(data_b64, str):
|
||||
raise ValueError("session init image data must be a base64 string")
|
||||
try:
|
||||
@@ -117,14 +87,6 @@ def persist_session_init_image(
|
||||
return SessionInitImage(path=path, display_name=display_name, mime=mime)
|
||||
|
||||
|
||||
def cleanup_session_init_image(session_image: SessionInitImage | None) -> None:
|
||||
"""Remove a persisted session init image and its temporary directory."""
|
||||
|
||||
if session_image is None:
|
||||
return
|
||||
shutil.rmtree(session_image.temp_dir, ignore_errors=True)
|
||||
|
||||
|
||||
def _sanitize_display_name(name: Any) -> str | None:
|
||||
if not isinstance(name, str):
|
||||
return None
|
||||
@@ -137,6 +99,5 @@ def _sanitize_display_name(name: Any) -> str | None:
|
||||
|
||||
__all__ = [
|
||||
"SessionInitImage",
|
||||
"cleanup_session_init_image",
|
||||
"persist_session_init_image",
|
||||
]
|
||||
|
||||
@@ -9,16 +9,12 @@ format; keeping the same shape makes log tooling portable.
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
from datetime import datetime, timezone
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, TextIO
|
||||
|
||||
_FILENAME_SANITIZE_RE = re.compile(r"[^A-Za-z0-9._-]")
|
||||
@@ -111,50 +107,7 @@ class SessionLogger:
|
||||
return handle, lock
|
||||
|
||||
|
||||
def _utc_now_iso() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
class SessionEventLogger:
|
||||
"""Dreamverse-compatible async session event logger.
|
||||
|
||||
This keeps the product server importable during the monorepo migration while
|
||||
the generic streaming logger remains :class:`SessionLogger`.
|
||||
"""
|
||||
|
||||
def __init__(self, root_dir: Path):
|
||||
self.hostname = socket.gethostname()
|
||||
timestamp = datetime.now(timezone.utc).strftime("%y%m%d_%H%M%S")
|
||||
self.directory = root_dir / self.hostname
|
||||
self.path = self.directory / f"{timestamp}.jsonl"
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
self.directory.mkdir(parents=True, exist_ok=True)
|
||||
self.path.touch(exist_ok=False)
|
||||
|
||||
async def write_event(
|
||||
self,
|
||||
*,
|
||||
event: str,
|
||||
client_id: str,
|
||||
payload: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
entry = {
|
||||
"ts": _utc_now_iso(),
|
||||
"event": event,
|
||||
"hostname": self.hostname,
|
||||
"client_id": client_id,
|
||||
}
|
||||
if payload:
|
||||
entry.update(payload)
|
||||
|
||||
async with self._lock:
|
||||
with self.path.open("a", encoding="utf-8") as fp:
|
||||
fp.write(json.dumps(entry, ensure_ascii=False) + "\n")
|
||||
|
||||
|
||||
__all__ = [
|
||||
"SessionLogEvent",
|
||||
"SessionEventLogger",
|
||||
"SessionLogger",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user