134 lines
3.9 KiB
Python
134 lines
3.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Per-GPU worker subprocess entry for :class:`SubprocessGpuPool`.
|
|
|
|
The pool manages binding, lifecycle, and message dispatch in the parent
|
|
process. The worker constructs its :class:`VideoGenerator` from a typed
|
|
:class:`GeneratorConfig`, runs the two-segment warmup so both
|
|
initial-segment and continuation-branch compile graphs are hot, and
|
|
then loops on the job queue.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import multiprocessing as mp
|
|
import queue
|
|
from typing import Any
|
|
|
|
from fastvideo.api.schema import (
|
|
GeneratorConfig,
|
|
GenerationRequest,
|
|
InputConfig,
|
|
OutputConfig,
|
|
SamplingConfig,
|
|
WarmupConfig,
|
|
)
|
|
from fastvideo.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
# Synthetic warmup dimensions: small enough to keep boot fast, big enough
|
|
# to exercise the real shape-dependent compile paths. Keep in sync with
|
|
# WarmupConfig if these become user-tunable.
|
|
_WARMUP_NUM_FRAMES = 8
|
|
_WARMUP_HEIGHT = 256
|
|
_WARMUP_WIDTH = 256
|
|
_WARMUP_NUM_INFERENCE_STEPS = 1
|
|
|
|
|
|
def worker_main(
|
|
*,
|
|
gpu_id: int,
|
|
worker_id: str,
|
|
generator_config: GeneratorConfig,
|
|
warmup_config: WarmupConfig,
|
|
job_queue: mp.Queue,
|
|
result_queue: mp.Queue,
|
|
shutdown_event: Any,
|
|
) -> None: # pragma: no cover - exercised via integration only
|
|
"""Per-worker subprocess entry.
|
|
|
|
Runs inside the child spawned by ``SubprocessGpuPool``. Blocking
|
|
``VideoGenerator`` construction + generation happens here, not in
|
|
the parent's event loop.
|
|
"""
|
|
import fastvideo.envs as envs
|
|
|
|
envs.set_external("CUDA_VISIBLE_DEVICES", str(gpu_id))
|
|
try:
|
|
from fastvideo import VideoGenerator
|
|
|
|
generator = VideoGenerator.from_pretrained(config=generator_config)
|
|
if warmup_config.enabled:
|
|
_warmup_worker(generator, warmup_config)
|
|
result_queue.put({"kind": "ready", "worker_id": worker_id})
|
|
except Exception as exc:
|
|
result_queue.put({"kind": "error", "error": repr(exc)})
|
|
return
|
|
|
|
while not shutdown_event.is_set():
|
|
try:
|
|
item = job_queue.get(timeout=0.5)
|
|
except queue.Empty:
|
|
continue
|
|
if item is None:
|
|
break
|
|
job_id = item["job_id"]
|
|
request = item["request"]
|
|
try:
|
|
result = generator.generate(request)
|
|
result_queue.put({
|
|
"kind": "result",
|
|
"job_id": job_id,
|
|
"result": result,
|
|
})
|
|
except Exception as exc:
|
|
result_queue.put({
|
|
"kind": "error",
|
|
"job_id": job_id,
|
|
"error": repr(exc),
|
|
})
|
|
|
|
|
|
def _warmup_worker(
|
|
generator: Any,
|
|
warmup_config: WarmupConfig,
|
|
) -> None:
|
|
"""Run two synthetic generations so both compile branches are primed.
|
|
|
|
Segment 1 is a fresh start (no continuation state) and exercises
|
|
the initial-segment graph. Segment 2 feeds segment 1's continuation
|
|
state back in so the conditioning branch is also compiled before
|
|
the first user request lands.
|
|
"""
|
|
sampling = SamplingConfig(
|
|
num_frames=_WARMUP_NUM_FRAMES,
|
|
height=_WARMUP_HEIGHT,
|
|
width=_WARMUP_WIDTH,
|
|
num_inference_steps=_WARMUP_NUM_INFERENCE_STEPS,
|
|
)
|
|
seg1 = GenerationRequest(
|
|
prompt=warmup_config.prompt,
|
|
sampling=sampling,
|
|
inputs=InputConfig(),
|
|
output=OutputConfig(save_video=False, return_frames=False, return_state=True),
|
|
)
|
|
seg1_result = generator.generate(seg1)
|
|
|
|
seg2 = GenerationRequest(
|
|
prompt=warmup_config.prompt,
|
|
sampling=sampling,
|
|
inputs=InputConfig(),
|
|
output=OutputConfig(save_video=False, return_frames=False),
|
|
state=_extract_continuation_state(seg1_result),
|
|
)
|
|
generator.generate(seg2)
|
|
|
|
|
|
def _extract_continuation_state(result: Any) -> Any:
|
|
state = getattr(result, "state", None)
|
|
if state is None and isinstance(result, dict):
|
|
state = result.get("state")
|
|
return state
|
|
|
|
|
|
__all__ = ["worker_main"]
|