144 lines
4.9 KiB
Python
144 lines
4.9 KiB
Python
"""The driven-loop contract — iteration the runtime can see.
|
|
|
|
A loop is a state machine the runner drives:
|
|
|
|
state = loop.init(request, instance, inputs)
|
|
while True:
|
|
plan = loop.next(state) # describe the next step; NO kernels run here
|
|
if isinstance(plan, Done):
|
|
break
|
|
result = plan.run() # the runner executes the thunk
|
|
state = loop.advance(state, result)
|
|
outputs = loop.finalize(state)
|
|
|
|
Two properties this buys, both load-bearing for the product:
|
|
|
|
* every denoise step is a named unit (``plan.label``) — the identity chain
|
|
``request/stage/loop.step`` reaches profilers (NVTX) and the trace without
|
|
any model code knowing about either;
|
|
* all per-request mutable state lives in ``LoopState``, never on the loop
|
|
object — a loop instance is shareable and a session/rollout engine can hold
|
|
many states against one instance.
|
|
|
|
Loop classes declare a ``semantics`` id (e.g. ``"wan.flow_euler.cfg/v1"``).
|
|
Cards pin the semantics their weights assume; ``ModelCard.validate()`` enforces
|
|
the match. Rename-proof, and a distilled student cannot silently run under a
|
|
base sampler.
|
|
|
|
This module is stdlib-only.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Callable, Mapping, Protocol, runtime_checkable
|
|
|
|
|
|
@dataclass
|
|
class LoopState:
|
|
"""All per-request mutable loop state. Nothing else may hold any."""
|
|
loop_id: str
|
|
request_id: str
|
|
step_idx: int = 0
|
|
latents: Any = None
|
|
sigmas: list[float] = field(default_factory=list)
|
|
cond: dict[str, Any] = field(default_factory=dict)
|
|
trajectory: list[Any] = field(default_factory=list) # filled when capture is on
|
|
scratch: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class WorkPlan:
|
|
"""A typed description of the next step. ``run`` is the kernel thunk,
|
|
built by ``next()`` but not called there — planning stays kernel-free."""
|
|
label: str # e.g. "denoise.7" — joins the identity chain
|
|
step: int
|
|
run: Callable[[], dict]
|
|
meta: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
class Done:
|
|
"""Sentinel returned by ``next`` when the loop is finished."""
|
|
|
|
|
|
@runtime_checkable
|
|
class Loop(Protocol):
|
|
"""The model-owned control flow. Four methods; ``next`` is kernel-free."""
|
|
semantics: str
|
|
|
|
def init(self, request: Any, instance: Any, inputs: Mapping[str, Any]) -> LoopState: ...
|
|
def next(self, state: LoopState) -> "WorkPlan | Done": ...
|
|
def advance(self, state: LoopState, result: dict) -> LoopState: ...
|
|
def finalize(self, state: LoopState) -> dict: ...
|
|
|
|
|
|
class LoopRunner:
|
|
"""Drives one loop for one request. ``observe(label, seconds, meta)`` is the
|
|
trace hook; the engine uses it for NVTX ranges and step timings."""
|
|
|
|
def __init__(self, loop: Loop, request: Any, instance: Any, inputs: Mapping[str, Any],
|
|
observe: Callable[[str, float, dict], None] | None = None):
|
|
self.loop = loop
|
|
self.state = loop.init(request, instance, inputs)
|
|
self.observe = observe
|
|
self._done = False
|
|
self._outputs: dict | None = None
|
|
|
|
@property
|
|
def done(self) -> bool:
|
|
return self._done
|
|
|
|
def step(self) -> bool:
|
|
"""Advance one step. Returns True when the loop has just finished."""
|
|
if self._done:
|
|
return True
|
|
plan = self.loop.next(self.state)
|
|
if isinstance(plan, Done):
|
|
self._outputs = self.loop.finalize(self.state)
|
|
self._done = True
|
|
return True
|
|
import time
|
|
nvtx = _nvtx_push(plan.label)
|
|
t0 = time.perf_counter()
|
|
try:
|
|
result = plan.run()
|
|
finally:
|
|
_nvtx_pop(nvtx)
|
|
if self.observe is not None:
|
|
self.observe(plan.label, time.perf_counter() - t0, plan.meta)
|
|
self.state = self.loop.advance(self.state, result)
|
|
return False
|
|
|
|
def run(self) -> dict:
|
|
while not self._done:
|
|
self.step()
|
|
assert self._outputs is not None
|
|
return self._outputs
|
|
|
|
|
|
def _nvtx_push(label: str) -> bool:
|
|
"""Open an NVTX range for one step when CUDA is live. Because the engine
|
|
already opened a ``request/stage`` range, step ranges nest under it — the
|
|
NVTX hierarchy *is* the identity chain, with zero model-code awareness."""
|
|
try:
|
|
import torch
|
|
if torch.cuda.is_available():
|
|
torch.cuda.nvtx.range_push(label)
|
|
return True
|
|
except ImportError:
|
|
pass
|
|
return False
|
|
|
|
|
|
def _nvtx_pop(pushed: bool) -> None:
|
|
if pushed:
|
|
import torch
|
|
torch.cuda.nvtx.range_pop()
|
|
|
|
|
|
def build_loop(spec: Any) -> Loop:
|
|
"""Construct the loop a ``LoopSpec`` declares: resolve the class ref,
|
|
apply the (plain-data) params."""
|
|
from fastvideo2.card import resolve_ref
|
|
cls = resolve_ref(spec.loop)
|
|
return cls(loop_id=spec.loop_id, **spec.params)
|