Skip to content

[feat]: decouple Dreamverse fMP4 streaming from generation #1425

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Davids048 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[feat]: decouple Dreamverse fMP4 streaming from generation #1425

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b4be1de
[feat]: decouple Dreamverse fMP4 streaming from generation
Davids048 Jun 2, 2026
53b1949
[feat]: cover async Dreamverse media relay
Davids048 Jun 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
142 changes: 113 additions & 29 deletions apps/dreamverse/dreamverse/gpu_pool.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import multiprocessing as mp
import os
import subprocess
import threading
import time
import traceback
from dataclasses import dataclass
Expand Down Expand Up @@ -35,6 +36,7 @@
LeaveAck,
MediaChunk,
MediaComplete,
MediaError,
MediaInit,
ReloadAck,
ReloadModelPayload,
Expand All @@ -47,6 +49,8 @@
WorkerEvent,
)

ACTIVE_STREAM_JOIN_TIMEOUT_SECONDS = 15.0


def _parse_requested_gpu_limit() -> int | None:
raw_value = os.getenv("FASTVIDEO_GPU_COUNT", "").strip().lower()
Expand Down Expand Up @@ -163,11 +167,99 @@ def gpu_worker_process(
from dreamverse.video_generation import VideoGenerationWorker

worker = VideoGenerationWorker(gpu_id)
# One active fMP4 stream per worker keeps binary media chunks ordered by segment.
active_stream_thread: threading.Thread | None = None
active_stream_segment_idx: int | None = None

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

def wait_for_active_stream() -> None:
"""Join the active fMP4 segment stream before leave, shutdown, or the next stream."""
nonlocal active_stream_thread
nonlocal active_stream_segment_idx
if active_stream_thread is None:
return
if active_stream_thread.is_alive():
print(f"[GPU {gpu_id}] Waiting for AV stream "
f"segment {active_stream_segment_idx} before continuing")
active_stream_thread.join(timeout=ACTIVE_STREAM_JOIN_TIMEOUT_SECONDS)
if active_stream_thread.is_alive():
raise RuntimeError(f"Timed out waiting for AV stream segment "
f"{active_stream_segment_idx} after "
f"{ACTIVE_STREAM_JOIN_TIMEOUT_SECONDS:.0f}s")
active_stream_thread = None
active_stream_segment_idx = None
Comment on lines +178 to +193

@gemini-code-assist gemini-code-assist Bot Jun 2, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Calling active_stream_thread.join() without a timeout can block the GPU worker event loop indefinitely if the background stream thread hangs (e.g., due to an ffmpeg pipe deadlock or process hang). If this happens during a USER_LEAVE command, the event loop will remain permanently blocked, rendering the GPU slot completely unusable for subsequent sessions. Adding a timeout to the join operation prevents permanent hangs.

        def wait_for_active_stream(timeout: float = 30.0) -> None:
            """Join the active fMP4 segment stream before leave, shutdown, or the next stream."""
            nonlocal active_stream_thread
            nonlocal active_stream_segment_idx
            if active_stream_thread is None:
                return
            if active_stream_thread.is_alive():
                print(f"[GPU {gpu_id}] Waiting for AV stream "
                      f"segment {active_stream_segment_idx} before continuing")
            active_stream_thread.join(timeout=timeout)
            if active_stream_thread.is_alive():
                print(f"[GPU {gpu_id}] WARNING: AV stream thread for segment "
                      f"{active_stream_segment_idx} did not finish within {timeout}s")
            else:
                active_stream_thread = None
                active_stream_segment_idx = None


def start_stream_thread(
*,
step_result,
user_id: str,
segment_idx: int,
stream_id: str,
head_trim_frames: int,
head_trim_audio_frames: int,
) -> None:
"""Start fMP4 encoding for one segment and publish media events to the controller."""
nonlocal active_stream_thread
nonlocal active_stream_segment_idx
wait_for_active_stream()
stream_timings = dict(step_result.timings)

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

def _run_stream() -> None:
try:
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=stream_timings,
head_trim_frames=head_trim_frames,
head_trim_audio_frames=head_trim_audio_frames,
shared_buffer=None,
shared_buffer_bytes=0,
publish=_publish,
log_prefix=f"[GPU {gpu_id}]",
)
if not av_ok:
response_queue.put(
MediaError(
user_id=user_id,
segment_idx=segment_idx,
stream_id=stream_id,
message=av_error or "worker av_fmp4 stream failed",
))
return
print(f"[GPU {gpu_id}] AV streamed segment {segment_idx}: "
f"encode_total={stream_timings.get('av_encode_stream_ms', 0):.0f}ms "
f"wav_write={stream_timings.get('av_wav_write_ms', 0):.1f}ms "
f"spawn={stream_timings.get('av_ffmpeg_spawn_ms', 0):.1f}ms "
f"first_chunk={stream_timings.get('av_first_chunk_ms', 0):.0f}ms "
f"chunk_interval_med={stream_timings.get('av_chunk_interval_ms_median', 0):.1f}ms "
f"chunk_interval_p95={stream_timings.get('av_chunk_interval_ms_p95', 0):.1f}ms "
f"publish_med={stream_timings.get('av_chunk_publish_ms_median', 0):.2f}ms "
f"read_med={stream_timings.get('av_chunk_read_ms_median', 0):.1f}ms")
except Exception as exc:
response_queue.put(
MediaError(
user_id=user_id,
segment_idx=segment_idx,
stream_id=stream_id,
message=str(exc),
))

active_stream_thread = threading.Thread(
target=_run_stream,
name=f"dreamverse-av-stream-gpu{gpu_id}-seg{segment_idx}",
daemon=False,
)
active_stream_segment_idx = segment_idx
active_stream_thread.start()

def handle_command(cmd: Command):
if cmd.type == CommandType.USER_JOIN:
print(f"[GPU {gpu_id}] User {cmd.user_id[:8]} joined")
Expand Down Expand Up @@ -200,40 +292,21 @@ def handle_command(cmd: Command):
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,
start_stream_thread(
step_result=step_result,
user_id=cmd.user_id,
segment_idx=segment_idx,
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")
print(f"[GPU {gpu_id}] AV streamed segment {segment_idx}: "
f"encode_total={step_result.timings.get('av_encode_stream_ms', 0):.0f}ms "
f"wav_write={step_result.timings.get('av_wav_write_ms', 0):.1f}ms "
f"spawn={step_result.timings.get('av_ffmpeg_spawn_ms', 0):.1f}ms "
f"first_chunk={step_result.timings.get('av_first_chunk_ms', 0):.0f}ms "
f"chunk_interval_med={step_result.timings.get('av_chunk_interval_ms_median', 0):.1f}ms "
f"chunk_interval_p95={step_result.timings.get('av_chunk_interval_ms_p95', 0):.1f}ms "
f"publish_med={step_result.timings.get('av_chunk_publish_ms_median', 0):.2f}ms "
f"read_med={step_result.timings.get('av_chunk_read_ms_median', 0):.1f}ms")
step_result.timings["ipc_put_start_ns"] = time.time_ns()
step_timings = dict(step_result.timings)
step_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,
timings=step_timings,
))
except Exception as e:
print(f"[GPU {gpu_id}] Step error: {e}")
Expand All @@ -245,8 +318,17 @@ def _publish(event: StreamEvent) -> None:

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))
try:
wait_for_active_stream()
worker.clear_conditioning()
response_queue.put(LeaveAck(user_id=cmd.user_id))
except Exception as e:
print(f"[GPU {gpu_id}] Leave error: {e}")
response_queue.put(WorkerError(
user_id=cmd.user_id,
message=str(e),
))
return False

elif cmd.type == CommandType.SHUTDOWN:
return False # Signal to exit
Expand Down Expand Up @@ -286,6 +368,7 @@ def _publish(event: StreamEvent) -> None:

if first_cmd is not None:
if first_cmd.type == CommandType.SHUTDOWN:
wait_for_active_stream()
worker.shutdown()
response_queue.put(ShutdownAck())
return
Expand All @@ -300,6 +383,7 @@ def _publish(event: StreamEvent) -> None:

if cmd.type == CommandType.SHUTDOWN:
print(f"[GPU {gpu_id}] Event loop shutting down")
wait_for_active_stream()
worker.shutdown()
response_queue.put(ShutdownAck())
return
Expand Down Expand Up @@ -600,7 +684,7 @@ def get_response_nonblocking():
continue

# AV streaming events → route to the user's stream queue.
if isinstance(event, (MediaInit, MediaChunk, MediaComplete)):
if isinstance(event, (MediaInit, MediaChunk, MediaComplete, MediaError)):
stream_queue = self._stream_queues.get(event.user_id)
if stream_queue is not None:
await stream_queue.put(event)
Expand Down
Loading
Loading