From 098bcf014a7ffc91087e19f7de06b3554e0e9fc0 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Tue, 16 Jun 2026 04:33:53 -0700 Subject: [PATCH] =?UTF-8?q?[fix]=20mini-fastvideo:=20address=20adversarial?= =?UTF-8?q?-review=20findings=20(admission=20liveness=20+=20=C2=A77.1=20ca?= =?UTF-8?q?che=20key)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Admission fails fast with AdmissionInfeasible on infeasible/deadlocked reservations instead of a 10M-iteration busy-spin: no-progress detection in run_to_completion/run_interleaved via a real progress token, plus feasibility pre-checks (need > pool capacity). - Compute budget is now a refundable concurrency gate (release() refunds spent), not a never-refunded lifetime cap that silently deadlocks. - §7.1: the text-encoder feature CacheKey carries adapter_versions + precision (no stale serve across te-LoRA stacks); per-component weight versions mean a transformer-only RL weight sync no longer flushes the frozen text-encoder cache (component-scoped invalidation, not wholesale). - Interleave gate flags symmetric-empty output instead of passing it vacuously. - skipped_steps counted only when the override is actually consumed; BatchScheduler wired for round batch-accounting (metric renamed stepped_units); ResidualCache.get cleanup; stream chunks carry a latent preview payload; dead progress-vars removed. - 5 regression tests added for the previously-untested paths. 54 tests pass (pytest + zero-dep runner). --- mini_fastvideo/cache/classes.py | 14 ++++- mini_fastvideo/cache/manager.py | 7 +++ mini_fastvideo/card/instance.py | 19 +++++- mini_fastvideo/models/backend.py | 37 +++++++++++ mini_fastvideo/models/common.py | 7 ++- mini_fastvideo/models/omni/ar_loop.py | 76 ++++++++++++++++++++++ mini_fastvideo/models/wan21/loop.py | 3 +- mini_fastvideo/models/wan_causal/loop.py | 2 +- mini_fastvideo/parity/interleave_gate.py | 6 ++ mini_fastvideo/runtime/__init__.py | 11 +++- mini_fastvideo/runtime/context.py | 3 +- mini_fastvideo/runtime/engine.py | 65 ++++++++++++++----- mini_fastvideo/runtime/scheduler.py | 54 ++++++++++------ mini_fastvideo/tests/test_review_fixes.py | 77 +++++++++++++++++++++++ mini_fastvideo/training/weight_sync.py | 6 +- 15 files changed, 339 insertions(+), 48 deletions(-) create mode 100644 mini_fastvideo/models/omni/ar_loop.py create mode 100644 mini_fastvideo/tests/test_review_fixes.py diff --git a/mini_fastvideo/cache/classes.py b/mini_fastvideo/cache/classes.py index 082b24805..4be7521be 100644 --- a/mini_fastvideo/cache/classes.py +++ b/mini_fastvideo/cache/classes.py @@ -79,6 +79,14 @@ class FeatureCache(_Pool): _v, nb, _k = self._store.pop(h) self.used_bytes -= nb + def invalidate_components(self, components: set[str]) -> None: + """Drop only entries produced by the changed components (design_v3 §7.1 partition-not-flush): + a transformer-only weight sync must NOT evict text-encoder embeddings.""" + drop = [h for h, (_v, _nb, k) in self._store.items() if k.component_id in components] + for h in drop: + _v, nb, _k = self._store.pop(h) + self.used_bytes -= nb + class ResidualCache(_Pool): """cache-dit residual store, scoped per ``LoopState`` AND per CFG branch (§5.1, §11). @@ -97,8 +105,10 @@ class ResidualCache(_Pool): def get(self, namespace: str, branch: str, name: str) -> Any | None: v = self._store.get((namespace, branch, name)) - (self.__dict__.__setitem__("hits", self.hits + 1) if v is not None - else self.__dict__.__setitem__("misses", self.misses + 1)) + if v is not None: + self.hits += 1 + else: + self.misses += 1 return v def clear_namespace(self, namespace: str) -> None: diff --git a/mini_fastvideo/cache/manager.py b/mini_fastvideo/cache/manager.py index 644dae636..ac371a4ae 100644 --- a/mini_fastvideo/cache/manager.py +++ b/mini_fastvideo/cache/manager.py @@ -44,6 +44,13 @@ class CacheManager: if hasattr(pool, "invalidate_weights"): pool.invalidate_weights(version) + def invalidate_components(self, components) -> None: + """Component-scoped invalidation (design_v3 §7.1): only drop caches for changed components.""" + comps = set(components) + for pool in self._pools.values(): + if hasattr(pool, "invalidate_components"): + pool.invalidate_components(comps) + def clear_namespace(self, namespace: str) -> None: """Release a finished request's per-request caches (residual/slab).""" for pool in self._pools.values(): diff --git a/mini_fastvideo/card/instance.py b/mini_fastvideo/card/instance.py index ad5d9b3bd..a603e6ddc 100644 --- a/mini_fastvideo/card/instance.py +++ b/mini_fastvideo/card/instance.py @@ -27,6 +27,9 @@ class ModelInstance: self.caches = cache_manager self.weights_version = weights_version self.adapter_versions: dict[str, str] = {} + # per-component weight versions (§7.1): a component's version changes only when IT is synced, + # so a transformer-only RL sync never invalidates the frozen text-encoder's feature cache. + self.component_versions: dict[str, str] = {cid: weights_version for cid in card.components} self._components: dict[str, Any] = {} self._loops: dict[str, Any] = {} self._asleep: set[str] = set() @@ -69,10 +72,20 @@ class ModelInstance: self._asleep.discard(cid) # --- weight sync (design_v3 §10): bump version + invalidate caches -------- # - def set_weights_version(self, version: str) -> None: + def version_of(self, component_id: str) -> str: + """The component's own weights version (defaults to the instance version).""" + return self.component_versions.get(component_id, self.weights_version) + + def set_weights_version(self, version: str, components: list[str] | None = None) -> None: + """Publish a new weights version. If ``components`` is given, only those components' versions + bump and only their caches are invalidated (design_v3 §7.1 partition-not-flush) — so a + transformer-only RL weight sync leaves the frozen text-encoder's feature cache intact.""" self.weights_version = version - if self.caches is not None and hasattr(self.caches, "invalidate_weights"): - self.caches.invalidate_weights(version) + changed = components if components is not None else list(self.card.components.keys()) + for c in changed: + self.component_versions[c] = version + if self.caches is not None and hasattr(self.caches, "invalidate_components"): + self.caches.invalidate_components(set(changed)) def __repr__(self) -> str: return (f"ModelInstance(card={self.card.model_id!r}, weights={self.weights_version!r}, " diff --git a/mini_fastvideo/models/backend.py b/mini_fastvideo/models/backend.py index b8831cc4d..7418155af 100644 --- a/mini_fastvideo/models/backend.py +++ b/mini_fastvideo/models/backend.py @@ -108,6 +108,43 @@ class ToyDiT: return loss, grad_norm +class ToyTokenizer: + """Deterministic byte-ish tokenizer for the omni AR pathway (phase 2).""" + EOS = 0 + VOCAB = 256 + + def encode(self, text: str) -> list[int]: + toks = [(ord(c) % (self.VOCAB - 1)) + 1 for c in (text or "")[:16]] + return toks or [1] + + def decode(self, tokens) -> str: + return "".join(chr(32 + (int(t) % 90)) for t in tokens) + + +class ToyMoTDiT(ToyDiT): + """Mixture-of-Transformers stand-in: ONE resident module that runs BOTH an AR (understanding) + pathway and a diffusion (generation) pathway on shared weights (design_v3 §4.2, §16). + + ``ar_forward`` is the und pathway (next-token); ``__call__`` (inherited) is the gen pathway + (velocity). Binding both the ar_decode and diffusion_denoise loops to one instance of this + component is the MoT requirement no DAG-of-engines can express. + """ + VOCAB = 256 + EOS = 0 + + def ar_forward(self, tokens) -> int: + # deterministic next-token from the context (greedy ⇒ trivially interleave-safe) + ctx = sum(int(t) for t in tokens) + nxt = (ctx * 7 + 13) % self.VOCAB + return int(nxt) + + def reasoner_embed(self, tokens) -> np.ndarray: + """Pack the und-pathway tokens into a conditioning embed the gen pathway consumes + (the Cosmos3 'prompt upsampling before diffusion in the same request').""" + rng = np.random.default_rng((sum(int(t) for t in tokens) + 1) % (1 << 31)) + return (rng.standard_normal((4, 8)) * 0.1).astype("float32") + + class ToyVAE: """Tiny deterministic VAE. encode: video→latent (mean-pool + channel proj); decode: latent→video.""" diff --git a/mini_fastvideo/models/common.py b/mini_fastvideo/models/common.py index 140fe6769..c2515a1f8 100644 --- a/mini_fastvideo/models/common.py +++ b/mini_fastvideo/models/common.py @@ -14,8 +14,13 @@ def cached_text_encode(instance: Any, text: str) -> Any: enc = instance.component("text_encoder") cache = instance.caches if cache is not None and cache.has("feature"): + # §7.1: every output-semantic field is in the key — adapter stack + precision partition it, + # and the text-encoder's OWN version (not the instance's) so a transformer weight sync + # does not invalidate the (frozen) text encoder's embeddings. key = CacheKey(model_id=instance.card.model_id, component_id="text_encoder", - weights_version=instance.weights_version, + weights_version=instance.version_of("text_encoder"), + adapter_versions=CacheKey.adapters(instance.adapter_versions), + precision=instance.card.precision.dtype_for("text_encoder"), input_hashes=(("text", content_hash(text)),)) hit = cache.pool("feature").get(key) if hit is not None: diff --git a/mini_fastvideo/models/omni/ar_loop.py b/mini_fastvideo/models/omni/ar_loop.py new file mode 100644 index 000000000..44d38c2e1 --- /dev/null +++ b/mini_fastvideo/models/omni/ar_loop.py @@ -0,0 +1,76 @@ +"""ARDecodeLoop — token decode until EOS/max_tokens (design_v3 §4.3, §5; §15b Cosmos3 reasoner). + +The autoregressive pathway: ``init`` prefills the prompt tokens (and allocates paged text-KV); +``next`` emits one ``ar_token`` WorkUnit (kernel-free); ``advance`` appends the sampled token and +streams it; ``done`` on EOS/max_tokens. It binds the SAME ``transformer`` component as the omni +model's ``diffusion_denoise`` loop — one resident MoT instance running both loop types (the §16 +claim). All per-request state (the token list, paged-KV namespace) lives in ``LoopState``. +""" +from __future__ import annotations + +from ..._enums import ExecutionProfile, WorkUnitKind +from ...loop.contracts import ( + Done, + LoopResult, + LoopState, + ResourceRequest, + ShapeSignature, + StepResult, + WorkPlan, +) +from ...request.streams import StreamChunk + + +class ARDecodeLoop: + def __init__(self, *, loop_id, transformer_id="transformer", cost, max_tokens=8): + self.loop_id = loop_id + self.transformer_id = transformer_id + self.cost = cost + self.max_tokens = max_tokens + + def init(self, req, model, ctx) -> LoopState: + seed = req.sampling.seed if req.sampling.seed is not None else 0 + st = LoopState(loop_id=self.loop_id, instance_id=model.card.model_id, + request_id=req.request_id, profile=ctx.profile, seed=seed) + prompt_tokens = ctx.slots.get("prompt_tokens") or [1] + st.scratch["tokens"] = list(prompt_tokens) # prefill (folded into init for the toy) + st.scratch["generated"] = [] + st.scratch["max_tokens"] = min(req.sampling.max_tokens, self.max_tokens) + caches = getattr(model, "caches", None) + if caches is not None and caches.has("paged_kv"): # reasoner paged text-KV (minority case) + caches.pool("paged_kv").allocate(req.request_id, n_blocks=1) + return st + + def next(self, st: LoopState): + if len(st.scratch["generated"]) >= st.scratch["max_tokens"]: + return Done() + tokens = list(st.scratch["tokens"]) + tid = self.transformer_id + + def run(model, override=None): + nxt = model.component(tid).ar_forward(tokens) # und pathway of the shared MoT weights + return StepResult(output={"token": int(nxt)}) + + pos = len(st.scratch["generated"]) + emits = [StreamChunk(stream_id=st.request_id, modality="text", seq=pos, preview=True)] + return WorkPlan( + loop_id=self.loop_id, instance_id=st.instance_id, kind=WorkUnitKind.AR_TOKEN, + shape_sig=ShapeSignature(WorkUnitKind.AR_TOKEN, dims=(len(tokens),)), + resources=ResourceRequest(compute_seconds=self.cost.predict(len(tokens)), + resident_bytes=8 * len(tokens), peak_activation_bytes=8), + payload={"branch": "ar", "pos": pos}, run=run, emits=emits, label=f"ar.tok{pos}") + + def advance(self, st: LoopState, result: StepResult) -> LoopState: + tok = int(result.output["token"]) + st.scratch["tokens"].append(tok) + st.scratch["generated"].append(tok) + st.step_idx += 1 + if st.profile == ExecutionProfile.ROLLOUT: + st.trajectory.append({"token": tok, "pos": st.step_idx}) + return st + + def finalize(self, st: LoopState) -> LoopResult: + toks = st.scratch["generated"] + return LoopResult(outputs={"tokens": toks, "text": "tok:" + ",".join(map(str, toks))}, + metrics={"tokens": float(len(toks))}, + behavior=st.trajectory or None) diff --git a/mini_fastvideo/models/wan21/loop.py b/mini_fastvideo/models/wan21/loop.py index d25f6b676..cbdac3508 100644 --- a/mini_fastvideo/models/wan21/loop.py +++ b/mini_fastvideo/models/wan21/loop.py @@ -94,7 +94,8 @@ class WanDenoiseLoop: peak_activation_bytes=int(x.nbytes)) # one step's transient working buffer emits = [] if st.scratch.get("stream_video"): - emits.append(StreamChunk(stream_id=st.request_id, modality="video", seq=i, preview=True)) + emits.append(StreamChunk(stream_id=st.request_id, modality="video", seq=i, + data=x, preview=True)) # carry the latent as a preview payload return WorkPlan( loop_id=self.loop_id, instance_id=st.instance_id, kind=WorkUnitKind.DIFFUSION_STEP, shape_sig=ShapeSignature(WorkUnitKind.DIFFUSION_STEP, dims=tuple(x.shape), diff --git a/mini_fastvideo/models/wan_causal/loop.py b/mini_fastvideo/models/wan_causal/loop.py index 14490ca43..82202af60 100644 --- a/mini_fastvideo/models/wan_causal/loop.py +++ b/mini_fastvideo/models/wan_causal/loop.py @@ -90,7 +90,7 @@ class ChunkRolloutLoop: emits = [] if i == self.steps_per_chunk - 1: # last step of this chunk → streamable emits.append(StreamChunk(stream_id=st.request_id, modality="video", - seq=st.scratch["chunk_idx"], preview=False)) + seq=st.scratch["chunk_idx"], data=x, preview=False)) res = ResourceRequest( compute_seconds=self.cost.predict(int(np.prod(x.shape)), float(len(branches))), resident_bytes=int(x.nbytes), peak_activation_bytes=int(x.nbytes * len(branches))) diff --git a/mini_fastvideo/parity/interleave_gate.py b/mini_fastvideo/parity/interleave_gate.py index 8210ae5d9..7eba9f3b0 100644 --- a/mini_fastvideo/parity/interleave_gate.py +++ b/mini_fastvideo/parity/interleave_gate.py @@ -45,6 +45,12 @@ def compare_outputs(serial: dict[str, Output], interleaved: dict[str, Output]) - continue for name in sorted(so.artifacts): a, b = _artifact_payload(so.artifacts[name]), _artifact_payload(io.artifacts[name]) + if a is None and b is None: + # both empty is NOT parity — a deferred/aborted request must not pass the gate vacuously + divs.append(Divergence(f"{rid}:{name}", ConsistencyLevel.C1, float("inf"), float("inf"), + "both serial and interleaved produced EMPTY output for a declared " + "artifact (deferred/aborted?) — suspicious, not parity")) + continue if not bit_identical(a, b): divs.append(Divergence(f"{rid}:{name}", ConsistencyLevel.C1, float("nan"), float("nan"), "serial vs interleaved output not bit-identical " diff --git a/mini_fastvideo/runtime/__init__.py b/mini_fastvideo/runtime/__init__.py index 4d20a72c2..09d3af27d 100644 --- a/mini_fastvideo/runtime/__init__.py +++ b/mini_fastvideo/runtime/__init__.py @@ -3,7 +3,14 @@ from __future__ import annotations from .context import RuntimeLoopContext from .engine import Engine, ProgramRunner -from .scheduler import AdmissionController, BatchScheduler, SchedulerMetrics, WorkUnit +from .scheduler import ( + AdmissionController, + AdmissionInfeasible, + BatchScheduler, + SchedulerMetrics, + StepTicket, + WorkUnit, +) __all__ = ["Engine", "ProgramRunner", "RuntimeLoopContext", "AdmissionController", - "BatchScheduler", "SchedulerMetrics", "WorkUnit"] + "AdmissionInfeasible", "BatchScheduler", "SchedulerMetrics", "StepTicket", "WorkUnit"] diff --git a/mini_fastvideo/runtime/context.py b/mini_fastvideo/runtime/context.py index 06515e8b9..6f833447d 100644 --- a/mini_fastvideo/runtime/context.py +++ b/mini_fastvideo/runtime/context.py @@ -41,7 +41,8 @@ class RuntimeLoopContext: override = None if self.interceptors.active and self.state is not None: override = self.interceptors.before(plan, self.state) - if override is not None: + # only count a skip the step body will actually honor (must carry a forward result) + if override is not None and "noise_pred" in override: self.metrics["skipped_steps"] = self.metrics.get("skipped_steps", 0) + 1 t0 = perf_counter() diff --git a/mini_fastvideo/runtime/engine.py b/mini_fastvideo/runtime/engine.py index c9d039e16..098e43bc6 100644 --- a/mini_fastvideo/runtime/engine.py +++ b/mini_fastvideo/runtime/engine.py @@ -30,7 +30,13 @@ from ..request.artifacts import ( from ..request.cancel import CancelKind, CancelScope from ..request.streams import Stream from .context import RuntimeLoopContext -from .scheduler import AdmissionController, SchedulerMetrics +from .scheduler import ( + AdmissionController, + AdmissionInfeasible, + BatchScheduler, + SchedulerMetrics, + WorkUnit, +) def _to_artifact(name: str, value: Any, producer: str = "") -> Any: @@ -108,6 +114,10 @@ class ProgramRunner: 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 @@ -120,13 +130,16 @@ class ProgramRunner: 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"{plan.resources.compute_seconds:.4g}s — exceeds pool/budget capacity") step_res = self.engine.admission.admit_step(plan) if step_res is None: - return False # deferred this tick + return False # deferred this tick (will retry; refundable) self.loop_runner.step() self.engine.admission.release(step_res) - self.engine.metrics.batched_units += 1 - self.engine.metrics.by_kind[plan.kind.value] += 1 + self.engine.metrics.stepped_units += 1 if self.loop_runner.done: self._finish_loop(node) return False # yielded one loop step → return for interleaving @@ -138,21 +151,30 @@ class ProgramRunner: 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.release(self._resident_res) + 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: - guard = 0 - limit = 10_000_000 + stuck = 0 while not self.done: - before = (self.node_idx, self.loop_runner.state.step_idx if self.loop_runner else -1) + p0 = self._progress self.tick() - guard += 1 - after = (self.node_idx, self.loop_runner.state.step_idx if self.loop_runner else -1) - if guard > limit: - raise RuntimeError(f"program {self.program.program_id} did not terminate") + 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] = {} @@ -201,13 +223,22 @@ class Engine: def run_interleaved(self, requests: list[Any]) -> dict[str, Output]: runners = [self._make_runner(req) for req in requests] - guard = 0 - limit = 10_000_000 while not all(r.done for r in runners): + # accounting: how many ready WorkPlans would co-batch this round (BatchScheduler, §6.3) + units = [WorkUnit(r.request.request_id, r.loop_runner.peek()) + for r in runners if not r.done and r.loop_runner is not None + and r.loop_runner.peek() is not None] + if units: + self.metrics.batches += len(BatchScheduler.group(units)) + progressed = False for r in runners: if not r.done: + p0 = r._progress r.tick() - guard += 1 - if guard > limit: - raise RuntimeError("interleaved scheduling did not terminate (admission deadlock?)") + if r.done or r._progress != p0: + progressed = True + if not progressed: # nobody could be admitted → real deadlock, fail fast + raise AdmissionInfeasible( + "interleaved scheduling deadlocked: no request could be admitted this round " + "(jointly infeasible — would require step-boundary preemption, §6.4)") return {r.request.request_id: r.output() for r in runners} diff --git a/mini_fastvideo/runtime/scheduler.py b/mini_fastvideo/runtime/scheduler.py index 7a43e9740..de9da7970 100644 --- a/mini_fastvideo/runtime/scheduler.py +++ b/mini_fastvideo/runtime/scheduler.py @@ -2,13 +2,12 @@ Mini-fastvideo implements the load-bearing layers: * ``AdmissionController`` — the reservation gate (§6.2): do not admit a WorkUnit unless its - compute budget AND memory (resident + worst-case peak) can be reserved. Two requests that - fit individually but jointly OOM are rejected at admission, not at step 37. - * ``BatchScheduler`` — groups compatible WorkPlans by ``shape_sig.batch_key`` (§6.3): image - diffusion and AR decode batch across requests; jumbo video stays batch-of-1. - * cost currency — predicted GPU-time (§6.1), accumulated for fairness/metrics. - -PlacementScheduler/TransferScheduler are single-node trivial here (one pool, in-proc edges). + compute budget AND memory (resident + worst-case peak) can be reserved. Reservations are + *refundable* (compute is a concurrency gate, not a lifetime cap), and infeasible reservations + (need > pool capacity) fail fast with ``AdmissionInfeasible`` instead of spinning. + * ``BatchScheduler`` — groups compatible WorkPlans by ``shape_sig.batch_key`` (§6.3). NOTE: in this + single-process mini, grouping drives *accounting* (how many units would co-batch); kernels still + execute per-plan. Real cross-request batched execution is a GPU-path concern (documented, not faked). """ from __future__ import annotations @@ -20,6 +19,10 @@ from ..loop.contracts import WorkPlan from ..memory.allocator import MemoryManager, Reservation +class AdmissionInfeasible(RuntimeError): + """A reservation can never be satisfied (need > pool capacity) — fail fast, don't spin.""" + + @dataclass class WorkUnit: """The smallest schedulable action (design_v3 §6.1).""" @@ -28,18 +31,25 @@ class WorkUnit: priority: int = 0 +@dataclass +class StepTicket: + """Returned by ``admit_step``; carries what ``release`` must refund (memory + compute).""" + memory_res: Reservation | None + compute_seconds: float + + @dataclass class SchedulerMetrics: admitted: int = 0 deferred: int = 0 batches: int = 0 - batched_units: int = 0 + stepped_units: int = 0 gpu_seconds: float = 0.0 by_kind: dict[str, int] = field(default_factory=lambda: defaultdict(int)) class AdmissionController: - """Reservation before admission (design_v3 §6.2). Compute budget + memory both gate.""" + """Reservation before admission (design_v3 §6.2). Compute budget + memory both gate; both refund.""" def __init__(self, memory: MemoryManager | None = None, compute_budget_seconds: float = float("inf")): self.memory = memory or MemoryManager() @@ -47,6 +57,14 @@ class AdmissionController: self.compute_spent = 0.0 self.metrics = SchedulerMetrics() + def feasible_resident(self, nbytes: int) -> bool: + return nbytes <= self.memory.total_bytes + + def feasible_step(self, plan: WorkPlan) -> bool: + """Could this step EVER be admitted on an empty pool? If not, the caller fails fast.""" + return (plan.resources.peak_activation_bytes <= self.memory.total_bytes + and plan.resources.compute_seconds <= self.compute_budget) + def reserve_resident(self, tag: str, nbytes: int) -> Reservation | None: if nbytes <= 0: return None @@ -54,11 +72,8 @@ class AdmissionController: return None return self.memory.reserve(tag, nbytes) - def admit_step(self, plan: WorkPlan) -> Reservation | None | bool: - """Reserve a step's peak activation + check the compute budget. - - Returns a Reservation (peak reserved), True (no peak to reserve but admitted), - or None (deferred — budget/memory unavailable).""" + def admit_step(self, plan: WorkPlan) -> StepTicket | None: + """Reserve a step's peak activation + check the compute budget. None ⇒ deferred.""" c = plan.resources.compute_seconds if self.compute_spent + c > self.compute_budget: self.metrics.deferred += 1 @@ -74,11 +89,14 @@ class AdmissionController: self.metrics.admitted += 1 self.metrics.gpu_seconds += c self.metrics.by_kind[plan.kind.value] += 1 - return res if res is not None else True + return StepTicket(res, c) - def release(self, res: Reservation | None | bool) -> None: - if isinstance(res, Reservation): - self.memory.release(res) + def release(self, ticket: StepTicket | None) -> None: + if ticket is None: + return + if ticket.memory_res is not None: + self.memory.release(ticket.memory_res) + self.compute_spent = max(0.0, self.compute_spent - ticket.compute_seconds) # refund class BatchScheduler: diff --git a/mini_fastvideo/tests/test_review_fixes.py b/mini_fastvideo/tests/test_review_fixes.py new file mode 100644 index 000000000..eb17e5df4 --- /dev/null +++ b/mini_fastvideo/tests/test_review_fixes.py @@ -0,0 +1,77 @@ +"""Regression tests for the adversarial-review findings (untested paths now covered). + +Locks in: fail-fast on infeasible/deadlocked admission (no busy-spin), refundable compute budget +(concurrency gate, not lifetime cap), §7.1 feature-key partitioning by adapter stack + precision, +component-scoped weight invalidation (transformer sync keeps the text-encoder cache), and the +interleave gate flagging symmetric-empty output. +""" +from __future__ import annotations + +import numpy as np + +from mini_fastvideo._enums import WorkUnitKind +from mini_fastvideo.cache import CacheManager +from mini_fastvideo.card import load_card +from mini_fastvideo.loop.contracts import ResourceRequest, ShapeSignature, WorkPlan +from mini_fastvideo.memory import MemoryManager +from mini_fastvideo.models import build_default_engine +from mini_fastvideo.models.common import cached_text_encode +from mini_fastvideo.models.wan21 import build_wan21_card +from mini_fastvideo.parity import compare_outputs +from mini_fastvideo.request import DiffusionParams, Output, TaskType, VideoArtifact, make_request +from mini_fastvideo.runtime import AdmissionController, AdmissionInfeasible, Engine + + +def test_infeasible_reservation_fails_fast_not_busy_spin(): + eng = build_default_engine(Engine(admission=AdmissionController(MemoryManager(total_bytes=16)))) + req = make_request(TaskType.T2V, "wan2.1-1.3b", "x", diffusion=DiffusionParams(num_steps=3, seed=1)) + try: + eng.run(req) + assert False, "expected AdmissionInfeasible (resident need >> 16 bytes)" + except AdmissionInfeasible: + pass + + +def test_compute_budget_is_a_refundable_concurrency_gate(): + ac = AdmissionController(compute_budget_seconds=1.0) + plan = WorkPlan(loop_id="l", instance_id="i", kind=WorkUnitKind.DIFFUSION_STEP, + shape_sig=ShapeSignature(WorkUnitKind.DIFFUSION_STEP), + resources=ResourceRequest(compute_seconds=0.6)) + t1 = ac.admit_step(plan) + assert t1 is not None and abs(ac.compute_spent - 0.6) < 1e-9 + assert ac.admit_step(plan) is None # 0.6+0.6 > 1.0 ⇒ deferred (concurrency gate) + ac.release(t1) + assert ac.compute_spent == 0.0 # refunded — NOT a permanent lifetime cap + assert ac.admit_step(plan) is not None # fits again after the refund + + +def test_feature_key_partitions_by_adapter_stack(): + card = build_wan21_card() + inst = load_card(card, cache_manager=CacheManager.from_card(card)) + inst.adapter_versions = {"te_lora": "1"} + cached_text_encode(inst, "p") # miss + store (adapter stack 1) + h = inst.caches.stats()["feature"]["hits"] + cached_text_encode(inst, "p") # same stack + text ⇒ hit + assert inst.caches.stats()["feature"]["hits"] == h + 1 + inst.adapter_versions = {"te_lora": "2"} # swap te-LoRA stack on the resident instance + m = inst.caches.stats()["feature"]["misses"] + cached_text_encode(inst, "p") # different stack ⇒ MISS, never a stale hit (§7.1) + assert inst.caches.stats()["feature"]["misses"] == m + 1 + + +def test_transformer_weight_sync_keeps_text_encoder_cache(): + card = build_wan21_card() + inst = load_card(card, cache_manager=CacheManager.from_card(card)) + cached_text_encode(inst, "p") + assert inst.caches.stats()["feature"]["used_bytes"] > 0 + inst.set_weights_version("w1", components=["transformer"]) # RL-style transformer-only sync + h = inst.caches.stats()["feature"]["hits"] + cached_text_encode(inst, "p") # text-encoder embed survived ⇒ hit (not flushed) + assert inst.caches.stats()["feature"]["hits"] == h + 1 + + +def test_interleave_gate_flags_symmetric_empty_output(): + empty = {"r": Output(request_id="r", + artifacts={"video": VideoArtifact(name="video", frames=None)})} + divs = compare_outputs(empty, empty) # identical BUT empty — must not pass vacuously + assert divs, "both-empty output must be flagged, not treated as parity" diff --git a/mini_fastvideo/training/weight_sync.py b/mini_fastvideo/training/weight_sync.py index 1acf47eb4..1f3997cc3 100644 --- a/mini_fastvideo/training/weight_sync.py +++ b/mini_fastvideo/training/weight_sync.py @@ -45,8 +45,10 @@ class WeightSyncPlan: dst_dit.blend_from(src_dit, self.decay) # decay·dst + (1-decay)·src else: dst_dit.copy_from(src_dit) # hard copy (student push / reference init) - # 4) bump version + 5) invalidate caches/graphs + 6) publish (set_weights_version does 4-6) + # 4) bump version + 5) invalidate the TRANSFORMER's caches/graphs + 6) publish. + # Scoped to the transformer (the synced weights) so frozen text/vision encoders keep + # their feature caches (design_v3 §7.1). version = f"w{next(_ver)}" if dst_instance is not None: - dst_instance.set_weights_version(version) + dst_instance.set_weights_version(version, components=["transformer"]) return version