SolitaryThinker
db866dfb04
Sweep all 296 v2 modules: simplify verbose docstrings/comments and remove 376 dangling "(design_vN §X)" citations to the now-deleted design docs (v2/README.md is the source of truth). Comment/docstring-only — AST-verified code-identical; the CPU suite holds at 237 passed / 5 pre-existing. Also applies yapf + ruff --fix auto-fixes (import ordering, forward-ref annotation de-quoting under `from __future__ import annotations`; behavior- neutral, suite-confirmed) and adds the legitimate domain terms mot/clen/te to the codespell ignore-list. Remaining ruff (24) + mypy (68 no-untyped-call) findings are pre-existing v2 debt, untouched here.
59 lines
2.2 KiB
Python
59 lines
2.2 KiB
Python
"""Admission = reservation before stepping.
|
|
|
|
Two requests that fit individually but jointly OOM are rejected at admission, not at step 37.
|
|
And crucially: admission-forced serialization does NOT change outputs (still bit-identical).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from v2.memory import MemoryManager, OutOfMemory
|
|
from v2.recipes import build_default_engine
|
|
from v2.recipes.wan21.loop import latent_shape
|
|
from v2.parity import compare_outputs
|
|
from v2.request import DiffusionParams, TaskType, make_request
|
|
from v2.runtime import AdmissionController, Engine
|
|
|
|
|
|
def test_memory_manager_reserve_release_and_oom():
|
|
mm = MemoryManager(total_bytes=100)
|
|
r = mm.reserve("t", 60)
|
|
assert mm.reserved == 60 and mm.available == 40
|
|
try:
|
|
mm.reserve("t", 60)
|
|
assert False, "expected OutOfMemory"
|
|
except OutOfMemory:
|
|
pass
|
|
mm.release(r)
|
|
assert mm.reserved == 0
|
|
|
|
|
|
def test_sleep_wake_frees_by_tag():
|
|
mm = MemoryManager(total_bytes=1000)
|
|
mm.reserve("dit", 400)
|
|
mm.reserve("vae", 200)
|
|
freed = mm.sleep(["dit"]) # component-granular: drop DiT, keep VAE
|
|
assert freed == 400 and mm.reserved == 200
|
|
|
|
|
|
def _req(prompt, seed):
|
|
return make_request(TaskType.T2V, "wan2.1-1.3b", prompt,
|
|
diffusion=DiffusionParams(num_steps=3, seed=seed))
|
|
|
|
|
|
def test_jointly_oom_defers_but_output_is_identical():
|
|
reqs = [_req("a", 1), _req("b", 2)]
|
|
# compute resident (latent + 2 text embeds) and peak (one latent) for the budget
|
|
nb = int(np.prod(latent_shape(reqs[0]))) * 4
|
|
cond = 2 * (4 * 8 * 4) # two (4,8) float32 embeds
|
|
R, P = nb + cond, nb
|
|
budget = R + P + (R - P) // 2 # fits ONE request's step, not TWO residents
|
|
|
|
big = build_default_engine() # ample memory
|
|
small = build_default_engine(Engine(admission=AdmissionController(MemoryManager(total_bytes=budget))))
|
|
|
|
serial = big.run_serial(reqs)
|
|
interleaved_small = small.run_interleaved(reqs)
|
|
assert small.metrics.deferred > 0, "second request should be deferred at admission"
|
|
assert not compare_outputs(serial, interleaved_small), "deferral must not change outputs"
|