SolitaryThinker
290795daf8
Second step of the runtime simplification (after cost removal). Drop the coordinated step-interleave scheduler + the interleave-parity gate; serving is now pooled run-to-completion. - Engine: remove run_interleaved + the WorkUnit/BatchScheduler imports; run / run_serial drive each request to completion (tick/run_to_completion kept as the per-request stepper). - scheduler.py: remove BatchScheduler + WorkUnit + the batches metric; the AdmissionController is now a pure refundable memory/OOM guard. - AsyncEngine: bound concurrency with a serving pool (asyncio.Semaphore, max_concurrent) — each request waits for a slot, then runs to completion. - parity: remove assert_interleave_parity (the run_serial==run_interleaved gate); rename interleave_gate.py -> compare.py (compare_outputs stays — bit-parity between execution paths, e.g. disaggregated==inline). - card specs: drop ParitySpec.interleave_required + LoopSpec.allows_interleaving; stripped interleave_required from all cards. - tests: delete the interleave-gate/parity tests; refocus the ones that exercised real behavior (residual-skip, compare_outputs symmetric-empty). - README: removed the design-doc references at the top; simplified the thesis / scheduler (§6) / parity (§9) / package-layout / comparison sections to pooled run-to-completion (no cost model, no interleave gate). CPU mini: 143 passed / 2 skipped. Native omni port (P3) still to come. (pyproject kernel hack excluded.)
261 lines
12 KiB
Python
261 lines
12 KiB
Python
"""The engine — drives Programs over resident ModelInstances.
|
|
|
|
Pooled run-to-completion: each request's program runs start-to-finish via one ProgramRunner
|
|
(``run`` for a single request, ``run_serial`` for a batch). Concurrency is bounded by the pool at
|
|
the serving layer (AsyncEngine), not by a step-interleaved scheduler.
|
|
|
|
The engine owns admission (a refundable memory/OOM reservation before stepping), the observer bus,
|
|
and the interceptor chain (deploy scope). It never imports ``training`` (dependency rule).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
from v2._enums import ExecutionProfile
|
|
from v2.extend.base import InterceptorChain, ObserverBus
|
|
from v2.loop.driver import LoopRunner
|
|
from v2.parity.aligner import ParityAligner # noqa: F401 (engine can attach it as an observer)
|
|
from v2.program.specs import ComponentNode, ModelLoopNode, Program
|
|
from v2.request.artifacts import (
|
|
AudioArtifact,
|
|
LatentArtifact,
|
|
Output,
|
|
TensorArtifact,
|
|
TextArtifact,
|
|
VideoArtifact,
|
|
)
|
|
from v2.request.cancel import CancelKind, CancelScope
|
|
from v2.request.streams import Stream
|
|
from v2.runtime.context import RuntimeLoopContext
|
|
from v2.runtime.scheduler import (
|
|
AdmissionController,
|
|
AdmissionInfeasible,
|
|
SchedulerMetrics,
|
|
)
|
|
|
|
|
|
def _host(v: Any) -> Any:
|
|
"""The single request-boundary marshal: a device (torch) tensor -> host numpy; numpy passes through.
|
|
On-device loops keep tensors resident the whole run and land here exactly once, at output."""
|
|
if not isinstance(v, np.ndarray) and hasattr(v, "detach") and hasattr(v, "cpu"):
|
|
return v.detach().float().cpu().numpy()
|
|
return v
|
|
|
|
|
|
def _to_artifact(name: str, value: Any, producer: str = "") -> Any:
|
|
if isinstance(value, dict):
|
|
value = value.get(name, value.get("latents", next(iter(value.values()), None)))
|
|
value = _host(value)
|
|
if name == "video":
|
|
return VideoArtifact(name=name, producer=producer, frames=value)
|
|
if name == "audio":
|
|
return AudioArtifact(name=name, producer=producer, samples=value)
|
|
if name == "text":
|
|
return TextArtifact(name=name, producer=producer, text=value if isinstance(value, str) else "")
|
|
if name == "latents":
|
|
return LatentArtifact(name=name, producer=producer, latent=value)
|
|
return TensorArtifact(name=name, producer=producer, tensor=value)
|
|
|
|
|
|
class ProgramRunner:
|
|
"""Drives one Program for one request to completion (``tick`` advances one unit at a time)."""
|
|
|
|
def __init__(self, engine: Engine, program: Program, instance: Any, request: Any):
|
|
self.engine = engine
|
|
self.program = program
|
|
self.instance = instance
|
|
self.request = request
|
|
self.slots: dict[str, Any] = {}
|
|
self.metrics: dict[str, float] = {}
|
|
self.cancel_scope = CancelScope(kind=CancelKind.REQUEST, target_id=request.request_id)
|
|
self.stream = Stream(request.request_id)
|
|
profile = ExecutionProfile.ROLLOUT if getattr(request.outputs, "capture", None) and \
|
|
request.outputs.capture.value == "behavior" else engine.profile
|
|
self.ctx = RuntimeLoopContext(instance,
|
|
observers=engine.observers,
|
|
interceptors=engine.interceptors,
|
|
slots=self.slots,
|
|
stream=self.stream,
|
|
cancel_scope=self.cancel_scope,
|
|
profile=profile,
|
|
metrics=self.metrics,
|
|
request_id=request.request_id)
|
|
self.node_idx = 0
|
|
self.loop_runner: LoopRunner | None = None
|
|
self._resident_res: Any = None
|
|
self.done = False
|
|
self.error: str | None = None
|
|
|
|
def _merge_metrics(self, m: dict[str, float]) -> None:
|
|
for k, v in m.items():
|
|
self.metrics[k] = self.metrics.get(k, 0.0) + v
|
|
|
|
def tick(self) -> bool:
|
|
if self.done:
|
|
return True
|
|
self.cancel_scope.check() # common-path cancellation (offline + all-component programs)
|
|
nodes = self.program.nodes
|
|
while True:
|
|
if self.node_idx >= len(nodes):
|
|
self.done = True
|
|
return True
|
|
node = nodes[self.node_idx]
|
|
if not node.when(self.request):
|
|
self.node_idx += 1
|
|
continue
|
|
|
|
if isinstance(node, ComponentNode):
|
|
if node.fn is not None:
|
|
node.fn(self.instance, self.slots, self.request, self.ctx)
|
|
self.node_idx += 1
|
|
continue # one-shots run within this tick
|
|
|
|
if isinstance(node, ModelLoopNode):
|
|
if self.loop_runner is None:
|
|
loop = self.instance.loop(node.loop_id)
|
|
self.loop_runner = LoopRunner(loop, self.ctx, self.request, self.instance)
|
|
self._resident_res = None
|
|
|
|
# reserve resident memory for the whole loop (concurrent-pressure admission)
|
|
if self._resident_res is None:
|
|
plan0 = self.loop_runner.peek()
|
|
if plan0 is None:
|
|
self._finish_loop(node)
|
|
return False
|
|
need = plan0.resources.resident_bytes
|
|
if need > 0:
|
|
if not self.engine.admission.feasible_resident(need):
|
|
raise AdmissionInfeasible(f"{node.loop_id}: resident {need}B exceeds pool capacity "
|
|
f"{self.engine.admission.memory.total_bytes}B")
|
|
res = self.engine.admission.reserve_resident(node.loop_id, need)
|
|
if res is None:
|
|
self.engine.metrics.deferred += 1
|
|
return False # defer: cannot fit resident yet (jointly-OOM guard)
|
|
self._resident_res = res
|
|
else:
|
|
self._resident_res = True
|
|
|
|
plan = self.loop_runner.peek()
|
|
if plan is None:
|
|
self._finish_loop(node)
|
|
return False
|
|
if not self.engine.admission.feasible_step(plan):
|
|
raise AdmissionInfeasible(
|
|
f"{node.loop_id}: step needs peak {plan.resources.peak_activation_bytes}B "
|
|
f"— exceeds pool memory capacity")
|
|
step_res = self.engine.admission.admit_step(plan)
|
|
if step_res is None:
|
|
return False # deferred this tick (will retry; refundable)
|
|
self.loop_runner.step()
|
|
self.engine.admission.release(step_res)
|
|
self.engine.metrics.stepped_units += 1
|
|
if self.loop_runner.done:
|
|
self._finish_loop(node)
|
|
return False # advanced one loop step; run_to_completion loops until done
|
|
# unknown node kind
|
|
self.node_idx += 1
|
|
|
|
def _finish_loop(self, node: ModelLoopNode) -> None:
|
|
if self.loop_runner is not None and self.loop_runner.result is not None:
|
|
self.slots[node.output_slot] = self.loop_runner.result.outputs
|
|
self._merge_metrics(self.loop_runner.result.metrics)
|
|
if self._resident_res not in (None, True):
|
|
self.engine.admission.memory.release(self._resident_res) # resident is a plain Reservation
|
|
self._resident_res = None
|
|
self.loop_runner = None
|
|
self.node_idx += 1
|
|
|
|
@property
|
|
def _progress(self) -> tuple:
|
|
"""A token that changes iff the runner advanced (node, loop step, or resident reservation)."""
|
|
sidx = self.loop_runner.state.step_idx if self.loop_runner is not None else -1
|
|
return (self.node_idx, sidx, self._resident_res is not None)
|
|
|
|
def run_to_completion(self) -> None:
|
|
stuck = 0
|
|
while not self.done:
|
|
p0 = self._progress
|
|
self.tick()
|
|
if not self.done and self._progress == p0:
|
|
stuck += 1
|
|
if stuck > 3: # fail fast, not a 10M busy-spin
|
|
raise AdmissionInfeasible(f"program {self.program.program_id} made no progress for several ticks "
|
|
"(unsatisfiable reservation or deadlock)")
|
|
else:
|
|
stuck = 0
|
|
|
|
def output(self) -> Output:
|
|
artifacts: dict[str, Any] = {}
|
|
for name, slot in self.program.output_artifacts.items():
|
|
artifacts[name] = _to_artifact(name, self.slots.get(slot), producer=slot)
|
|
return Output(request_id=self.request.request_id,
|
|
artifacts=artifacts,
|
|
metrics=dict(self.metrics),
|
|
error=self.error)
|
|
|
|
|
|
@dataclass
|
|
class Engine:
|
|
"""A single-pool engine. Pools are single-node; fleet scale is multiple pools."""
|
|
profile: ExecutionProfile = ExecutionProfile.SERVE
|
|
observers: ObserverBus = field(default_factory=ObserverBus)
|
|
interceptors: InterceptorChain = field(default_factory=InterceptorChain)
|
|
admission: AdmissionController = field(default_factory=AdmissionController)
|
|
metrics: SchedulerMetrics = field(default_factory=SchedulerMetrics)
|
|
_registry: dict[str, tuple[Any, Program]] = field(default_factory=dict)
|
|
_workflows: dict[str, Any] = field(default_factory=dict) # workflow_id -> validated Workflow
|
|
_wf_running: set = field(default_factory=set) # workflow ids in-flight (nested cycle guard)
|
|
|
|
def register(self, model_id: str, instance: Any, program: Program) -> None:
|
|
self._registry[model_id] = (instance, program)
|
|
|
|
def register_workflow(self, workflow: Any) -> Any:
|
|
"""Register a cross-model Workflow as a servable, addressable by ``workflow.workflow_id`` in
|
|
the SAME namespace as model ids. Validates its required cards are already registered."""
|
|
workflow.validate(self)
|
|
if workflow.workflow_id in self._registry:
|
|
raise ValueError(f"workflow id {workflow.workflow_id!r} collides with a registered model")
|
|
self._workflows[workflow.workflow_id] = workflow
|
|
return workflow
|
|
|
|
def serves(self, name: str) -> bool:
|
|
return name in self._registry or name in self._workflows
|
|
|
|
def _resolve(self, request: Any) -> tuple[Any, Program]:
|
|
if request.model_id not in self._registry:
|
|
raise KeyError(f"no model registered for {request.model_id!r} (have {list(self._registry)})")
|
|
return self._registry[request.model_id]
|
|
|
|
def _make_runner(self, request: Any) -> ProgramRunner:
|
|
instance, program = self._resolve(request)
|
|
return ProgramRunner(self, program, instance, request)
|
|
|
|
# --- offline single-request (the VideoGenerator path) -------------------- #
|
|
def run(self, request: Any) -> Output:
|
|
if request.model_id in self._workflows: # a workflow servable → orchestrate stages
|
|
wf = self._workflows[request.model_id]
|
|
if request.model_id in self._wf_running: # nested workflows: cycle guard
|
|
raise ValueError(f"workflow cycle: {request.model_id!r} is already running "
|
|
f"(a workflow cannot transitively invoke itself)")
|
|
seed = request.diffusion.seed if request.diffusion.seed is not None else 0
|
|
self._wf_running.add(request.model_id)
|
|
try:
|
|
return wf.run(self, prompt=request.prompt(), seed=seed, request=request)
|
|
finally:
|
|
self._wf_running.discard(request.model_id)
|
|
r = self._make_runner(request)
|
|
r.run_to_completion()
|
|
return r.output()
|
|
|
|
# --- offline batch: each request runs to completion ---------------------- #
|
|
def run_serial(self, requests: list[Any]) -> dict[str, Output]:
|
|
out: dict[str, Output] = {}
|
|
for req in requests:
|
|
r = self._make_runner(req)
|
|
r.run_to_completion()
|
|
out[req.request_id] = r.output()
|
|
return out
|