Will Lin
1fe50c0092
The v2 top level had grown to ~27 dirs + 13 loose files — half of them
single-concept abstraction shells (memory/ 86 LOC, transport/ 222, parity/ 178,
extend/ 226) and half vendored fastvideo code sitting as peers to the actual v2
design. Regroup to mirror README section 3 "Planes & dependency order":
core/ enums+types, card, loop, program, parity, request, parallel
(the model-native contracts; no kernels)
runtime/ + folded-in substrate: cache, memory, transport, extend
(the import graph shows only runtime consumes them; compile/cudagraph
already lived here)
serving/ + deploy/ (products / fleet)
_vendor/ all copied fastvideo: models, layers, attention, configs, distributed,
platforms, api, hooks, logging_utils, third_party + fastvideo_args/
utils/logger/envs/forward_context — internal layout unchanged, still
mirrors upstream for diffing
28 dirs + 13 root files -> 8 dirs + 6 root files. Pure mechanical move: 947
absolute-import paths rewritten (v2.X -> v2.{core,runtime,serving,_vendor}.X),
boundary-anchored so platform/ (native dispatch) and platforms/ (vendored CUDA
detect) no longer collide and neither does hooks/ vs extend/. Deleted the empty
v2/loader/. README section 3 + 16 updated; stale test count corrected
(34 files/216 tests -> 22 files/143 tests).
Validated: `import v2` stays torch-free; `v2/run_tests.py` and `pytest v2/tests/`
-> 143 passed, 2 skipped on the numpy toy backend. (On a GPU box force the toy
backend with CUDA_VISIBLE_DEVICES="" or detect() picks cuda.) No external importer
changed — examples use the re-exported `from v2 import VideoGenerator`.
81 lines
2.5 KiB
Python
81 lines
2.5 KiB
Python
"""Tagged memory pools with reservation-before-admission.
|
|
|
|
Reservation must be pre-flight: a scheduler that admits a diffusion step and then discovers it
|
|
cannot allocate the VAE tile is wrong.
|
|
|
|
Sleep/wake is component-granular (tags are component names) for RL: drop DiT + caches,
|
|
keep VAE/text-encoder resident (CuMem-style).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import itertools
|
|
from dataclasses import dataclass
|
|
|
|
|
|
class OutOfMemory(Exception):
|
|
pass
|
|
|
|
|
|
_res_ctr = itertools.count(1)
|
|
|
|
|
|
@dataclass
|
|
class Reservation:
|
|
res_id: int
|
|
tag: str
|
|
nbytes: int
|
|
active: bool = True
|
|
|
|
|
|
class MemoryManager:
|
|
|
|
def __init__(self, total_bytes: int = 1 << 40, per_tag_budget: dict[str, int] | None = None):
|
|
self.total_bytes = total_bytes
|
|
self.per_tag_budget = dict(per_tag_budget or {})
|
|
self.reserved = 0
|
|
self._by_tag: dict[str, int] = {}
|
|
self._reservations: dict[int, Reservation] = {}
|
|
self._asleep: set[str] = set()
|
|
|
|
@property
|
|
def available(self) -> int:
|
|
return self.total_bytes - self.reserved
|
|
|
|
def can_reserve(self, tag: str, nbytes: int) -> bool:
|
|
if nbytes > self.available:
|
|
return False
|
|
budget = self.per_tag_budget.get(tag)
|
|
return budget is None or self._by_tag.get(tag, 0) + nbytes <= budget
|
|
|
|
def reserve(self, tag: str, nbytes: int) -> Reservation:
|
|
if not self.can_reserve(tag, nbytes):
|
|
raise OutOfMemory(f"cannot reserve {nbytes} bytes for tag {tag!r} "
|
|
f"(available={self.available}, used={self._by_tag.get(tag, 0)})")
|
|
res = Reservation(next(_res_ctr), tag, nbytes)
|
|
self._reservations[res.res_id] = res
|
|
self.reserved += nbytes
|
|
self._by_tag[tag] = self._by_tag.get(tag, 0) + nbytes
|
|
return res
|
|
|
|
def release(self, res: Reservation) -> None:
|
|
if not res.active:
|
|
return
|
|
res.active = False
|
|
self._reservations.pop(res.res_id, None)
|
|
self.reserved -= res.nbytes
|
|
self._by_tag[res.tag] = max(0, self._by_tag.get(res.tag, 0) - res.nbytes)
|
|
|
|
# component-granular sleep/wake (tags = component names) ------------------- #
|
|
def sleep(self, tags: list[str]) -> int:
|
|
freed = 0
|
|
for tag in tags:
|
|
self._asleep.add(tag)
|
|
for res in [r for r in self._reservations.values() if r.tag == tag]:
|
|
freed += res.nbytes
|
|
self.release(res)
|
|
return freed
|
|
|
|
def wake(self, tags: list[str]) -> None:
|
|
for tag in tags:
|
|
self._asleep.discard(tag)
|