90 lines
2.8 KiB
Python
90 lines
2.8 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
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from fastvideo.api.schema import (
|
|
ContinuationState,
|
|
GenerationRequest,
|
|
WarmupConfig,
|
|
)
|
|
from fastvideo.entrypoints.streaming.worker import (
|
|
_extract_continuation_state,
|
|
_warmup_worker,
|
|
)
|
|
|
|
|
|
@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
|