8530fdc550
Co-authored-by: Will Lin <willlin@nvidia.com> Co-authored-by: William Lin <8941107+SolitaryThinker@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
122 lines
3.9 KiB
Python
122 lines
3.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Unit coverage for :mod:`fastvideo.entrypoints.streaming.worker`.
|
|
|
|
The full ``worker_main`` loop runs in a subprocess and is exercised via
|
|
``test_gpu_pool.py``'s subprocess integration tests. This file covers
|
|
the in-process pieces:
|
|
|
|
* the two-segment warmup feeds segment 1's continuation state into
|
|
segment 2 so both compile branches are primed before the worker
|
|
reports ready
|
|
* result-shape extractors handle both attribute-style and dict-style
|
|
generator returns
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import queue
|
|
import threading
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
import fastvideo.envs as envs
|
|
from fastvideo.api.schema import (
|
|
ContinuationState,
|
|
GeneratorConfig,
|
|
GenerationRequest,
|
|
WarmupConfig,
|
|
)
|
|
from fastvideo.entrypoints.streaming.worker import (
|
|
_extract_continuation_state,
|
|
_warmup_worker,
|
|
worker_main,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class _RecordingGenerator:
|
|
"""Captures the requests passed to ``generate`` and returns a
|
|
canned :class:`ContinuationState` on the first call so the warmup
|
|
can feed it into the second call.
|
|
"""
|
|
|
|
state_to_return: ContinuationState | None = field(default_factory=lambda: ContinuationState(
|
|
kind="ltx2.v1",
|
|
payload={
|
|
"schema_version": 1,
|
|
"segment_index": 1
|
|
},
|
|
))
|
|
requests: list[GenerationRequest] = field(default_factory=list)
|
|
|
|
def generate(self, request: GenerationRequest) -> dict[str, Any]:
|
|
self.requests.append(request)
|
|
return {"frames": [], "state": self.state_to_return}
|
|
|
|
|
|
class TestWarmupTwoSegment:
|
|
|
|
def test_warmup_runs_segment_one_then_segment_two_with_returned_state(self) -> None:
|
|
gen = _RecordingGenerator()
|
|
_warmup_worker(gen, WarmupConfig(enabled=True, prompt="warm"))
|
|
|
|
assert len(gen.requests) == 2
|
|
|
|
seg1 = gen.requests[0]
|
|
assert seg1.state is None
|
|
assert seg1.output.return_state is True
|
|
|
|
seg2 = gen.requests[1]
|
|
assert seg2.state is gen.state_to_return
|
|
|
|
def test_warmup_passes_through_when_no_state_returned(self) -> None:
|
|
gen = _RecordingGenerator(state_to_return=None)
|
|
_warmup_worker(gen, WarmupConfig(enabled=True, prompt="warm"))
|
|
|
|
assert len(gen.requests) == 2
|
|
assert gen.requests[1].state is None
|
|
|
|
|
|
class TestExtractContinuationState:
|
|
|
|
def test_extracts_from_attribute(self) -> None:
|
|
|
|
class _R:
|
|
state = ContinuationState(kind="k", payload={})
|
|
|
|
assert _extract_continuation_state(_R()).kind == "k"
|
|
|
|
def test_extracts_from_dict(self) -> None:
|
|
state = ContinuationState(kind="k", payload={})
|
|
assert _extract_continuation_state({"state": state}) is state
|
|
|
|
def test_returns_none_for_missing(self) -> None:
|
|
assert _extract_continuation_state({}) is None
|
|
assert _extract_continuation_state(object()) is None
|
|
|
|
|
|
def test_worker_rejects_external_launcher_before_generator_construction(env_overrides) -> None:
|
|
# worker_main() pins CUDA_VISIBLE_DEVICES for the worker process it normally
|
|
# runs in. Here it runs in the pytest process, so restore the test
|
|
# process's value afterwards.
|
|
env_overrides.enter_context(envs.override_external("CUDA_VISIBLE_DEVICES",
|
|
os.environ.get("CUDA_VISIBLE_DEVICES")))
|
|
result_queue: queue.Queue = queue.Queue()
|
|
config = GeneratorConfig(model_path="/models/fake")
|
|
config.engine.execution_backend = "external_launcher"
|
|
|
|
worker_main(
|
|
gpu_id=0,
|
|
worker_id="worker-0",
|
|
generator_config=config,
|
|
warmup_config=WarmupConfig(enabled=False),
|
|
job_queue=queue.Queue(),
|
|
result_queue=result_queue,
|
|
shutdown_event=threading.Event(),
|
|
)
|
|
|
|
result = result_queue.get_nowait()
|
|
assert result["kind"] == "error"
|
|
assert "external_launcher is supported only" in result["error"]
|
|
assert "streaming worker" in result["error"]
|