diff --git a/src/gen_worker/models/checkpoint_juggle.py b/src/gen_worker/models/checkpoint_juggle.py index ca5632143..a12a7e9e5 100644 --- a/src/gen_worker/models/checkpoint_juggle.py +++ b/src/gen_worker/models/checkpoint_juggle.py @@ -89,7 +89,7 @@ from dataclasses import dataclass from enum import Enum from pathlib import Path -from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple +from typing import Any, Callable, Dict, List, Optional, Set, Tuple from .._vendor.tensorfs.layout2 import ExpectedHeader, LayoutTensor from .arena_residency import ( @@ -100,11 +100,16 @@ dlpack_dtype, ) from .safetensors_header import header_len_ok +from ..serving.streaming.fill_client import ( + AddressSource, + CudaFillClient, + Destination, + FileSource, + HostFillClient, +) logger = logging.getLogger(__name__) -_DL_UINT8 = (1, 8) - DEFAULT_HOST_FLOOR_BYTES = 4 << 30 _SAFETENSORS_DTYPES: Dict[str, Tuple[int, int]] = { @@ -420,7 +425,6 @@ def __init__( *, torch_mod: Any, varena_mod: Any = None, - engine: Any = None, pin: bool = True, ) -> None: self.checkpoint_id = checkpoint_id @@ -456,27 +460,43 @@ def __init__( if buffer is None: buffer = torch_mod.empty(n, dtype=torch_mod.uint8) self.buffer = buffer - self._build(engine) + self._build() - def _build(self, engine: Any) -> None: + def _build(self) -> None: torch = self._torch - straight: List[Tuple[str, int, int, int, int]] = [] casts: List[Tuple[SlotSpec, SlotSource]] = [] base_ptr = int(self.buffer.data_ptr()) - for region in self.layout.regions: - for slot in region.slots: - src = self.manifest[_slot_key(slot)] - if (src.dtype_code, src.dtype_bits) == (slot.dtype_code, slot.dtype_bits): - straight.append( - (str(src.path), src.offset, src.length, base_ptr + slot.offset, 0) - ) - else: - casts.append((slot, src)) - if straight: - if engine is not None: - engine.submit(straight, 0).wait() + sources: List[FileSource] = [] + cursor = 0 + slots = sorted( + (slot for region in self.layout.regions for slot in region.slots), + key=lambda slot: slot.offset, + ) + for slot in slots: + if slot.offset > cursor: + sources.append(FileSource(None, 0, slot.offset - cursor)) + src = self.manifest[_slot_key(slot)] + if (src.dtype_code, src.dtype_bits) == (slot.dtype_code, slot.dtype_bits): + sources.append(FileSource(str(src.path), src.offset, src.length)) else: - self._pread(straight, base_ptr) + sources.append(FileSource(None, 0, slot.nbytes)) + casts.append((slot, src)) + cursor = slot.offset + slot.nbytes + if cursor < self.layout.virtual_bytes: + sources.append(FileSource(None, 0, self.layout.virtual_bytes - cursor)) + if sources: + HostFillClient().fill_files( + sources, + Destination( + name=self.checkpoint_id, + pointer=base_ptr, + capacity=self.layout.virtual_bytes, + source_offset=0, + shape=(self.layout.virtual_bytes,), + element_bytes=1, + layout="torch.contiguous@1", + ), + ) for slot, src in casts: file_dtype = _torch_dtype(torch, src.dtype_code, src.dtype_bits) lane_dtype = _torch_dtype(torch, slot.dtype_code, slot.dtype_bits) @@ -495,21 +515,6 @@ def _build(self, engine: Any) -> None: for region in self.layout.regions: self.region_digests[region.name] = self.region_digest(region) - def _pread( - self, requests: Sequence[Tuple[str, int, int, int, int]], base_ptr: int - ) -> None: - mv = memoryview(self.buffer.numpy()) - for path, offset, length, host_ptr, _dev in requests: - start = host_ptr - base_ptr - with open(path, "rb") as fh: - fh.seek(offset) - got = fh.readinto(mv[start : start + length]) - if got != length: - raise JuggleRefusal( - f"{self.checkpoint_id}: short read from {path} " - f"({got} of {length} bytes)" - ) - def region_bytes(self, region: RegionSpec) -> Any: return self.buffer[region.offset : region.offset + region.span] @@ -559,14 +564,12 @@ def __init__( *, torch_mod: Any, varena_mod: Any = None, - engine_factory: Optional[Callable[[], Any]] = None, host_floor_bytes: int = DEFAULT_HOST_FLOOR_BYTES, mem_available: Callable[[], int] = _mem_available_bytes, ) -> None: self.layout = layout self._torch = torch_mod self._varena = varena_mod - self._engine_factory = engine_factory self.host_floor_bytes = int(host_floor_bytes) self._mem_available = mem_available self.manifests: Dict[str, Dict[str, SlotSource]] = {} @@ -630,10 +633,9 @@ def ensure_warm(self, checkpoint_id: str) -> Optional[CheckpointImage]: self.host_floor_bytes / (1 << 30), self.pressure_epoch, ) return None - engine = self._engine_factory() if self._engine_factory else None image = CheckpointImage( checkpoint_id, self.layout, manifest, - torch_mod=self._torch, varena_mod=self._varena, engine=engine, + torch_mod=self._torch, varena_mod=self._varena, ) self.images[checkpoint_id] = image self._lru.append(checkpoint_id) @@ -775,7 +777,6 @@ def __init__( self.layout, torch_mod=residency._torch, varena_mod=residency._varena, - engine_factory=residency._engine_for, ) if not residency.adopted: raise ValueError( @@ -791,6 +792,9 @@ def __init__( self.rearms = 0 self.reports: List[SwitchReport] = [] self._residue_bytes = self._count_residue() + self._fill_client = CudaFillClient( + 64 << 20, int(residency.device.index or 0) + ) def admit(self, checkpoint_id: str, manifest: Dict[str, SlotSource]) -> None: self.catalog.admit(checkpoint_id, manifest) @@ -838,18 +842,13 @@ def switch_to(self, checkpoint_id: str) -> SwitchReport: from_id = self.serving_id device = self.residency.device - stream = torch.cuda.Stream(device=device) - stream.wait_stream(torch.cuda.current_stream(device)) with torch.no_grad(): self.residency.ring.drain() for region in self.layout.regions: if self.residency.is_resident(region.name): self.ledger.begin(region.name, checkpoint_id) try: - with torch.cuda.stream(stream): - bytes_moved += self._refill_backed( - region, image, manifest, stream - ) + bytes_moved += self._refill_backed(region, image, manifest) except Exception: self.ledger.poison(region.name) logger.error( @@ -876,7 +875,6 @@ def switch_to(self, checkpoint_id: str) -> SwitchReport: key: (src.path, src.offset, src.length) for key, src in manifest.items() } - torch.cuda.current_stream(device).wait_stream(stream) torch.cuda.synchronize(device) backing_verified = self._verify_backing(checkpoint_id) @@ -923,17 +921,31 @@ def _refill_backed( region: RegionSpec, image: Optional[CheckpointImage], manifest: Dict[str, SlotSource], - stream: Any, ) -> int: - if image is not None: - dst = self._device_bytes(region) - src = image.region_bytes(region) - dst.copy_(src, non_blocking=image.pinned) - return region.span base = int(self.residency.reservation.base_ptr) - requests = [] - moved = 0 - for slot in region.slots: + destination = Destination( + name=region.name, + pointer=base + region.offset, + capacity=region.span, + source_offset=0, + shape=(region.span,), + element_bytes=1, + layout="torch.contiguous@1", + ) + if image is not None: + stats = self._fill_client.fill_address( + AddressSource( + pointer=int(image.buffer.data_ptr()) + region.offset, + capacity=region.span, + ), + destination, + ) + return int(stats.destination_bytes) + sources: List[FileSource] = [] + cursor = region.offset + for slot in sorted(region.slots, key=lambda item: item.offset): + if slot.offset > cursor: + sources.append(FileSource(None, 0, slot.offset - cursor)) src = manifest[_slot_key(slot)] if (src.dtype_code, src.dtype_bits) != (slot.dtype_code, slot.dtype_bits): raise JuggleRefusal( @@ -941,23 +953,13 @@ def _refill_backed( f"flight; warm this checkpoint first (the image is where " f"casts happen, once)" ) - requests.append( - (str(src.path), src.offset, src.length, 0, base + slot.offset) - ) - moved += src.length - handle = self.residency._engine_for().submit( - requests, int(stream.cuda_stream) if stream is not None else 0 - ) - handle.wait() - return moved - - def _device_bytes(self, region: RegionSpec) -> Any: - torch = self.residency._torch - return torch.from_dlpack( - self.residency.reservation.tensor( - region.offset, [region.span], *_DL_UINT8 - ) - ) + sources.append(FileSource(str(src.path), src.offset, src.length)) + cursor = slot.offset + slot.nbytes + end = region.offset + region.span + if cursor < end: + sources.append(FileSource(None, 0, end - cursor)) + stats = self._fill_client.fill_files(sources, destination) + return int(stats.destination_bytes) def _count_residue(self) -> int: from .stream_residency import own_tensors, tensor_bytes diff --git a/src/gen_worker/serving/streaming/__init__.py b/src/gen_worker/serving/streaming/__init__.py index 1fa19ae2c..572e03165 100644 --- a/src/gen_worker/serving/streaming/__init__.py +++ b/src/gen_worker/serving/streaming/__init__.py @@ -13,18 +13,16 @@ TensorRow, ) from .engine import LoadError, LoadReport, NameMismatch, StreamingLoader +from .fill_client import Destination from .skeleton import Skeleton, SkeletonError from .source import ( - BridgeWeightStore, NativeWeightStore, StreamedTensor, TensorStream, WeightStore, WeightStoreUnavailable, - native_available, store_for, ) -from .staging import StagingPool logger = logging.getLogger(__name__) @@ -57,12 +55,12 @@ def engine_for( __all__ = [ "engine_for", - "BridgeWeightStore", "CENSUS_KIND", "Census", "CensusError", "CensusMismatch", "ComponentCensus", + "Destination", "TensorRow", "LoadError", "LoadReport", @@ -70,12 +68,10 @@ def engine_for( "NativeWeightStore", "Skeleton", "SkeletonError", - "StagingPool", "StreamedTensor", "StreamingLoader", "TensorStream", "WeightStore", "WeightStoreUnavailable", - "native_available", "store_for", ] diff --git a/src/gen_worker/serving/streaming/engine.py b/src/gen_worker/serving/streaming/engine.py index 204c3495b..71aaf5c38 100644 --- a/src/gen_worker/serving/streaming/engine.py +++ b/src/gen_worker/serving/streaming/engine.py @@ -1,4 +1,4 @@ -"""``ctx.load``'s engine: chunk store -> pinned staging -> device, no files. +"""``ctx.load``'s engine: plan destinations, ask tensorfs to fill, fence. The load's postcondition is not asserted here. It is stated by the CONSTRUCTION CENSUS (:mod:`.census`, pgw#1647) and this engine REPLAYS it: after the fill and @@ -22,20 +22,22 @@ from dataclasses import dataclass from pathlib import Path from typing import ( - TYPE_CHECKING, Any, Dict, FrozenSet, List, Mapping, Optional, Sequence, Tuple, + TYPE_CHECKING, Any, Dict, FrozenSet, List, Mapping, Optional, Tuple, ) from . import census as _census from . import keymap as _keymap from . import skeleton as _skeleton -from .source import StreamedTensor, TensorStream, WeightStore, component_of -from .staging import DEFAULT_BUFFER_BYTES, DEFAULT_BUFFERS, StagingPool +from .fill_client import Destination, FillClient, client_for +from .source import WeightStore, component_of if TYPE_CHECKING: # pragma: no cover - typing only import torch logger = logging.getLogger(__name__) +DEFAULT_STAGING_BYTES = 64 * 1024 * 1024 + _TORCH_DTYPE: Mapping[str, str] = { "F64": "float64", "F32": "float32", @@ -79,12 +81,12 @@ class LoadReport: weights_streamed_bytes: int = 0 weights_stream_gbps: float = 0.0 - source: str = "bridge" - staging: str = "pageable" + source: str = "native" + staging: str = "destination" io: str = "buffered" containers: int = 0 tensors: int = 0 - windows: int = 0 + chunks: int = 0 seconds: float = 0.0 dtypes: Tuple[str, ...] = () cast_to_lane: int = 0 @@ -102,19 +104,6 @@ def attributes(self) -> Dict[str, object]: } -@dataclass(slots=True) -class _Placement: - - name: str - offset: int - nbytes: int - flat: "torch.Tensor" - - @property - def end(self) -> int: - return self.offset + self.nbytes - - @dataclass(slots=True) class _Slot: @@ -172,12 +161,6 @@ def _install(slot: _Slot, tensor: "torch.Tensor") -> None: slot.owner._buffers[slot.leaf] = tensor -def _flat_bytes(tensor: "torch.Tensor") -> "torch.Tensor": - import torch - - return tensor.view(-1).view(torch.uint8) - - class StreamingLoader: def __init__( @@ -186,16 +169,16 @@ def __init__( *, device: Any = "cuda", io: str = "buffered", - buffer_bytes: int = DEFAULT_BUFFER_BYTES, - buffers: int = DEFAULT_BUFFERS, + staging_bytes: int = DEFAULT_STAGING_BYTES, ) -> None: if io not in ("buffered", "direct"): raise LoadError(f"io must be 'buffered' or 'direct', not {io!r}") self._store = store self._device = device self._io = io - self._buffer_bytes = int(buffer_bytes) - self._buffers = int(buffers) + if staging_bytes <= 0: + raise LoadError("tensorfs staging must hold bytes") + self._staging_bytes = int(staging_bytes) self._planned: List[Tuple[str, str]] = [] self.last_report: Optional[LoadReport] = None #: The census the last load actually produced and the fence accepted. @@ -206,7 +189,14 @@ def last_census(self) -> Optional["_census.Census"]: """What the last load BUILT, as data — the fence's own answer.""" return self._census - def build(self, pipeline_cls: type, *, checkpoint_dir: Path, lane: Any) -> Any: + def build( + self, + pipeline_cls: type, + *, + checkpoint_dir: Path, + lane: Any, + expected_census: Optional[_census.Census] = None, + ) -> Any: """Meta skeleton, then weights streamed into it.""" import torch @@ -222,32 +212,36 @@ def build(self, pipeline_cls: type, *, checkpoint_dir: Path, lane: Any) -> Any: # release's census beside the resolved variant, `expected` is read from # it instead and the same predicate then also catches an image that # builds a different module than the one the release was derived from. - expected = self._expected_census(built) + expected = ( + expected_census + if expected_census is not None + else self._expected_census(built) + ) report = LoadReport( io=self._io, source=str(getattr(self._store, "KIND", type(self._store).__name__)), ) recasts: List[Tuple[_Slot, str, Any]] = [] - with StagingPool( - device, - buffer_bytes=self._buffer_bytes, - buffers=self._buffers, - ) as pool: - report.staging = pool.staging - for component, container in self._plan(built.modules): - self._stream_container( - built.modules[component], - component, - container, - pool=pool, - device=device, - report=report, - compute_dtype=compute_dtype, - recasts=recasts, - lane_exempt=_lane_exempt(built, component), - ) - report.containers = len(self._planned) + fill_client = client_for( + device.type, + device_index=int(device.index or 0), + staging_bytes=self._staging_bytes, + ) + report.staging = fill_client.staging + for component, container in self._plan(built.modules): + self._fill_container( + built.modules[component], + component, + container, + fill_client=fill_client, + device=device, + report=report, + compute_dtype=compute_dtype, + recasts=recasts, + lane_exempt=_lane_exempt(built, component), + ) + report.containers = len(self._planned) self._cast_to_lane(recasts, compute_dtype, report) # The prepare seam's TRAILING half, once per component and in one @@ -281,7 +275,7 @@ def build(self, pipeline_cls: type, *, checkpoint_dir: Path, lane: Any) -> Any: self._assert_lane_dtype(built.modules, compute_dtype, built) logger.info( "ctx.load: %s resident on %s — %.2f GiB streamed in %.2fs " - "(%.2f GB/s, staging=%s io=%s, %d tensors over %d windows)", + "(%.2f GB/s, staging=%s io=%s, %d tensors over %d chunks)", pipeline_cls.__name__, device, report.weights_streamed_bytes / (1 << 30), @@ -290,7 +284,7 @@ def build(self, pipeline_cls: type, *, checkpoint_dir: Path, lane: Any) -> Any: report.staging, report.io, report.tensors, - report.windows, + report.chunks, ) return built.pipeline @@ -326,13 +320,13 @@ def _plan(self, modules: Mapping[str, Any]) -> List[Tuple[str, str]]: self._planned = planned return planned - def _stream_container( + def _fill_container( self, module: Any, component: str, container: str, *, - pool: StagingPool, + fill_client: FillClient, device: Any, report: LoadReport, compute_dtype: Any = None, @@ -341,16 +335,16 @@ def _stream_container( ) -> None: import torch - stream: TensorStream = self._store.open( + stream = self._store.open( container, direct=self._io == "direct" ) - entries: Sequence[StreamedTensor] = stream.tensors + entries = stream.tensors if not entries: return slots = _slots(module) renames = _keymap.migration(module, (entry.name for entry in entries)) - placements: List[_Placement] = [] + destinations: List[Destination] = [] unexpected: List[str] = [] seen: set[str] = set() dtypes: set[str] = set(report.dtypes) @@ -366,18 +360,17 @@ def _stream_container( continue dtype = _torch_dtype(entry.dtype, f"{component}/{entry.name}") dtypes.add(entry.dtype.upper()) - destination = torch.empty( + tensor = torch.empty( tuple(int(dim) for dim in entry.shape), dtype=dtype, device=device ) - flat = _flat_bytes(destination) - if flat.numel() != entry.nbytes: + capacity = int(tensor.numel() * tensor.element_size()) + if capacity != entry.nbytes: raise LoadError( f"{component}/{entry.name}: the container says " f"{entry.nbytes} bytes, a {dtype} tensor of " - f"{tuple(entry.shape)} holds {flat.numel()}" + f"{tuple(entry.shape)} holds {capacity}" ) - pool.track(destination) - _install(slot, destination) + _install(slot, tensor) # pgw#1638: a tensor a QUANTIZER owns is out of the lane cast's # scope, whatever its dtype. Property 3 above already says a # quantized container is the lane's own bytes — but it said it of @@ -391,12 +384,15 @@ def _stream_container( and name not in lane_exempt): recasts.append((slot, f"{component}/{entry.name}", dtype)) seen.add(entry.name) - placements.append( - _Placement( + destinations.append( + Destination( name=entry.name, - offset=int(entry.offset), - nbytes=int(entry.nbytes), - flat=flat, + pointer=int(tensor.data_ptr()), + capacity=capacity, + source_offset=int(entry.offset), + shape=tuple(int(dim) for dim in entry.shape), + element_bytes=int(tensor.element_size()), + layout="torch.contiguous@1", ) ) @@ -419,55 +415,18 @@ def _stream_container( ) report.dtypes = tuple(sorted(dtypes)) - placements.sort(key=lambda placement: placement.offset) - report.tensors += len(placements) - report.windows += self._walk(stream, placements, pool=pool, report=report) - - def _walk( - self, - stream: TensorStream, - placements: List[_Placement], - *, - pool: StagingPool, - report: LoadReport, - ) -> int: - position = placements[0].offset - finish = placements[-1].end - window = pool.buffer_bytes - first = 0 - windows = 0 - - while position < finish: - count = min(window, finish - position) - slot = pool.acquire() - stream.readinto(position, count, slot.view[:count]) - windows += 1 - report.weights_streamed_bytes += count - - index = first - while index < len(placements) and placements[index].offset < position + count: - placement = placements[index] - if placement.end <= position: - index += 1 - first = index - continue - low = max(placement.offset, position) - high = min(placement.end, position + count) - pool.copy_out( - slot, - low - position, - placement.flat, - low - placement.offset, - high - low, + destinations.sort(key=lambda destination: destination.source_offset) + report.tensors += len(destinations) + for destination_data in destinations: + stats = fill_client.fill(stream, destination_data) + if int(stats.source_bytes) != destination_data.capacity: + raise LoadError( + f"{component}/{destination_data.name}: tensorfs filled " + f"{stats.source_bytes} source bytes into a " + f"{destination_data.capacity}-byte destination" ) - if placement.end <= position + count: - first = index + 1 - index += 1 - - pool.release(slot) - position += count - - return windows + report.weights_streamed_bytes += int(stats.source_bytes) + report.chunks += int(stats.chunks) @staticmethod def _cast_to_lane( @@ -565,6 +524,7 @@ def _lane_compute_dtype(lane: Any) -> Any: __all__ = [ "LaneDtypeUnmet", + "DEFAULT_STAGING_BYTES", "LoadError", "LoadReport", "NameMismatch", diff --git a/src/gen_worker/serving/streaming/fill_client.py b/src/gen_worker/serving/streaming/fill_client.py new file mode 100644 index 000000000..7e18f6277 --- /dev/null +++ b/src/gen_worker/serving/streaming/fill_client.py @@ -0,0 +1,165 @@ +"""The torch-free data seam into tensorfs's one fill implementation.""" + +from __future__ import annotations + +import importlib +from dataclasses import dataclass +from typing import Any, Optional, Protocol, Sequence, Tuple + + +@dataclass(frozen=True, slots=True) +class AddressSource: + """One contiguous source allocation, described without its owner type.""" + + pointer: int + capacity: int + + +@dataclass(frozen=True, slots=True) +class FileSource: + """One immutable file range, or a zero-filled hole when ``path`` is None.""" + + path: Optional[str] + offset: int + length: int + + +@dataclass(frozen=True, slots=True) +class Destination: + """Where one named tensor lands; every field is plain data.""" + + name: str + pointer: int + capacity: int + source_offset: int + shape: Tuple[int, ...] + element_bytes: int + layout: str + + +class FillClient(Protocol): + """Fill destinations without learning which allocator produced them.""" + + staging: str + + def fill(self, reader: Any, destination: Destination) -> Any: ... + + def fill_address( + self, source: AddressSource, destination: Destination + ) -> Any: ... + + def fill_files( + self, sources: Sequence[FileSource], destination: Destination + ) -> Any: ... + + +class HostFillClient: + """The host destination is its own staging allocation.""" + + staging = "destination" + + def fill(self, reader: Any, destination: Destination) -> Any: + return reader.fill_host_address( + destination.name, + destination.pointer, + destination.capacity, + layout=destination.layout, + ) + + @staticmethod + def _native() -> Any: + native = importlib.import_module("tensorfs.native") + NativeHostFillClient = getattr(native, "HostFillClient") + return NativeHostFillClient() + + def fill_address( + self, source: AddressSource, destination: Destination + ) -> Any: + return self._native().fill_address( + source.pointer, + source.capacity, + destination.pointer, + destination.capacity, + destination.shape, + destination.element_bytes, + layout=destination.layout, + ) + + def fill_files( + self, sources: Sequence[FileSource], destination: Destination + ) -> Any: + records = [(source.path, source.offset, source.length) for source in sources] + return self._native().fill_files( + records, + destination.pointer, + destination.capacity, + destination.shape, + destination.element_bytes, + layout=destination.layout, + ) + + +class CudaFillClient: + """One reusable tensorfs-owned pinned slab for all CUDA destinations.""" + + staging = "tensorfs-pinned" + + def __init__(self, staging_bytes: int, device: int) -> None: + native = importlib.import_module("tensorfs.native") + NativeCudaFillClient = getattr(native, "CudaFillClient") + self._client = NativeCudaFillClient(staging_bytes, device) + + def fill(self, reader: Any, destination: Destination) -> Any: + native_reader = getattr(reader, "native_reader", reader) + return self._client.fill( + native_reader, + destination.name, + destination.pointer, + destination.capacity, + layout=destination.layout, + ) + + def fill_address( + self, source: AddressSource, destination: Destination + ) -> Any: + return self._client.fill_address( + source.pointer, + source.capacity, + destination.pointer, + destination.capacity, + destination.shape, + destination.element_bytes, + layout=destination.layout, + ) + + def fill_files( + self, sources: Sequence[FileSource], destination: Destination + ) -> Any: + records = [(source.path, source.offset, source.length) for source in sources] + return self._client.fill_files( + records, + destination.pointer, + destination.capacity, + destination.shape, + destination.element_bytes, + layout=destination.layout, + ) + + +def client_for(device_type: str, *, device_index: int, staging_bytes: int) -> FillClient: + """Bind the one backend implied by the granted destination device.""" + + if device_type == "cpu": + return HostFillClient() + if device_type == "cuda": + return CudaFillClient(staging_bytes, device_index) + raise ValueError(f"tensorfs fill has no destination backend for {device_type!r}") + + +__all__ = [ + "AddressSource", + "Destination", + "FileSource", + "FillClient", + "client_for", +] diff --git a/src/gen_worker/serving/streaming/source.py b/src/gen_worker/serving/streaming/source.py index 88c284a5f..84d78cc38 100644 --- a/src/gen_worker/serving/streaming/source.py +++ b/src/gen_worker/serving/streaming/source.py @@ -2,19 +2,15 @@ from __future__ import annotations -import logging from pathlib import Path from typing import ( Any, Mapping, - Optional, Protocol, Sequence, runtime_checkable, ) -logger = logging.getLogger(__name__) - TENSOR_PLANNERS = frozenset({"safetensors-v1", "gguf-v1"}) TENSOR_SUFFIXES = (".safetensors", ".gguf") @@ -37,16 +33,21 @@ def nbytes(self) -> int: ... @runtime_checkable class TensorStream(Protocol): - """One tensor container, readable at arbitrary offsets into caller memory.""" + """One tensor container consumable by tensorfs's fill path.""" @property def tensors(self) -> Sequence[StreamedTensor]: ... @property def length(self) -> int: ... - def readinto(self, offset: int, length: int, buffer: Any) -> int: - """Copy ``[offset, offset+length)`` into a writable C-contiguous buffer — typically CUDA-pinned host memory, which the store neither knows nor cares about.""" - ... + def fill_host_address( + self, + name: str, + destination_ptr: int, + destination_bytes: int, + destination_offset: int = 0, + layout: str = "torch.contiguous@1", + ) -> Any: ... class WeightStore(Protocol): @@ -131,120 +132,14 @@ def open(self, container: str, *, direct: bool = False) -> TensorStream: return reader -class _BridgeStream: - - def __init__(self, reader: Any, container: str) -> None: - self._reader = reader - self._container = container - views = sorted(reader.values(), key=lambda view: view.offset) - self._tensors: list[StreamedTensor] = list(views) - self._length = max((v.offset + v.nbytes for v in views), default=0) - - @property - def tensors(self) -> Sequence[StreamedTensor]: - return tuple(self._tensors) - - @property - def length(self) -> int: - return self._length - - def readinto(self, offset: int, length: int, buffer: Any) -> int: - target = memoryview(buffer).cast("B") - if len(target) < length: - raise ValueError( - f"{self._container}: buffer holds {len(target)} bytes, " - f"the read needs {length}" - ) - at = 0 - for piece in self._reader._pieces(self._container, offset, length): - target[at : at + len(piece)] = piece - at += len(piece) - return at - - -class BridgeWeightStore: - """The interim byte source over ``_vendor.tensorfs``'s reader.""" - - KIND = "bridge" - - SUFFIXES = TENSOR_SUFFIXES - - def __init__(self, cas: Any, manifest: Any, *, verify: bool = False) -> None: - self._cas = cas - self._manifest = manifest - self._verify = verify - self._open: list[Any] = [] - self._entries = { - entry.path: entry - for entry in manifest.files - if entry.path.endswith(self.SUFFIXES) - } - - def containers(self) -> Sequence[str]: - return tuple(self._entries) - - def open(self, container: str, *, direct: bool = False) -> TensorStream: - from gen_worker._vendor.tensorfs.manifest import RepositoryManifest - from gen_worker._vendor.tensorfs.tensors import TensorReader - - if direct: - raise WeightStoreUnavailable( - "io=direct needs the native reader's O_DIRECT open " - "(tensorfs#115); the interim bridge is buffered only" - ) - entry = self._entries.get(container) - if entry is None: - raise WeightStoreUnavailable( - f"{container!r} is not a tensor container of this manifest" - ) - reader = TensorReader( - self._cas, - RepositoryManifest(files=(entry,)), - verify=self._verify, - ) - self._open.append(reader) - return _BridgeStream(reader, container) - - def close(self) -> None: - for reader in self._open: - reader.close() - self._open.clear() - - -def store_for( - checkpoint_dir: Path | str, *, native: Optional[bool] = None -) -> Optional[WeightStore]: +def store_for(checkpoint_dir: Path | str) -> WeightStore | None: """The byte source backing a projected checkpoint tree, or ``None``.""" from ...models import projection projected = projection.resolve_projection(checkpoint_dir) if projected is None: return None - if native is None: - native = native_available() - if native: - try: - return NativeWeightStore.from_manifest( - projected.cas.root, projected.manifest - ) - except Exception: - logger.warning( - "ctx.load: the native tensorfs store would not open over %s " - "— falling back to the GIL-bound bridge, which is ~10x " - "slower (tensorfs#115)", - projected.cas.root, - exc_info=True, - ) - return BridgeWeightStore(projected.cas, projected.manifest) - - -def native_available() -> bool: - """True when a tensorfs carrying the #115 stream surface is importable.""" - try: - _native_stream_reader() - except (WeightStoreUnavailable, AttributeError): - return False - return True + return NativeWeightStore.from_manifest(projected.cas.root, projected.manifest) def component_of(container: str) -> str: @@ -256,13 +151,11 @@ def component_of(container: str) -> str: __all__ = [ "TENSOR_PLANNERS", "TENSOR_SUFFIXES", - "BridgeWeightStore", "NativeWeightStore", "StreamedTensor", "TensorStream", "WeightStore", "WeightStoreUnavailable", "component_of", - "native_available", "store_for", ] diff --git a/src/gen_worker/serving/streaming/staging.py b/src/gen_worker/serving/streaming/staging.py deleted file mode 100644 index e7c24b7e9..000000000 --- a/src/gen_worker/serving/streaming/staging.py +++ /dev/null @@ -1,140 +0,0 @@ -"""Double-buffered staging: store bytes -> pinned host memory -> device.""" - -from __future__ import annotations - -import logging -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, List, Optional - -if TYPE_CHECKING: # pragma: no cover - typing only - import torch - -logger = logging.getLogger(__name__) - -DEFAULT_BUFFER_BYTES = 64 * 1024 * 1024 -DEFAULT_BUFFERS = 4 - - -class StagingError(RuntimeError): - """The staging pool could not be built or used as declared.""" - - -def _writable_view(tensor: "torch.Tensor") -> memoryview: - import ctypes - - count = tensor.numel() - block = (ctypes.c_ubyte * count).from_address(tensor.data_ptr()) - return memoryview(block).cast("B") - - -@dataclass(slots=True) -class _Slot: - - index: int - tensor: "torch.Tensor" - view: memoryview - event: Optional[Any] = None - - -class StagingPool: - """A ring of host buffers and the stream the copies ride.""" - - def __init__( - self, - device: "torch.device", - *, - buffer_bytes: int = DEFAULT_BUFFER_BYTES, - buffers: int = DEFAULT_BUFFERS, - ) -> None: - import torch - - if buffers < 2: - raise StagingError( - f"a staging pool of {buffers} buffer(s) cannot double-buffer; " - "the read and the copy would serialize" - ) - if buffer_bytes <= 0: - raise StagingError("staging buffers must hold bytes") - - self.device = device - self.buffer_bytes = int(buffer_bytes) - self.pinned = device.type == "cuda" - self._torch = torch - self._stream: Optional[Any] = ( - torch.cuda.Stream(device=device) # type: ignore[no-untyped-call] - if self.pinned - else None - ) - self._slots: List[_Slot] = [] - for index in range(buffers): - host = torch.empty( - self.buffer_bytes, dtype=torch.uint8, pin_memory=self.pinned - ) - self._slots.append( - _Slot(index=index, tensor=host, view=_writable_view(host)) - ) - self._next = 0 - - def acquire(self) -> _Slot: - """The next buffer, once its previous copy has actually landed.""" - slot = self._slots[self._next] - self._next = (self._next + 1) % len(self._slots) - if slot.event is not None: - slot.event.synchronize() - slot.event = None - return slot - - def copy_out( - self, slot: _Slot, src_offset: int, dst: "torch.Tensor", dst_offset: int, count: int - ) -> None: - """Enqueue ``count`` bytes of ``slot`` into a flat uint8 destination.""" - source = slot.tensor[src_offset : src_offset + count] - target = dst[dst_offset : dst_offset + count] - if self._stream is None: - target.copy_(source) - return - with self._torch.cuda.stream(self._stream): - target.copy_(source, non_blocking=True) - - def release(self, slot: _Slot) -> None: - """Mark every copy enqueued out of this buffer, so reuse can wait.""" - if self._stream is None: - return - event = self._torch.cuda.Event() # type: ignore[no-untyped-call] - event.record(self._stream) - slot.event = event - - def track(self, tensor: "torch.Tensor") -> None: - """Tell the caching allocator this tensor is written on the copy stream, not the stream it was allocated on.""" - if self._stream is not None: - tensor.record_stream(self._stream) - - def finish(self) -> None: - """Block until every enqueued copy has landed, then make the compute stream see them.""" - if self._stream is None: - return - self._stream.synchronize() - self._torch.cuda.current_stream(self.device).wait_stream(self._stream) - - def close(self) -> None: - self._slots.clear() - - @property - def staging(self) -> str: - """The telemetry token: ``pinned`` or ``pageable``.""" - return "pinned" if self.pinned else "pageable" - - def __enter__(self) -> "StagingPool": - return self - - def __exit__(self, *_exc: object) -> None: - self.finish() - self.close() - - -__all__ = [ - "DEFAULT_BUFFERS", - "DEFAULT_BUFFER_BYTES", - "StagingError", - "StagingPool", -] diff --git a/tests/streaming_fixture.py b/tests/streaming_fixture.py index a0d9ccb9c..6e2d7777b 100644 --- a/tests/streaming_fixture.py +++ b/tests/streaming_fixture.py @@ -211,13 +211,31 @@ def tensors(self) -> Any: def length(self) -> int: return self._inner.length - def readinto(self, offset: int, length: int, buffer: Any) -> int: - self._log.append((self._container, int(offset), int(length))) - return self._inner.readinto(offset, length, buffer) + @property + def native_reader(self) -> Any: + return self._inner + + def fill_host_address( + self, + name: str, + destination_ptr: int, + destination_bytes: int, + destination_offset: int = 0, + layout: str = "torch.contiguous@1", + ) -> Any: + tensor = next(item for item in self.tensors if item.name == name) + self._log.append((self._container, int(tensor.offset), int(tensor.nbytes))) + return self._inner.fill_host_address( + name, + destination_ptr, + destination_bytes, + destination_offset, + layout, + ) class TracedStore: - """Records every byte range the engine asks the store for.""" + """Records every tensor range the engine asks tensorfs to fill.""" def __init__(self, inner: Any) -> None: self._inner = inner @@ -233,7 +251,7 @@ def open(self, container: str, *, direct: bool = False) -> Any: def assert_file_order(self) -> int: """One forward pass per container, never a seek backwards.""" - assert self.reads, "the engine read nothing" + assert self.reads, "the engine filled nothing" per_container: Dict[str, List[Tuple[int, int]]] = {} for container, offset, length in self.reads: per_container.setdefault(container, []).append((offset, length)) diff --git a/tests/test_checkpoint_juggle.py b/tests/test_checkpoint_juggle.py index 915a67503..9e6fae83f 100644 --- a/tests/test_checkpoint_juggle.py +++ b/tests/test_checkpoint_juggle.py @@ -4,7 +4,9 @@ import json import struct +from contextlib import nullcontext from pathlib import Path +from types import SimpleNamespace from typing import Any, Dict, List, Tuple import pytest @@ -20,6 +22,7 @@ from gen_worker.models.checkpoint_juggle import ( # noqa: E402 CheckpointCatalog, CheckpointImage, + CheckpointJuggler, JuggleRefusal, RegionInvalid, RegionValidity, @@ -181,7 +184,7 @@ def build_image( layout = layout_for(template) manifest = read_manifest(checkpoint_dir(tmp_path, name, module, **kw)) return layout, CheckpointImage( - name, layout, manifest, torch_mod=torch, varena_mod=None, engine=None + name, layout, manifest, torch_mod=torch, varena_mod=None ) @@ -323,3 +326,119 @@ def test_a_partial_switch_is_never_servable_under_either_identity(tmp_path: Path for identity in ("a", "b"): with pytest.raises(RegionInvalid): ledger.assert_servable(identity) + + +class _FillCapture: + staging = "tensorfs-pinned" + + def __init__(self) -> None: + self.addresses: List[Tuple[Any, Any]] = [] + self.files: List[Tuple[Any, Any]] = [] + + def fill_address(self, source: Any, destination: Any) -> Any: + self.addresses.append((source, destination)) + return SimpleNamespace(destination_bytes=destination.capacity) + + def fill_files(self, sources: Any, destination: Any) -> Any: + self.files.append((tuple(sources), destination)) + return SimpleNamespace(destination_bytes=destination.capacity) + + +def _fake_residency(layout: Any, compile_calls: List[int]) -> Any: + def compile_(value: Any) -> Any: + compile_calls.append(1) + return value + + fake_torch = SimpleNamespace( + no_grad=nullcontext, + cuda=SimpleNamespace(synchronize=lambda _device: None), + compile=compile_, + ) + return SimpleNamespace( + adopted=True, + layout=layout, + _torch=fake_torch, + _varena=None, + device=SimpleNamespace(index=0), + reservation=SimpleNamespace(base_ptr=0x40000000), + ring=SimpleNamespace(drain=lambda: None), + is_resident=lambda _name: True, + _host={}, + _triples={}, + _roots=[], + ) + + +def test_warm_swaps_reuse_tensorfs_at_stable_addresses_without_recompile( + tmp_path: Path, monkeypatch: Any +) -> None: + import gen_worker.models.checkpoint_juggle as checkpoint_juggle + + template = Net(seed=0) + mem = {"available": 64 << 30} + layout, catalog = make_catalog(tmp_path, template, mem, floor=4 << 30) + manifests = {} + for seed, name in ((30, "a"), (31, "b")): + manifest = read_manifest(checkpoint_dir(tmp_path, name, Net(seed=seed))) + manifests[name] = manifest + catalog.admit(name, manifest) + assert catalog.ensure_warm(name) is not None + + fills = _FillCapture() + monkeypatch.setattr(checkpoint_juggle, "CudaFillClient", lambda *_args: fills) + compile_calls: List[int] = [] + juggler = CheckpointJuggler( + _fake_residency(layout, compile_calls), + "a", + manifests["a"], + catalog=catalog, + ) + + juggler.switch_to("b") + juggler.switch_to("a") + + region_count = len(layout.regions) + first = [destination.pointer for _, destination in fills.addresses[:region_count]] + second = [destination.pointer for _, destination in fills.addresses[region_count:]] + assert first == second + assert first == [0x40000000 + region.offset for region in layout.regions] + assert compile_calls == [] + assert all(isinstance(source.pointer, int) for source, _ in fills.addresses) + assert all(isinstance(destination.shape, tuple) for _, destination in fills.addresses) + + +def test_cold_swap_is_file_records_through_the_same_tensorfs_client( + tmp_path: Path, monkeypatch: Any +) -> None: + import gen_worker.models.checkpoint_juggle as checkpoint_juggle + + template = Net(seed=0) + layout, catalog = make_catalog( + tmp_path, template, {"available": 0}, floor=4 << 30 + ) + manifests = { + name: read_manifest(checkpoint_dir(tmp_path, name, Net(seed=seed))) + for seed, name in ((40, "a"), (41, "b")) + } + catalog.admit("b", manifests["b"]) + + fills = _FillCapture() + monkeypatch.setattr(checkpoint_juggle, "CudaFillClient", lambda *_args: fills) + juggler = CheckpointJuggler( + _fake_residency(layout, []), + "a", + manifests["a"], + catalog=catalog, + ) + + report = juggler.switch_to("b") + + assert fills.addresses == [] + assert len(fills.files) == len(layout.regions) + for sources, destination in fills.files: + assert sum(source.length for source in sources) == destination.capacity + assert all( + source.path is None or isinstance(source.path, str) for source in sources + ) + assert report.tier == "disk-cold" + assert report.bytes_moved == sum(region.span for region in layout.regions) diff --git a/tests/test_lane_dtype_pgw1623.py b/tests/test_lane_dtype_pgw1623.py index c9fef1d24..5f82c9eb5 100644 --- a/tests/test_lane_dtype_pgw1623.py +++ b/tests/test_lane_dtype_pgw1623.py @@ -36,7 +36,7 @@ from gen_worker._vendor.tensorfs import LocalCAS, project_snapshot # noqa: E402 from gen_worker.models.projection import REF_PREFIX, SNAPSHOTS_DIR # noqa: E402 from gen_worker.serving.streaming import ( # noqa: E402 - BridgeWeightStore, + NativeWeightStore, StreamingLoader, ) from gen_worker.serving.streaming.engine import LaneDtypeUnmet # noqa: E402 @@ -114,8 +114,10 @@ def loaded(tmp_path_factory: pytest.TempPathFactory) -> Any: source = base / "source-model" pipeline_cls = _heterogeneous_source(source) tree = _project(base, source, key="c" * 64) + cas, manifest = _cas_manifest(tree) loader = StreamingLoader( - BridgeWeightStore(*_cas_manifest(tree)), device="cpu", buffer_bytes=4096 + NativeWeightStore.from_manifest(cas.root, manifest), + device="cpu", ) pipeline = loader.build(pipeline_cls, checkpoint_dir=tree, lane=Lane) return {"pipeline": pipeline, "report": loader.last_report, "tree": tree, diff --git a/tests/test_tensorfs_fill_client_pgw1648.py b/tests/test_tensorfs_fill_client_pgw1648.py new file mode 100644 index 000000000..461af3b5d --- /dev/null +++ b/tests/test_tensorfs_fill_client_pgw1648.py @@ -0,0 +1,87 @@ +"""pgw#1648: tensorfs owns bytes; pgw hands it destination data.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, get_args, get_origin, get_type_hints + +from gen_worker.serving.streaming.fill_client import ( + AddressSource, + Destination, + FileSource, + HostFillClient, +) + + +class _Reader: + def __init__(self) -> None: + self.call: tuple[Any, ...] = () + + def fill_host_address(self, *args: Any, **kwargs: Any) -> object: + self.call = (*args, kwargs) + return object() + + +def test_destination_map_crosses_as_plain_data_only() -> None: + destination = Destination( + name="layer.weight", + pointer=1234, + capacity=4096, + source_offset=8192, + shape=(16, 32), + element_bytes=2, + layout="torch.contiguous@1", + ) + annotations = get_type_hints(Destination) + assert set(annotations) == { + "name", + "pointer", + "capacity", + "source_offset", + "shape", + "element_bytes", + "layout", + } + assert all("torch" not in str(annotation).lower() for annotation in annotations.values()) + shape = annotations["shape"] + assert get_origin(shape) in (tuple, None) or tuple in get_args(shape) + + reader = _Reader() + result = HostFillClient().fill(reader, destination) + + assert result is not None + assert reader.call == ( + "layer.weight", + 1234, + 4096, + {"layout": "torch.contiguous@1"}, + ) + assert all(isinstance(value, (str, int, tuple, dict)) for value in reader.call) + + for seam_type in (AddressSource, FileSource): + assert all( + "torch" not in str(annotation).lower() + for annotation in get_type_hints(seam_type).values() + ) + + +def test_the_replaced_pgw_byte_plane_is_deleted() -> None: + root = Path(__file__).parents[1] / "src/gen_worker/serving/streaming" + assert not (root / "staging.py").exists() + sources = "\n".join(path.read_text(encoding="utf-8") for path in sorted(root.glob("*.py"))) + for dead in ( + "BridgeWeightStore", + "StagingPool", + "cudaMemcpyAsync", + "class _Placement", + "def _walk(", + ): + assert dead not in sources + + +def test_the_fill_seam_module_has_no_torch_import_or_type() -> None: + import gen_worker.serving.streaming.fill_client as fill_client + + source = Path(fill_client.__file__).read_text(encoding="utf-8") + assert "import torch" not in source + assert "torch.Tensor" not in source diff --git a/tests/test_weight_streaming.py b/tests/test_weight_streaming.py index f366625c3..82747c5c2 100644 --- a/tests/test_weight_streaming.py +++ b/tests/test_weight_streaming.py @@ -15,7 +15,7 @@ from cas_fixture import ingest_repository # noqa: E402 from gen_worker.models.projection import REF_PREFIX, SNAPSHOTS_DIR # noqa: E402 from gen_worker.serving.streaming import ( # noqa: E402 - BridgeWeightStore, + NativeWeightStore, NameMismatch, StreamingLoader, engine_for, @@ -38,9 +38,6 @@ write_bytes_now, ) -WINDOW = 4096 - - def _project(base: Path, source: Path, key: str) -> Path: cas = LocalCAS(base) manifest = ingest_repository(cas, source) @@ -70,6 +67,11 @@ def _cas_manifest(tree: Path) -> Tuple[Any, Any]: return projected.cas, projected.manifest +def _native_store(tree: Path) -> NativeWeightStore: + cas, manifest = _cas_manifest(tree) + return NativeWeightStore.from_manifest(cas.root, manifest) + + def test_the_fixture_can_actually_witness_a_scrambled_walk( article: dict[str, Any] ) -> None: @@ -95,9 +97,15 @@ def test_ctx_load_streams_store_to_memory_writing_nothing( stubs = [p for p in sorted(tree.rglob("*.safetensors")) if stub_at(p) is not None] assert stubs, f"{tree} projected no pointer stubs — nothing to stream" - store = TracedStore(BridgeWeightStore(*_cas_manifest(tree))) - loader = StreamingLoader(store, device="cpu", buffer_bytes=WINDOW, buffers=3) + store = TracedStore(_native_store(tree)) + loader = StreamingLoader(store, device="cpu") + # Warm library-level config/import caches before the process-I/O arm. The + # measured load below is still a fresh skeleton and a fresh destination + # map; only unrelated one-time interpreter writes are outside the fence. + StreamingLoader(_native_store(tree), device="cpu").build( + pipeline_cls, checkpoint_dir=tree, lane=Lane() + ) before = write_bytes_now() pipeline = loader.build(pipeline_cls, checkpoint_dir=tree, lane=Lane()) written = write_bytes_now() - before @@ -106,7 +114,7 @@ def test_ctx_load_streams_store_to_memory_writing_nothing( report = loader.last_report assert report is not None assert report.weights_streamed_bytes > 0 - assert report.staging == "pageable" + assert report.staging == "destination" assert report.io == "buffered" assert report.containers == 4 @@ -115,14 +123,14 @@ def test_ctx_load_streams_store_to_memory_writing_nothing( for component in ("unet", "vae", "text_encoder", "text_encoder_2"): assert on_meta(getattr(pipeline, component)) == () - windows = store.assert_file_order() - assert windows > 20, ( - f"only {windows} window(s) were read; a walk that fits in one window " - f"is ordered by accident and cannot witness a scrambled one" + fills = store.assert_file_order() + assert fills > 20, ( + f"only {fills} tensor(s) were filled; a one-tensor container " + f"cannot witness a scrambled walk" ) - assert report.windows == windows + assert report.tensors == fills - assert written < 1 << 20, ( + assert written == 0, ( f"the streamed load wrote {written} bytes; the whole point of the " f"2026-08-19 ruling is that it writes none" ) @@ -132,8 +140,8 @@ def test_every_container_is_read_end_to_end_exactly_once( article: dict[str, Any] ) -> None: """The windows tile each container's data range: no gap (a tensor read from nowhere) and no overlap (a byte paid for twice).""" - store = TracedStore(BridgeWeightStore(*_cas_manifest(article["tree"]))) - loader = StreamingLoader(store, device="cpu", buffer_bytes=WINDOW, buffers=3) + store = TracedStore(_native_store(article["tree"])) + loader = StreamingLoader(store, device="cpu") loader.build(article["pipeline_cls"], checkpoint_dir=article["tree"], lane=Lane()) per_container: dict[str, list[tuple[int, int]]] = {}