SolitaryThinker
145082cf74
Directory rename (git mv, history preserved) plus rewrite of all references: absolute imports in tests, the zero-dep runner, docstrings, comments, and the README. No behavior change. Run: python3 -m pytest v2/tests/ -q ; python3 v2/run_tests.py ; python3 -m v2.examples
204 lines
8.3 KiB
Python
204 lines
8.3 KiB
Python
"""The driven-loop contract (design_v3 §5) — the keystone.
|
|
|
|
A loop is a **serializable state machine the runtime drives**:
|
|
|
|
state = loop.init(req, model_state, ctx)
|
|
while True:
|
|
plan = loop.next(state) # describe next step; NO GPU kernels
|
|
if isinstance(plan, Done): break
|
|
result = ctx.execute(plan) # ← THE INVERSION POINT (runtime owns it)
|
|
state = loop.advance(state, result) # fold result in; decide what's next
|
|
for chunk in plan.emits: ctx.emit(chunk)
|
|
return loop.finalize(state)
|
|
|
|
Two properties this contract buys (design_v3 §5.1):
|
|
* content-adaptive steps are natural — ``next`` reads ``state``, which already folded
|
|
in the previous ``StepResult`` via ``advance``;
|
|
* cross-request state safety is *structural* — all per-request mutable state lives in
|
|
``LoopState`` (incl. ``plugin_state`` per request/CFG-branch), never module globals,
|
|
so interleaving requests through one ModelInstance cannot smear state. This is what
|
|
the batch-of-N interleave gate (§9.3) verifies.
|
|
|
|
This module is pure stdlib (tensors typed as ``TensorLike``) so it imports no backend.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Protocol, runtime_checkable
|
|
|
|
from .._enums import ExecutionProfile, WorkUnitKind
|
|
from .._types import TensorLike
|
|
from ..request.streams import StreamChunk
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Shape / resources / cache plan / placement #
|
|
# --------------------------------------------------------------------------- #
|
|
@dataclass(frozen=True)
|
|
class ShapeSignature:
|
|
"""Batch-compatibility + graph-capture key (design_v3 §6.2)."""
|
|
kind: WorkUnitKind
|
|
dims: tuple[int, ...] = ()
|
|
dtype: str = "float32"
|
|
extra: tuple[tuple[str, Any], ...] = () # e.g. (("cfg","classic"),("expert","e0"))
|
|
|
|
@property
|
|
def work_units(self) -> int:
|
|
u = 1
|
|
for d in self.dims:
|
|
u *= max(int(d), 1)
|
|
return u
|
|
|
|
@property
|
|
def batch_key(self) -> tuple:
|
|
"""Two WorkPlans batch together iff their batch_keys are equal (§6.3 BatchScheduler)."""
|
|
return (self.kind, self.dims, self.dtype, self.extra)
|
|
|
|
|
|
@dataclass
|
|
class ResourceRequest:
|
|
"""Everything admission must reserve (design_v3 §6.2 admission rule)."""
|
|
compute_seconds: float = 0.0
|
|
resident_bytes: int = 0
|
|
peak_activation_bytes: int = 0
|
|
cache_blocks: dict[str, int] = field(default_factory=dict)
|
|
transfer_bytes: int = 0
|
|
graph_capture_size: int = 0
|
|
|
|
|
|
@dataclass
|
|
class CacheOp:
|
|
cache_class: str
|
|
key: Any # CacheKey
|
|
nbytes: int = 0
|
|
|
|
|
|
@dataclass
|
|
class CachePlan:
|
|
reads: list[CacheOp] = field(default_factory=list)
|
|
writes: list[CacheOp] = field(default_factory=list)
|
|
|
|
|
|
@dataclass
|
|
class PlacementHint:
|
|
pool: str = "default"
|
|
role: str = "denoise"
|
|
device: str = "cpu"
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# WorkPlan / Done / StepResult / LoopResult #
|
|
# --------------------------------------------------------------------------- #
|
|
@dataclass
|
|
class WorkPlan:
|
|
"""A typed description of the next step (design_v3 §6.2).
|
|
|
|
``run`` is the kernel thunk built by ``next()`` but NOT called there (kernel-free
|
|
planning). The runtime's ``ctx.execute`` calls it — possibly after batching with
|
|
other compatible plans — preserving 'the runtime owns iteration'.
|
|
"""
|
|
loop_id: str
|
|
instance_id: str
|
|
kind: WorkUnitKind
|
|
shape_sig: ShapeSignature
|
|
resources: ResourceRequest = field(default_factory=ResourceRequest)
|
|
cache: CachePlan = field(default_factory=CachePlan)
|
|
placement: PlacementHint = field(default_factory=PlacementHint)
|
|
emits: list[StreamChunk] = field(default_factory=list)
|
|
payload: dict[str, Any] = field(default_factory=dict) # inspectable inputs (debug/parity)
|
|
# the kernel thunk: run(model_instance, override=None) -> StepResult | dict.
|
|
# Built by next() but NOT called there (kernel-free planning). ``override`` is an optional
|
|
# interceptor-supplied forward result (e.g. a cached prediction); the step body still runs
|
|
# the cheap solver step with it.
|
|
run: Any = None
|
|
label: str = ""
|
|
|
|
|
|
@dataclass
|
|
class Done:
|
|
"""Sentinel returned by ``next`` when the loop is finished. ``finalize`` produces the
|
|
actual LoopResult; ``result`` here is optional (kept for loops that want to carry it)."""
|
|
result: "LoopResult | None" = None
|
|
|
|
|
|
@dataclass
|
|
class StepResult:
|
|
output: dict[str, Any] = field(default_factory=dict) # typed per-loop (e.g. {"noise_pred": ...})
|
|
actual_seconds: float = 0.0
|
|
cache_writes: list[CacheOp] = field(default_factory=list)
|
|
behavior: Any = None # BehaviorRecord slice (rollout profile)
|
|
|
|
|
|
@dataclass
|
|
class LoopResult:
|
|
outputs: dict[str, Any] = field(default_factory=dict) # final latents / tokens / artifacts
|
|
metrics: dict[str, float] = field(default_factory=dict)
|
|
behavior: Any = None # full BehaviorRecord (rollout)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# LoopState — the per-request mutable container (NEVER module globals) #
|
|
# --------------------------------------------------------------------------- #
|
|
@dataclass
|
|
class LoopState:
|
|
"""All per-request mutable state lives here (design_v3 §5.1 structural safety)."""
|
|
loop_id: str
|
|
instance_id: str
|
|
request_id: str
|
|
profile: ExecutionProfile = ExecutionProfile.SERVE
|
|
step_idx: int = 0
|
|
done: bool = False
|
|
rng: Any = None # seeded numpy Generator (per request)
|
|
seed: int | None = None
|
|
# common typed fields (resolved at init)
|
|
latents: dict[str, TensorLike] = field(default_factory=dict)
|
|
cond: dict[str, Any] = field(default_factory=dict) # conditioning written by ConditioningInjector
|
|
timesteps: list[float] = field(default_factory=list)
|
|
sigmas: list[float] = field(default_factory=list)
|
|
# per-model extension (Cosmos3PackedSeq, MatrixGameState, ...) — typed by LoopSpec.extension_schema
|
|
extension: Any = None
|
|
# interceptor/policy state, keyed per plugin id AND per CFG branch (design_v3 §5.1, §11)
|
|
plugin_state: dict[str, Any] = field(default_factory=dict)
|
|
cache_handles: dict[str, Any] = field(default_factory=dict)
|
|
# capture buffers (rollout): trajectory of per-step records
|
|
trajectory: list[Any] = field(default_factory=list)
|
|
scratch: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# StepContext — what policies read each step (design_v3 §6.2.3 rule 2) #
|
|
# --------------------------------------------------------------------------- #
|
|
@dataclass
|
|
class StepContext:
|
|
step_idx: int
|
|
timestep: float
|
|
sigma: float
|
|
branch: str = "cond" # current guidance branch
|
|
active_expert_id: str | None = None # set by ExpertRouting; observed by AdaptiveGateCFG
|
|
sampler_coeffs: dict[str, float] = field(default_factory=dict)
|
|
extra: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Protocols: LoopContext (the runtime seam) and Loop (the model contract) #
|
|
# --------------------------------------------------------------------------- #
|
|
@runtime_checkable
|
|
class LoopContext(Protocol):
|
|
"""The single seam the runtime exposes to a loop. The model never sees the
|
|
scheduler; the scheduler never sees the model's math."""
|
|
profile: ExecutionProfile
|
|
|
|
def execute(self, plan: WorkPlan) -> StepResult: ... # THE INVERSION POINT
|
|
def emit(self, chunk: StreamChunk) -> None: ...
|
|
def check_cancel(self) -> None: ... # raises request.Cancelled at boundary
|
|
def observe(self, event: str, **kw) -> None: ... # observer bus hook
|
|
|
|
|
|
@runtime_checkable
|
|
class Loop(Protocol):
|
|
"""The model-owned control flow (design_v3 §5.1). Four methods; ``next`` is kernel-free."""
|
|
def init(self, req: Any, model: Any, ctx: LoopContext) -> LoopState: ...
|
|
def next(self, state: LoopState) -> "WorkPlan | Done": ...
|
|
def advance(self, state: LoopState, result: StepResult) -> LoopState: ...
|
|
def finalize(self, state: LoopState) -> LoopResult: ...
|