diff --git a/daser/connector/daser_connector.py b/daser/connector/daser_connector.py index 46dc178..659780b 100644 --- a/daser/connector/daser_connector.py +++ b/daser/connector/daser_connector.py @@ -2,7 +2,6 @@ # Standard import asyncio -from concurrent.futures import ThreadPoolExecutor import threading from typing import TYPE_CHECKING, Any @@ -40,7 +39,7 @@ from daser.connector.staging import ( copy_staging_to_kv_cache as _copy_staging_to_kv_cache, ) -from daser.connector.worker import WorkerConnectorMixin +from daser.connector.worker import _LOAD_REQUEST_MAX_INFLIGHT, WorkerConnectorMixin from daser.logging import init_logger logger = init_logger(__name__) @@ -124,6 +123,13 @@ def __init__( self._transfer_ready = False self._transfer_mode = str(extra.get("transfer_mode", "iouring")) self._ipc_load_async = IPCClientAsync(self._socket_path) + self._ipc_load_async_pool = [ + self._ipc_load_async, + *[ + IPCClientAsync(self._socket_path) + for _ in range(max(0, _LOAD_REQUEST_MAX_INFLIGHT - 1)) + ], + ] self._ipc_store_async = IPCClientAsync(self._socket_path) self._kv_caches: dict[str, torch.Tensor] = {} self._layer_names: list[str] = [] @@ -133,15 +139,15 @@ def __init__( self._pending_save_staging_bytes = 0 self._store_staging_bytes = 0 self._pending_store_staging_limit_bytes = 0 - self._staging_pool = None + self._store_staging_pool = None self._pending_commits: set[str] = set() self._pending_finished_saves: dict[str, Any] = {} self._pending_loads: dict[str, Any] = {} self._invalid_load_block_ids: set[int] = set() - self._load_executor = ThreadPoolExecutor( - max_workers=1, - thread_name_prefix="daser-load", - ) + self._load_request_queue = None + self._load_request_dispatcher = None + self._load_request_queue_lock = threading.Lock() + self._load_request_dispatcher_future = None self._load_loop = asyncio.new_event_loop() self._store_loop = asyncio.new_event_loop() self._load_thread = threading.Thread( diff --git a/daser/connector/ipc_client.py b/daser/connector/ipc_client.py index 0e514b9..15bd404 100644 --- a/daser/connector/ipc_client.py +++ b/daser/connector/ipc_client.py @@ -547,3 +547,81 @@ async def transfer_load_cuda( "spans": spans, } ) + + async def register_load_staging_cuda( + self, + buffer_index: int, + cuda_ipc_handle: bytes, + allocation_bytes: int, + device_id: int, + device_ptr: int, + allocation_base_ptr: int, + allocation_offset: int, + producer_pid: int, + ) -> None: + """Register one fixed load staging CUDA allocation with the server. + + Args: + buffer_index: Worker-local fixed staging buffer index. + cuda_ipc_handle: exported CUDA IPC memory handle. + allocation_bytes: byte size of the CUDA allocation to map. + device_id: CUDA device ordinal for the exported allocation. + device_ptr: raw device pointer for same-process server harnesses. + allocation_base_ptr: base pointer of the CUDA allocation owning + ``device_ptr``. + allocation_offset: byte offset of ``device_ptr`` from + ``allocation_base_ptr``. + producer_pid: process ID that exported the pointer. + + Async/thread-safety: + Serializes with other calls on the dedicated async client + connection. Intended for worker initialization before hot-path + cache-hit loads. + """ + await self.call( + { + "op": "register_load_staging", + "payload": { + "buffer_index": int(buffer_index), + "cuda_ipc_handle": cuda_ipc_handle, + "allocation_bytes": int(allocation_bytes), + "device_id": int(device_id), + "device_ptr": int(device_ptr), + "allocation_base_ptr": int(allocation_base_ptr), + "allocation_offset": int(allocation_offset), + "producer_pid": int(producer_pid), + }, + } + ) + + async def transfer_load_registered_cuda( + self, + buffer_index: int, + nbytes: int, + spans: list[dict[str, int]], + ) -> dict[str, Any]: + """Load into a previously registered fixed CUDA staging buffer. + + Args: + buffer_index: Worker-local fixed staging buffer index registered + through ``register_load_staging_cuda``. + nbytes: logical bytes to write for this transfer. + spans: byte spans containing target_offset, nbytes, and file_offset. + + Returns: + Server response including transferred bytes and timing counters. + + Async/thread-safety: + Runs on the worker load event loop and avoids per-load CUDA IPC + handle export/open payloads on the hot path. + """ + return await self.call( + { + "op": "transfer_load", + "payload": { + "load_staging_buffer_index": int(buffer_index), + "nbytes": int(nbytes), + }, + "spans": spans, + } + ) diff --git a/daser/connector/staging.py b/daser/connector/staging.py index cf664bd..4394394 100644 --- a/daser/connector/staging.py +++ b/daser/connector/staging.py @@ -3,7 +3,9 @@ from __future__ import annotations # Standard +from collections.abc import Callable from dataclasses import dataclass, replace +from typing import Any, Protocol # Third Party import torch @@ -36,9 +38,25 @@ ] = {} +class _CudaStagingLeaseOwner(Protocol): + """Pool protocol required by ``CudaStagingLease``. + + Async/thread-safety: + Implementations define their own ownership model; release is invoked + by the worker thread after CUDA IPC users are finished with the lease. + """ + + def release(self, lease: "CudaStagingLease") -> None: + """Return a lease to its owning pool. + + Args: + lease: Lease previously returned by that pool. + """ + + @dataclass class CudaStagingLease: - """One logical staging allocation leased from ``CudaStagingPool``. + """One logical staging allocation leased from a worker-side staging pool. Args: pool: Owning pool that will receive the allocation on release. @@ -50,7 +68,7 @@ class CudaStagingLease: completion. It must not be reused while an async store future owns it. """ - pool: "CudaStagingPool" + pool: _CudaStagingLeaseOwner tensor: torch.Tensor nbytes: int _released: bool = False @@ -79,91 +97,154 @@ def release(self) -> None: self.pool.release(self) -class CudaStagingPool: - """Reusable worker-side GPU staging buffers for CUDA IPC transfer. +class FixedCudaStagingPool: + """Fixed-size worker-side CUDA staging buffers. Args: device: Device on which staging tensors are allocated. - initial_bytes: Size of the buffer allocated during initialization. - max_buffer_bytes: Maximum size of one staging buffer. + buffer_bytes: Size of each fixed staging buffer. + depth: Number of fixed buffers to allocate. Async/thread-safety: - The pool is owned by one vLLM worker thread. It does not use locks; the - worker tracks async store futures and releases leases only after those - futures complete. + The pool is owned by one vLLM worker thread. It never allocates after + construction, so load and store staging cannot grow into unexpected + CUDA OOM. Callers that need blocking semantics can pass a release + callback to ``acquire`` when all buffers are currently leased. """ def __init__( self, device: torch.device, - initial_bytes: int, - max_buffer_bytes: int, + buffer_bytes: int, + depth: int, ) -> None: - if initial_bytes < 0: - raise ValueError("initial_bytes must be non-negative") - if max_buffer_bytes <= 0: - raise ValueError("max_buffer_bytes must be positive") + if buffer_bytes <= 0: + raise ValueError("buffer_bytes must be positive") + if depth <= 0: + raise ValueError("depth must be positive") self._device = device - self._max_buffer_bytes = max_buffer_bytes - self._free: list[torch.Tensor] = [] - if initial_bytes: - self._free.append( - torch.empty(initial_bytes, dtype=torch.uint8, device=device) - ) + self._buffer_bytes = buffer_bytes + self._buffers: list[torch.Tensor] = [ + torch.empty(buffer_bytes, dtype=torch.uint8, device=device) + for _ in range(depth) + ] + self._free_indices: list[int] = list(range(depth)) + + @property + def buffer_bytes(self) -> int: + """Return the fixed capacity of each staging buffer.""" + return self._buffer_bytes @property - def max_buffer_bytes(self) -> int: - """Return the maximum bytes allowed for a single staging buffer.""" - return self._max_buffer_bytes + def available(self) -> int: + """Return the number of currently free staging buffers.""" + return len(self._free_indices) - def acquire(self, nbytes: int) -> CudaStagingLease: - """Lease a staging tensor with at least ``nbytes`` capacity. + @property + def depth(self) -> int: + """Return the number of fixed staging buffers in this pool.""" + return len(self._buffers) + + def buffer(self, index: int) -> torch.Tensor: + """Return a fixed staging backing tensor by index. + + Args: + index: Fixed buffer index to inspect. + + Returns: + The full preallocated backing tensor for the requested index. + + Async/thread-safety: + Intended for one-time CUDA IPC registration before load traffic. + Callers must not mutate the returned tensor independently of the + pool lease protocol. + """ + if index < 0 or index >= len(self._buffers): + raise ValueError(f"fixed staging buffer index out of range: {index}") + return self._buffers[index] + + def acquire( + self, + nbytes: int, + wait_for_release: Callable[[], None] | None = None, + ) -> CudaStagingLease: + """Lease one preallocated staging buffer. Args: nbytes: Logical transfer byte count. + wait_for_release: Optional callback invoked once when no buffer is + currently free. Store callers use this to wait for an older + asynchronous store to finish and release its lease. Returns: - Lease whose ``view`` is sized to ``nbytes``. + Lease whose view is limited to ``nbytes``. Raises: - ValueError: If ``nbytes`` exceeds the configured single-buffer cap. + ValueError: If ``nbytes`` exceeds the fixed buffer size. + RuntimeError: If all fixed buffers are currently in use. + """ + if nbytes < 0: + raise ValueError("nbytes must be non-negative") + if nbytes > self._buffer_bytes: + raise ValueError( + f"staging request {nbytes} exceeds fixed staging buffer " + f"{self._buffer_bytes}" + ) + if not self._free_indices and wait_for_release is not None: + wait_for_release() + if not self._free_indices: + raise RuntimeError("no fixed staging buffers available") + return self.acquire_index(self._free_indices[0], nbytes) - Async/thread-safety: - Synchronous worker-thread allocation path. Reuses preallocated - buffers first; only grows when the current workload exceeds the - initialization size. + def acquire_index(self, index: int, nbytes: int) -> CudaStagingLease: + """Lease a specific preallocated staging buffer. + + Args: + index: Fixed buffer index to lease. + nbytes: Logical transfer byte count. + + Returns: + Lease whose view is limited to ``nbytes``. + + Raises: + ValueError: If ``index`` is invalid or ``nbytes`` exceeds the fixed + buffer size. + RuntimeError: If the requested buffer is currently in use. """ + if index < 0 or index >= len(self._buffers): + raise ValueError(f"fixed staging buffer index out of range: {index}") if nbytes < 0: raise ValueError("nbytes must be non-negative") - if nbytes > self._max_buffer_bytes: + if nbytes > self._buffer_bytes: raise ValueError( - f"staging request {nbytes} exceeds cap {self._max_buffer_bytes}" + f"staging request {nbytes} exceeds fixed staging buffer " + f"{self._buffer_bytes}" ) - selected_idx = -1 - selected_size = 0 - for idx, candidate in enumerate(self._free): - capacity = int(candidate.numel()) - if capacity >= nbytes and (selected_idx < 0 or capacity < selected_size): - selected_idx = idx - selected_size = capacity - if selected_idx >= 0: - tensor = self._free.pop(selected_idx) - else: - capacity = max(nbytes, min(self._max_buffer_bytes, nbytes)) - tensor = torch.empty(capacity, dtype=torch.uint8, device=self._device) - return CudaStagingLease(pool=self, tensor=tensor, nbytes=nbytes) + if index not in self._free_indices: + raise RuntimeError(f"fixed staging buffer {index} is not available") + self._free_indices.remove(index) + return CudaStagingLease( + pool=self, + tensor=self._buffers[index], + nbytes=nbytes, + ) def release(self, lease: CudaStagingLease) -> None: - """Return a completed staging lease to the free list. + """Return a fixed staging lease to the free list. Args: lease: Lease previously returned by ``acquire``. - - Async/thread-safety: - Caller must ensure no CUDA operation or server CUDA IPC mapping is - still using the tensor. """ - self._free.append(lease.tensor) + for index, tensor in enumerate(self._buffers): + if tensor is lease.tensor: + if index not in self._free_indices: + self._free_indices.append(index) + self._free_indices.sort() + return + raise ValueError("lease does not belong to this fixed staging pool") + + +StoreCudaStagingPool = FixedCudaStagingPool @dataclass(frozen=True) @@ -824,12 +905,15 @@ def copy_cross_layer_kv_cache_to_staging( def build_load_read_plan( reqs_to_load: dict[str, ReqLoadSpec], slot_size: int, -) -> tuple[int, list[dict[str, int]], list[tuple[int, int, ReqLoadSpec]]]: + include_req_ids: bool = False, +) -> tuple[int, list[dict[str, int]], list[Any]]: """Build a combined transfer-load plan for one forward step. Args: reqs_to_load: request ID to load spec from scheduler metadata. slot_size: bytes per vLLM KV slot. + include_req_ids: when True, include request IDs in per-request ranges + for worker-side request completion tracking. Returns: ``(total_bytes, spans, per_req_ranges)`` where spans target one @@ -838,9 +922,9 @@ def build_load_read_plan( """ total_bytes = 0 spans: list[dict[str, int]] = [] - per_req_ranges: list[tuple[int, int, ReqLoadSpec]] = [] + per_req_ranges: list[Any] = [] source_ranges: dict[tuple[str, int, int, int, int], tuple[int, int]] = {} - for spec in reqs_to_load.values(): + for req_id, spec in reqs_to_load.items(): num_slots = len(spec.block_ids) if num_slots == 0: continue @@ -861,7 +945,10 @@ def build_load_read_plan( total_bytes = end else: start, end = existing - per_req_ranges.append((start, end, spec)) + if include_req_ids: + per_req_ranges.append((start, end, req_id, spec)) + else: + per_req_ranges.append((start, end, spec)) return total_bytes, spans, per_req_ranges @@ -869,13 +956,16 @@ def build_load_read_batches( reqs_to_load: dict[str, ReqLoadSpec], slot_size: int, max_batch_bytes: int, -) -> list[tuple[int, list[dict[str, int]], list[tuple[int, int, ReqLoadSpec]]]]: + include_req_ids: bool = False, +) -> list[tuple[int, list[dict[str, int]], list[Any]]]: """Build bounded load staging plans for one forward step. Args: reqs_to_load: request ID to load spec from scheduler metadata. slot_size: bytes per vLLM KV slot. max_batch_bytes: Maximum staging bytes for one transfer batch. + include_req_ids: when True, include request IDs in per-request ranges + for worker-side request completion tracking. Returns: List of ``build_load_read_plan``-style tuples. Individual requests are @@ -889,9 +979,7 @@ def build_load_read_batches( if max_batch_bytes <= 0: raise ValueError("max_batch_bytes must be positive") max_slots = max(1, max_batch_bytes // slot_size) - batches: list[ - tuple[int, list[dict[str, int]], list[tuple[int, int, ReqLoadSpec]]] - ] = [] + batches: list[tuple[int, list[dict[str, int]], list[Any]]] = [] current: dict[str, ReqLoadSpec] = {} current_slots = 0 synthetic_id = 0 @@ -899,7 +987,13 @@ def build_load_read_batches( def flush() -> None: nonlocal current, current_slots if current: - batches.append(build_load_read_plan(current, slot_size)) + batches.append( + build_load_read_plan( + current, + slot_size, + include_req_ids=include_req_ids, + ) + ) current = {} current_slots = 0 @@ -971,7 +1065,11 @@ def flush() -> None: run_pos_offset = 0 run_block_ids = [] - for start, end, spec in per_req_ranges: + for item in per_req_ranges: + if len(item) == 3: + start, end, spec = item + else: + start, end, _req_id, spec = item if not spec.block_ids: continue if run_start >= 0 and start == run_end and spec.pos_offset == run_pos_offset: diff --git a/daser/connector/worker.py b/daser/connector/worker.py index 6defaca..42e420a 100644 --- a/daser/connector/worker.py +++ b/daser/connector/worker.py @@ -4,8 +4,14 @@ # Standard import asyncio +from collections import deque +from collections.abc import Callable, Sequence +from concurrent.futures import FIRST_COMPLETED +from concurrent.futures import wait as wait_futures from dataclasses import dataclass import os +import queue +import threading import time from typing import TYPE_CHECKING, Any @@ -32,8 +38,9 @@ DEFAULT_STORE_STAGING_BYTES, FUSED_RESTORE_MIN_SLOTS, CudaStagingLease, - CudaStagingPool, + FixedCudaStagingPool, StagedStoreBatch, + StoreCudaStagingPool, ) from daser.connector.staging import ( build_load_copy_runs as _build_load_copy_runs, @@ -79,6 +86,173 @@ logger = init_logger(__name__) _ROPE_WARMUP_BLOCKS = 1 +_LOAD_PIPELINE_DEPTH = 2 +_LOAD_REQUEST_MAX_INFLIGHT = 8 +_LOAD_DISPATCH_WAIT_TIMEOUT_S = 0.001 +_LOAD_STAGING_RESERVE_BYTES = 1 << 30 +_MIN_STORE_STAGING_POOL_DEPTH = 1 +_LoadBatch = tuple[int, list[dict[str, int]], list[Any]] + + +def _run_fixed_depth_pipeline( + batches: Sequence[Any], + submit: Callable[[Any, int], Any], + is_complete: Callable[[Any], bool], + consume: Callable[[Any], int], + buffer_indices: Sequence[int], + wait_for_completion: Callable[[Sequence[Any]], Any] | None = None, +) -> int: + """Run a fixed-depth completion-driven submit/consume pipeline. + + Args: + batches: Batch payloads to submit. + submit: Function called with ``(batch, buffer_index)``. It returns an + opaque in-flight state. + is_complete: Predicate that returns True when an in-flight state can be + consumed without blocking on its load future. + consume: Function called with an in-flight state. It must return the + reusable buffer index. + buffer_indices: Fixed buffer identifiers available to the pipeline. + wait_for_completion: Optional blocking wait that returns one completed + in-flight state when none are immediately complete. + + Returns: + Maximum number of in-flight batches observed. + + Async/thread-safety: + Pure synchronous helper. Callers own any synchronization required to + wait for completion and to consume a completed state safely. + """ + if not buffer_indices: + raise ValueError("buffer_indices must not be empty") + next_batch = 0 + max_inflight = 0 + in_flight: deque[Any] = deque() + + while next_batch < len(batches) and len(in_flight) < len(buffer_indices): + buffer_index = buffer_indices[len(in_flight)] + in_flight.append(submit(batches[next_batch], buffer_index)) + next_batch += 1 + max_inflight = max(max_inflight, len(in_flight)) + + while in_flight: + ready_state = next((state for state in in_flight if is_complete(state)), None) + if ready_state is None: + ready_state = ( + wait_for_completion(tuple(in_flight)) + if wait_for_completion is not None + else in_flight[0] + ) + in_flight.remove(ready_state) + state = ready_state + reusable_buffer = consume(state) + if next_batch < len(batches): + in_flight.append(submit(batches[next_batch], reusable_buffer)) + next_batch += 1 + max_inflight = max(max_inflight, len(in_flight)) + + return max_inflight + + +def _run_depth_two_pipeline( + batches: Sequence[Any], + submit: Callable[[Any, int], Any], + consume: Callable[[Any], int], + buffer_indices: Sequence[int], +) -> int: + """Run a fixed-depth pipeline, consuming states immediately. + + Args: + batches: Batch payloads to submit. + submit: Function called with ``(batch, buffer_index)``. It returns an + opaque in-flight state. + consume: Function called with an in-flight state. It must return the + reusable buffer index. + buffer_indices: Fixed buffer identifiers available to the pipeline. + + Returns: + Maximum number of in-flight batches observed. + + Async/thread-safety: + Compatibility wrapper for tests and callers whose states are always + ready when selected. + """ + return _run_fixed_depth_pipeline( + batches=batches, + submit=submit, + is_complete=lambda _state: True, + consume=consume, + buffer_indices=buffer_indices, + ) + + +def _store_staging_pool_depth(buffer_bytes: int, pending_limit_bytes: int) -> int: + """Return fixed store staging pool depth for the configured byte budget. + + Args: + buffer_bytes: Capacity of one fixed staging buffer. + pending_limit_bytes: Total pending store staging byte budget. + + Returns: + Number of fixed store staging buffers to preallocate. + + Async/thread-safety: + Pure helper used during worker-side pool initialization. + """ + if buffer_bytes <= 0: + raise ValueError("buffer_bytes must be positive") + if pending_limit_bytes <= 0: + return _MIN_STORE_STAGING_POOL_DEPTH + return max(_MIN_STORE_STAGING_POOL_DEPTH, pending_limit_bytes // buffer_bytes) + + +def _load_staging_pool_depth( + buffer_bytes: int, + pending_limit_bytes: int, + device: torch.device, +) -> int: + """Return fixed load staging depth under memory and inflight constraints. + + Args: + buffer_bytes: Capacity of one fixed staging buffer. + pending_limit_bytes: Existing staging byte budget from store limits. + device: CUDA device used for staging allocation. + + Returns: + Number of fixed load staging buffers to preallocate. + + Async/thread-safety: + Pure helper except for querying CUDA free memory. Called during worker + initialization before request traffic. + """ + depth = min( + _LOAD_REQUEST_MAX_INFLIGHT, + _store_staging_pool_depth(buffer_bytes, pending_limit_bytes), + ) + if device.type != "cuda": + return max(1, depth) + try: + free_bytes, _total_bytes = torch.cuda.mem_get_info(device) + except Exception: # noqa: BLE001 + return max(1, depth) + usable_bytes = max(0, int(free_bytes) - _LOAD_STAGING_RESERVE_BYTES) + memory_depth = max(1, usable_bytes // buffer_bytes) + return max(1, min(depth, memory_depth)) + + +def _load_completion_req_id(req_id: str) -> str: + """Return the vLLM request ID that owns a worker-side load range. + + Args: + req_id: Scheduler request ID or staging-local split ID. + + Returns: + Base vLLM request ID used for request-level load completion. + + Async/thread-safety: + Pure string helper. Safe to call from worker and executor threads. + """ + return base_req_id(req_id.split("#", 1)[0]) @dataclass @@ -133,6 +307,69 @@ def release(self) -> None: self.lease = None +@dataclass +class _InflightLoadBatch: + """One submitted load batch and its fixed staging lease. + + Attributes: + buffer_index: Fixed load staging buffer index. + total_bytes: Logical bytes in this transfer batch. + per_req_ranges: Ranges used to restore staging bytes into vLLM KV cache. + staging_lease: Fixed staging lease retained until restore completes. + future: Future returned by the load event loop. + submitted_at: Wall-clock timestamp used for wait timing. + """ + + buffer_index: int + total_bytes: int + per_req_ranges: list[Any] + staging_lease: CudaStagingLease + future: Any + submitted_at: float + + +@dataclass +class _InflightRequestLoad: + """One request-level load submitted by the dispatcher. + + Attributes: + item: Original queued request item. + buffer_index: Fixed staging buffer owned by this request state. + batches: Bounded read batches for this request. + next_batch: Index of the next unsubmitted read batch. + remaining_batches: Number of read batches not yet consumed. + active: Current in-flight read batch for this request, if any. + completed: Consumed batch timing rows for logging. + """ + + item: "_QueuedLoadRequest" + buffer_index: int + batches: list[_LoadBatch] + next_batch: int + remaining_batches: int + active: _InflightLoadBatch | None + completed: list["_ConsumedLoadBatch"] + + +@dataclass +class _ConsumedLoadBatch: + """Timing and accounting data for one restored load batch.""" + + buffer_index: int + bytes: int + copies: int + copy_runs: int + ipc_ms: float + wait_ms: float + copy_ms: float + transfer_open_ms: float + transfer_load_ms: float + transfer_sync_ms: float + l1_hits: int + l1_misses: int + l2_reads: int + + class _ImmediateLoadError: """Completed future used when a load cannot be submitted. @@ -161,6 +398,159 @@ def result(self, timeout: float | None = None) -> None: raise RuntimeError(self._message) +class _RequestLoadFuture: + """Small future used to release request loads independently. + + Async/thread-safety: + The connector load executor marks completion and the vLLM worker thread + polls ``done``/``result`` from ``get_finished``. ``threading.Event`` + provides cross-thread visibility for the result state. + """ + + def __init__(self) -> None: + self._event = threading.Event() + self._error: BaseException | None = None + + def done(self) -> bool: + """Return whether this request load has completed.""" + return self._event.is_set() + + def result(self, timeout: float | None = None) -> None: + """Wait for completion and raise a stored load error, if any. + + Args: + timeout: Optional timeout in seconds. + """ + if not self._event.wait(timeout): + raise TimeoutError("request load did not complete before timeout") + if self._error is not None: + raise self._error + + def set_result(self) -> None: + """Mark this request load as successful.""" + self._event.set() + + def set_exception(self, error: BaseException) -> None: + """Mark this request load as failed. + + Args: + error: Exception to re-raise from ``result``. + """ + self._error = error + self._event.set() + + +@dataclass +class _QueuedLoadRequest: + """One request waiting in the worker-side load queue. + + Attributes: + req_id: Base vLLM request ID used for completion. + spec_id: Scheduler metadata ID for this load spec. + spec: Load spec to restore into the vLLM KV cache. + future: Per-request completion future observed by ``get_finished``. + """ + + req_id: str + spec_id: str + spec: Any + future: _RequestLoadFuture + + +class LoadRequestDispatcher: + """Bound request-level load concurrency by max in-flight and staging depth. + + Args: + max_inflight: Maximum request loads allowed in flight. + staging_depth: Number of fixed staging buffers available to request + loads. + + Async/thread-safety: + The dispatcher object is owned by the connector load asyncio loop. Test + helpers may call its pure synchronous scheduling helpers directly. + """ + + def __init__( + self, + max_inflight: int, + staging_depth: int, + ) -> None: + self._effective_inflight = max(1, min(max_inflight, staging_depth)) + self._free_buffers: deque[int] = deque(range(self._effective_inflight)) + + @property + def effective_inflight(self) -> int: + """Return the active request limit after applying staging depth.""" + return self._effective_inflight + + def submit_ready( + self, + connector: Any, + queued: list[_QueuedLoadRequest], + sample_tensor: torch.Tensor, + ) -> list[_InflightRequestLoad]: + """Submit queued requests while in-flight slots and buffers are free. + + Args: + connector: Worker connector that owns submit helpers. + queued: Mutable FIFO list of queued request work. + sample_tensor: Representative KV cache tensor. + + Returns: + Newly submitted request states. + """ + submitted: list[_InflightRequestLoad] = [] + while queued and self._free_buffers: + buffer_index = self._free_buffers.popleft() + item = queued.pop(0) + state = connector._submit_request_load_for_dispatcher( # noqa: SLF001 + item, + buffer_index, + sample_tensor, + ) + if state.active is None: + self._free_buffers.append(state.buffer_index) + continue + submitted.append(state) + return submitted + + def consume_ready( + self, + connector: Any, + active: list[_InflightRequestLoad], + sample_tensor: torch.Tensor, + ) -> list[_InflightRequestLoad]: + """Consume completed request loads and release their buffers. + + Args: + connector: Worker connector that owns consume helpers. + active: Mutable list of active request load states. + sample_tensor: Representative KV cache tensor. + + Returns: + States consumed during this call. + """ + consumed: list[_InflightRequestLoad] = [] + for state in list(active): + active_batch = state.active + if active_batch is None: + active.remove(state) + self._free_buffers.append(state.buffer_index) + consumed.append(state) + continue + if not active_batch.future.done(): + continue + reusable_buffer, request_done = connector._consume_dispatcher_load( # noqa: SLF001 + state, + sample_tensor, + ) + if request_done: + active.remove(state) + self._free_buffers.append(reusable_buffer) + consumed.append(state) + return consumed + + def _cuda_allocation_base_and_offset(device_ptr: int) -> tuple[int, int]: """Return CUDA allocation base pointer and byte offset for ``device_ptr``. @@ -354,16 +744,32 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES, self._slot_size, ) - self._staging_pool = CudaStagingPool( + self._store_staging_pool = StoreCudaStagingPool( + device=sample.device, + buffer_bytes=self._store_staging_bytes, + depth=_store_staging_pool_depth( + self._store_staging_bytes, + self._pending_store_staging_limit_bytes, + ), + ) + self._load_staging_pool = FixedCudaStagingPool( device=sample.device, - initial_bytes=self._store_staging_bytes, - max_buffer_bytes=self._store_staging_bytes, + buffer_bytes=self._store_staging_bytes, + depth=_load_staging_pool_depth( + self._store_staging_bytes, + self._pending_store_staging_limit_bytes, + sample.device, + ), ) + self._load_staging_registered = False logger.info( - "[CONNECTOR] preallocated staging buffer=%d cap=%d pending=%d", + "[CONNECTOR] preallocated staging buffer=%d cap=%d pending=%d " + "load_request_max_inflight=%d load_staging_depth=%d", self._store_staging_bytes, self._store_staging_bytes, self._pending_store_staging_limit_bytes, + _LOAD_REQUEST_MAX_INFLIGHT, + self._load_staging_pool.depth, ) if sample.dim() >= 5: _warm_rope_apply_backends( @@ -428,16 +834,32 @@ def register_cross_layers_kv_cache( self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES, self._slot_size, ) - self._staging_pool = CudaStagingPool( + self._store_staging_pool = StoreCudaStagingPool( + device=kv_cache.device, + buffer_bytes=self._store_staging_bytes, + depth=_store_staging_pool_depth( + self._store_staging_bytes, + self._pending_store_staging_limit_bytes, + ), + ) + self._load_staging_pool = FixedCudaStagingPool( device=kv_cache.device, - initial_bytes=self._store_staging_bytes, - max_buffer_bytes=self._store_staging_bytes, + buffer_bytes=self._store_staging_bytes, + depth=_load_staging_pool_depth( + self._store_staging_bytes, + self._pending_store_staging_limit_bytes, + kv_cache.device, + ), ) + self._load_staging_registered = False logger.info( - "[CONNECTOR] register_cross_layers_kv_cache: layers=%d shape=%s dtype=%s", + "[CONNECTOR] register_cross_layers_kv_cache: layers=%d shape=%s " + "dtype=%s load_request_max_inflight=%d load_staging_depth=%d", len(self._layer_names), tuple(kv_cache.shape), kv_cache.dtype, + _LOAD_REQUEST_MAX_INFLIGHT, + self._load_staging_pool.depth, ) _warm_rope_apply_backends( device=kv_cache.device, @@ -512,25 +934,29 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non self._mark_load_start_failed(reqs_to_load, "no registered KV cache tensor") return - block_ids = [ - block_id for spec in reqs_to_load.values() for block_id in spec.block_ids - ] - base_req_ids = {base_req_id(req_id) for req_id in reqs_to_load} - future = self._load_executor.submit( - self._load_kv_specs, - reqs_to_load, - sample_tensor, - ) pending_loads = getattr(self, "_pending_loads", None) if pending_loads is None: pending_loads = {} self._pending_loads = pending_loads - for req_id in base_req_ids: + load_queue = self._ensure_load_request_queue(sample_tensor) + for spec_id, spec in reqs_to_load.items(): + req_id = base_req_id(spec_id) + request_future = _RequestLoadFuture() pending_loads[req_id] = _PendingLoad( - future=future, - block_ids=list(block_ids), + future=request_future, + block_ids=list(spec.block_ids), lease=None, ) + self._enqueue_load_request( + load_queue, + _QueuedLoadRequest( + req_id=req_id, + spec_id=spec_id, + spec=spec, + future=request_future, + ), + ) + self._ensure_load_request_dispatcher(sample_tensor) def _mark_load_start_failed( self, @@ -564,10 +990,184 @@ def _mark_load_start_failed( lease=None, ) + def _ensure_load_request_queue( + self, sample_tensor: torch.Tensor | None = None + ) -> Any: + """Return the worker-side request load queue, creating it if needed. + + Returns: + Queue used to group request loads across scheduler windows. + + Async/thread-safety: + Called on the vLLM worker thread. Attribute installation is guarded + by a small lock because multiple model-runner calls may race during + startup. + """ + load_queue = getattr(self, "_load_request_queue", None) + if load_queue is not None: + return load_queue + lock = getattr(self, "_load_request_queue_lock", None) + if lock is None: + lock = threading.Lock() + self._load_request_queue_lock = lock + with lock: + load_queue = getattr(self, "_load_request_queue", None) + if load_queue is None: + if sample_tensor is not None and hasattr(self, "_load_loop"): + future = asyncio.run_coroutine_threadsafe( + self._create_load_request_queue(), + self._load_loop, + ) + load_queue = future.result(timeout=10.0) + else: + load_queue = queue.Queue() + self._load_request_queue = load_queue + return load_queue + + async def _create_load_request_queue(self) -> asyncio.Queue[Any]: + """Create an asyncio request queue on the load event loop.""" + return asyncio.Queue() + + def _enqueue_load_request( + self, + load_queue: Any, + request: _QueuedLoadRequest, + ) -> None: + """Enqueue one request from a worker thread into the load queue. + + Args: + load_queue: Queue created by ``_ensure_load_request_queue``. + request: Request-level load work item. + + Async/thread-safety: + Called on vLLM worker threads. Production queues are asyncio queues + owned by ``_load_loop``; test queues may be synchronous ``queue.Queue`` + instances. + """ + if isinstance(load_queue, asyncio.Queue): + self._load_loop.call_soon_threadsafe(load_queue.put_nowait, request) + else: + load_queue.put(request) + + def _ensure_load_request_dispatcher(self, sample_tensor: torch.Tensor) -> None: + """Start the persistent request load dispatcher if it is not running. + + Args: + sample_tensor: Representative KV cache tensor used by load workers. + + Async/thread-safety: + Called on the vLLM worker thread. The dispatcher itself runs on the + connector load asyncio loop and keeps request-level loads in flight + up to the fixed max-inflight/staging-pool limit. + """ + dispatcher_future = getattr(self, "_load_request_dispatcher_future", None) + if dispatcher_future is not None and not dispatcher_future.done(): + return + lock = getattr(self, "_load_request_queue_lock", None) + if lock is None: + lock = threading.Lock() + self._load_request_queue_lock = lock + with lock: + dispatcher_future = getattr(self, "_load_request_dispatcher_future", None) + if dispatcher_future is not None and not dispatcher_future.done(): + return + self._load_request_dispatcher_future = asyncio.run_coroutine_threadsafe( + self._run_load_request_dispatcher(sample_tensor), + self._load_loop, + ) + + async def _run_load_request_dispatcher(self, sample_tensor: torch.Tensor) -> None: + """Continuously dispatch request loads as in-flight slots become free. + + Args: + sample_tensor: Representative KV cache tensor used for load staging. + + Async/thread-safety: + Runs on the connector load asyncio loop. The dispatcher never waits + to coalesce requests; it submits each queued request immediately when + both an in-flight slot and a fixed staging buffer are available. + """ + load_queue = self._ensure_load_request_queue(sample_tensor) + load_staging_pool = self._ensure_load_staging_pool(sample_tensor) + dispatcher = LoadRequestDispatcher( + max_inflight=_LOAD_REQUEST_MAX_INFLIGHT, + staging_depth=load_staging_pool.depth, + ) + self._load_request_dispatcher = dispatcher + queued: list[_QueuedLoadRequest] = [] + active: list[_InflightRequestLoad] = [] + try: + while True: + if not queued: + if active: + await self._drain_load_queue(load_queue, queued) + else: + item = await load_queue.get() + if item is None: + return + queued.append(item) + active.extend(dispatcher.submit_ready(self, queued, sample_tensor)) + consumed = dispatcher.consume_ready(self, active, sample_tensor) + if consumed: + continue + if not active: + continue + await self._wait_for_dispatcher_completion(active) + except BaseException as exc: + for state in active: + if not state.item.future.done(): + state.item.future.set_exception(exc) + active_batch = state.active + if active_batch is not None: + active_batch.staging_lease.release() + for item in queued: + if not item.future.done(): + item.future.set_exception(exc) + raise + + async def _drain_load_queue( + self, + load_queue: asyncio.Queue[Any], + queued: list[_QueuedLoadRequest], + ) -> None: + """Move immediately available queue items into a local FIFO list.""" + while True: + try: + item = load_queue.get_nowait() + except asyncio.QueueEmpty: + return + if item is None: + await load_queue.put(None) + return + queued.append(item) + + async def _wait_for_dispatcher_completion( + self, + active: list[_InflightRequestLoad], + ) -> None: + """Wait until any request-level active read future completes.""" + wrapped: dict[asyncio.Future[Any], _InflightRequestLoad] = {} + for state in active: + active_batch = state.active + if active_batch is None: + continue + wrapped[asyncio.wrap_future(active_batch.future)] = state + if not wrapped: + await asyncio.sleep(0) + return + done, _pending = await asyncio.wait( + wrapped.keys(), + timeout=_LOAD_DISPATCH_WAIT_TIMEOUT_S, + return_when=asyncio.FIRST_COMPLETED, + ) + if not done: + await asyncio.sleep(0) + def _load_kv_specs( self, reqs_to_load: dict[str, Any], sample_tensor: torch.Tensor, + request_futures: dict[str, _RequestLoadFuture] | None = None, ) -> None: """Load cross-layer KV specs into vLLM cache on a background thread. @@ -575,117 +1175,456 @@ def _load_kv_specs( reqs_to_load: Worker load specs from scheduler metadata. sample_tensor: Representative KV cache tensor used for device and synchronization. + request_futures: Optional base request ID to completion future map. + When provided, each request future is marked as soon as all + batches containing that request have been restored. Async/thread-safety: Runs on the connector load executor. It submits server transfer RPCs to the dedicated load asyncio loop and waits inside the background thread, not on vLLM's worker thread. """ + try: + load_staging_pool = self._ensure_load_staging_pool(sample_tensor) + load_batches = _build_load_read_batches( + reqs_to_load, + self._slot_size, + max_batch_bytes=load_staging_pool.buffer_bytes, + include_req_ids=request_futures is not None, + ) + if not load_batches: + for request_future in (request_futures or {}).values(): + request_future.set_result() + return + remaining_batches = self._request_load_batch_counts(load_batches) + + total_copies = 0 + total_copy_runs = 0 + total_bytes_loaded = 0 + total_ipc_ms = 0.0 + total_copy_ms = 0.0 + total_sync_ms = 0.0 + total_l1_hits = 0 + total_l1_misses = 0 + total_l2_reads = 0 + total_transfer_open_ms = 0.0 + total_transfer_load_ms = 0.0 + total_transfer_sync_ms = 0.0 + consumed_batches: list[_ConsumedLoadBatch] = [] + + def consume_load_batch(state: _InflightLoadBatch) -> int: + consumed = self._consume_loaded_batch(state, sample_tensor) + consumed_batches.append(consumed) + self._mark_loaded_batch_requests( + state, + request_futures, + remaining_batches, + ) + return consumed.buffer_index + + def wait_for_load_completion( + states: Sequence[_InflightLoadBatch], + ) -> _InflightLoadBatch: + """Wait until any in-flight load future completes.""" + future_to_state = {state.future: state for state in states} + done, _pending = wait_futures( + future_to_state.keys(), + return_when=FIRST_COMPLETED, + ) + first_done = next(iter(done)) + return future_to_state[first_done] + + max_inflight_batches = _run_fixed_depth_pipeline( + batches=load_batches, + submit=lambda batch, buffer_index: self._submit_load_batch( + batch, + buffer_index, + sample_tensor, + ), + is_complete=lambda state: state.future.done(), + consume=consume_load_batch, + buffer_indices=range(_LOAD_PIPELINE_DEPTH), + wait_for_completion=wait_for_load_completion, + ) + + for consumed in consumed_batches: + total_bytes_loaded += consumed.bytes + total_copies += consumed.copies + total_copy_runs += consumed.copy_runs + total_ipc_ms += consumed.ipc_ms + total_copy_ms += consumed.copy_ms + total_transfer_open_ms += consumed.transfer_open_ms + total_transfer_load_ms += consumed.transfer_load_ms + total_transfer_sync_ms += consumed.transfer_sync_ms + total_l1_hits += consumed.l1_hits + total_l1_misses += consumed.l1_misses + total_l2_reads += consumed.l2_reads + + sync_start = time.perf_counter() + _synchronize_cuda_tensor(sample_tensor) + total_sync_ms += (time.perf_counter() - sync_start) * 1000 + + logger.debug( + "[CONNECTOR] start_load_kv timing: reqs=%d batches=%d bytes=%d " + "copy_runs=%d gpu_copies=%d ipc_ms=%.3f copy_ms=%.3f " + "sync_ms=%.3f transfer_open_ms=%.3f transfer_load_ms=%.3f " + "transfer_sync_ms=%.3f l1_hits=%d l1_misses=%d l2_reads=%d " + "load_pipeline_depth=%d max_inflight_batches=%d " + "load_buffer_bytes=%d dynamic_load_allocations=%d", + len(reqs_to_load), + len(load_batches), + total_bytes_loaded, + total_copy_runs, + total_copies, + total_ipc_ms, + total_copy_ms, + total_sync_ms, + total_transfer_open_ms, + total_transfer_load_ms, + total_transfer_sync_ms, + total_l1_hits, + total_l1_misses, + total_l2_reads, + _LOAD_PIPELINE_DEPTH, + max_inflight_batches, + load_staging_pool.buffer_bytes, + 0, + ) + except BaseException as exc: + for request_future in (request_futures or {}).values(): + if not request_future.done(): + request_future.set_exception(exc) + raise + + def _ensure_load_staging_pool( + self, + sample_tensor: torch.Tensor, + ) -> FixedCudaStagingPool: + """Return the fixed load staging pool, creating it before traffic if needed. + + Args: + sample_tensor: Representative KV cache tensor used for device + placement when the pool was not initialized during registration. + + Returns: + Fixed load staging pool with two preallocated buffers. + + Async/thread-safety: + Called from the worker load executor. Normal production flow creates + the pool during KV-cache registration; this fallback keeps tests and + deferred initialization paths explicit while still allocating only + once before batch submission. + """ + pool = getattr(self, "_load_staging_pool", None) + if pool is not None: + return pool + buffer_bytes = max( + self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES, + self._slot_size, + ) + pool = FixedCudaStagingPool( + device=sample_tensor.device, + buffer_bytes=buffer_bytes, + depth=_load_staging_pool_depth( + buffer_bytes, + getattr(self, "_pending_store_staging_limit_bytes", 0), + sample_tensor.device, + ), + ) + self._load_staging_pool = pool + return pool + + def _request_load_batch_counts( + self, + load_batches: Sequence[_LoadBatch], + ) -> dict[str, int]: + """Return how many load batches each request must wait for. + + Args: + load_batches: Bounded load batches whose ranges may include request + IDs when request-level completion is enabled. + + Returns: + Base request ID to number of batches containing that request. + + Async/thread-safety: + Pure CPU helper used on the connector load executor before the + completion-driven load pipeline starts. + """ + counts: dict[str, int] = {} + for _total_bytes, _spans, per_req_ranges in load_batches: + batch_req_ids: set[str] = set() + for item in per_req_ranges: + if len(item) < 4: + continue + _start, _end, req_id, _spec = item + batch_req_ids.add(_load_completion_req_id(str(req_id))) + for req_id in batch_req_ids: + counts[req_id] = counts.get(req_id, 0) + 1 + return counts + + def _mark_loaded_batch_requests( + self, + state: _InflightLoadBatch, + request_futures: dict[str, _RequestLoadFuture] | None, + remaining_batches: dict[str, int], + ) -> None: + """Mark request futures whose final load batch has been restored. + + Args: + state: Consumed load batch state. + request_futures: Base request ID to request completion future map. + remaining_batches: Mutable base request ID to outstanding batch + count map. + + Async/thread-safety: + Called on the connector load executor immediately after staging + bytes have been copied into the vLLM KV cache and synchronized. + """ + if not request_futures: + return + batch_req_ids: set[str] = set() + for item in state.per_req_ranges: + if len(item) < 4: + continue + _start, _end, req_id, _spec = item + batch_req_ids.add(_load_completion_req_id(str(req_id))) + for req_id in batch_req_ids: + remaining = max(0, remaining_batches.get(req_id, 1) - 1) + remaining_batches[req_id] = remaining + if remaining == 0: + future = request_futures.get(req_id) + if future is not None and not future.done(): + future.set_result() + + def _submit_load_batch( + self, + batch: _LoadBatch, + buffer_index: int, + sample_tensor: torch.Tensor, + ) -> _InflightLoadBatch: + """Submit one load batch into a fixed staging buffer. + + Args: + batch: Tuple from ``build_load_read_batches``. + buffer_index: Fixed load staging buffer to use. + sample_tensor: Representative KV cache tensor for device context. + + Returns: + In-flight batch state consumed by ``_consume_loaded_batch``. + + Async/thread-safety: + Runs on the connector load executor. The returned state owns the + fixed staging lease until consumption finishes. + """ + del sample_tensor + total_bytes, spans, per_req_ranges = batch + pool = self._ensure_load_staging_pool(next(iter(self._kv_caches.values()))) + staging_lease = pool.acquire_index(buffer_index, total_bytes) + staging = staging_lease.view + + submitted_at = time.perf_counter() + if getattr(self, "_load_staging_registered", False): + transfer_coro = self._transfer_load_registered_cuda( + buffer_index=buffer_index, + nbytes=total_bytes, + spans=spans, + ) + else: + cp_staging = cupy.asarray(staging) + cuda_handle = export_cuda_ipc_handle(cp_staging) + device_id = cuda_array_device_id(cp_staging) + device_ptr = cuda_array_pointer(cp_staging) + ipc_base_ptr, ipc_offset = _cuda_allocation_base_and_offset(device_ptr) + transfer_coro = self._transfer_load_cuda( + buffer_index=buffer_index, + cuda_ipc_handle=cuda_handle, + nbytes=total_bytes, + device_id=device_id, + device_ptr=device_ptr, + allocation_base_ptr=ipc_base_ptr, + allocation_offset=ipc_offset, + producer_pid=os.getpid(), + spans=spans, + ) + future = self._submit_load_coroutine(transfer_coro) + return _InflightLoadBatch( + buffer_index=buffer_index, + total_bytes=total_bytes, + per_req_ranges=per_req_ranges, + staging_lease=staging_lease, + future=future, + submitted_at=submitted_at, + ) + + def _submit_request_load_for_dispatcher( + self, + item: _QueuedLoadRequest, + buffer_index: int, + sample_tensor: torch.Tensor, + ) -> _InflightRequestLoad: + """Submit the first read batch for one queued request load. + + Args: + item: Request-level queued load. + buffer_index: Fixed staging buffer index reserved by the dispatcher. + sample_tensor: Representative KV cache tensor. + + Returns: + Request state tracked by ``LoadRequestDispatcher``. + + Async/thread-safety: + Runs on the connector load asyncio loop. Each request owns at most one + active staging buffer at a time; multi-segment requests submit the + next segment only after the previous segment has been restored. + """ + load_staging_pool = self._ensure_load_staging_pool(sample_tensor) load_batches = _build_load_read_batches( - reqs_to_load, + {item.spec_id: item.spec}, self._slot_size, - max_batch_bytes=(self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES), + max_batch_bytes=load_staging_pool.buffer_bytes, + include_req_ids=True, ) if not load_batches: - return + item.future.set_result() + return _InflightRequestLoad( + item=item, + buffer_index=buffer_index, + batches=[], + next_batch=0, + remaining_batches=0, + active=None, + completed=[], + ) + active_batch = self._submit_load_batch( + load_batches[0], + buffer_index, + sample_tensor, + ) + return _InflightRequestLoad( + item=item, + buffer_index=buffer_index, + batches=load_batches, + next_batch=1, + remaining_batches=len(load_batches), + active=active_batch, + completed=[], + ) - total_copies = 0 - total_copy_runs = 0 - total_bytes_loaded = 0 - total_ipc_ms = 0.0 - total_copy_ms = 0.0 - total_sync_ms = 0.0 - total_l1_hits = 0 - total_l1_misses = 0 - total_l2_reads = 0 - total_transfer_open_ms = 0.0 - total_transfer_load_ms = 0.0 - total_transfer_sync_ms = 0.0 - for total_bytes, spans, per_req_ranges in load_batches: - total_bytes_loaded += total_bytes - staging_lease = self._acquire_staging(total_bytes, sample_tensor.device) - staging = staging_lease.view - try: - cp_staging = cupy.asarray(staging) - cuda_handle = export_cuda_ipc_handle(cp_staging) - device_id = cuda_array_device_id(cp_staging) - device_ptr = cuda_array_pointer(cp_staging) - ipc_base_ptr, ipc_offset = _cuda_allocation_base_and_offset(device_ptr) + def _consume_dispatcher_load( + self, + state: _InflightRequestLoad, + sample_tensor: torch.Tensor, + ) -> tuple[int, bool]: + """Consume one completed request read and finish or advance the request. - ipc_start = time.perf_counter() - load_response = self._submit_load_coroutine( - self._transfer_load_cuda( - cuda_ipc_handle=cuda_handle, - nbytes=total_bytes, - device_id=device_id, - device_ptr=device_ptr, - allocation_base_ptr=ipc_base_ptr, - allocation_offset=ipc_offset, - producer_pid=os.getpid(), - spans=spans, - ) - ).result(timeout=120.0) - total_ipc_ms += (time.perf_counter() - ipc_start) * 1000 - total_transfer_open_ms += float( - load_response.get("transfer_open_ms", 0.0) - ) - total_transfer_load_ms += float( - load_response.get("transfer_load_ms", 0.0) - ) - total_transfer_sync_ms += float( - load_response.get("transfer_sync_ms", 0.0) - ) - stats = load_response.get("transfer_stats_delta", {}) - if isinstance(stats, dict): - total_l1_hits += int(stats.get("l1_hits", 0)) - total_l1_misses += int(stats.get("l1_misses", 0)) - total_l2_reads += int(stats.get("l2_reads", 0)) - - copy_runs = _build_load_copy_runs(per_req_ranges) - total_copy_runs += len(copy_runs) - copy_start = time.perf_counter() - for run in copy_runs: - total_copies += _copy_staging_to_kv_cache( - staging=staging[run.start : run.end], - kv_caches=self._kv_caches, - layer_names=self._layer_names, - block_ids=run.block_ids, - slot_size=self._slot_size, - load_key_scale=self._load_key_scale, - load_value_scale=self._load_value_scale, - pos_offset=run.pos_offset, - rope_delta_scale=self._rope_delta_scale, - rope_base=self._rope_base, - rope_rotary_dim=self._rope_rotary_dim, - rope_is_neox_style=self._rope_is_neox_style, - ) - total_copy_ms += (time.perf_counter() - copy_start) * 1000 - sync_start = time.perf_counter() - _synchronize_cuda_tensor(sample_tensor) - total_sync_ms += (time.perf_counter() - sync_start) * 1000 - finally: - staging_lease.release() + Args: + state: Request-level load state with a completed active batch. + sample_tensor: Representative KV cache tensor. - logger.debug( - "[CONNECTOR] start_load_kv timing: reqs=%d batches=%d bytes=%d " - "copy_runs=%d gpu_copies=%d ipc_ms=%.3f copy_ms=%.3f " - "sync_ms=%.3f transfer_open_ms=%.3f transfer_load_ms=%.3f " - "transfer_sync_ms=%.3f l1_hits=%d l1_misses=%d l2_reads=%d", - len(reqs_to_load), - len(load_batches), - total_bytes_loaded, - total_copy_runs, - total_copies, - total_ipc_ms, - total_copy_ms, - total_sync_ms, - total_transfer_open_ms, - total_transfer_load_ms, - total_transfer_sync_ms, - total_l1_hits, - total_l1_misses, - total_l2_reads, + Returns: + Tuple of released staging buffer index and whether the request is + fully complete. + + Async/thread-safety: + Runs on the connector load asyncio loop after the associated transfer + future is complete. The staging buffer is not reused until restore + kernels have synchronized and the lease has been released. + """ + active_batch = state.active + if active_batch is None: + return state.buffer_index, True + consumed = self._consume_loaded_batch(active_batch, sample_tensor) + state.completed.append(consumed) + state.remaining_batches = max(0, state.remaining_batches - 1) + reusable_buffer = consumed.buffer_index + state.buffer_index = reusable_buffer + if state.remaining_batches == 0: + if not state.item.future.done(): + state.item.future.set_result() + return reusable_buffer, True + state.active = self._submit_load_batch( + state.batches[state.next_batch], + reusable_buffer, + sample_tensor, ) + state.next_batch += 1 + return reusable_buffer, False + + def _consume_loaded_batch( + self, + state: _InflightLoadBatch, + sample_tensor: torch.Tensor, + ) -> _ConsumedLoadBatch: + """Wait for one submitted load batch and restore it into vLLM KV cache. + + Args: + state: In-flight load batch returned by ``_submit_load_batch``. + sample_tensor: Representative KV cache tensor for synchronization. + + Returns: + Timing and accounting data for the restored batch. + + Async/thread-safety: + Runs on the connector load executor. It releases the fixed staging + lease only after restore kernels that read the staging view have + been synchronized. + """ + try: + wait_start = time.perf_counter() + load_response = state.future.result(timeout=120.0) + wait_ms = (time.perf_counter() - wait_start) * 1000 + ipc_ms = (time.perf_counter() - state.submitted_at) * 1000 + transfer_open_ms = float(load_response.get("transfer_open_ms", 0.0)) + transfer_load_ms = float(load_response.get("transfer_load_ms", 0.0)) + transfer_sync_ms = float(load_response.get("transfer_sync_ms", 0.0)) + stats = load_response.get("transfer_stats_delta", {}) + l1_hits = 0 + l1_misses = 0 + l2_reads = 0 + if isinstance(stats, dict): + l1_hits = int(stats.get("l1_hits", 0)) + l1_misses = int(stats.get("l1_misses", 0)) + l2_reads = int(stats.get("l2_reads", 0)) + + copy_runs = _build_load_copy_runs(state.per_req_ranges) + copy_start = time.perf_counter() + copies = 0 + staging = state.staging_lease.view + for run in copy_runs: + copies += _copy_staging_to_kv_cache( + staging=staging[run.start : run.end], + kv_caches=self._kv_caches, + layer_names=self._layer_names, + block_ids=run.block_ids, + slot_size=self._slot_size, + load_key_scale=self._load_key_scale, + load_value_scale=self._load_value_scale, + pos_offset=run.pos_offset, + rope_delta_scale=self._rope_delta_scale, + rope_base=self._rope_base, + rope_rotary_dim=self._rope_rotary_dim, + rope_is_neox_style=self._rope_is_neox_style, + ) + copy_ms = (time.perf_counter() - copy_start) * 1000 + _synchronize_cuda_tensor(sample_tensor) + return _ConsumedLoadBatch( + buffer_index=state.buffer_index, + bytes=state.total_bytes, + copies=copies, + copy_runs=len(copy_runs), + ipc_ms=ipc_ms, + wait_ms=wait_ms, + copy_ms=copy_ms, + transfer_open_ms=transfer_open_ms, + transfer_load_ms=transfer_load_ms, + transfer_sync_ms=transfer_sync_ms, + l1_hits=l1_hits, + l1_misses=l1_misses, + l2_reads=l2_reads, + ) + finally: + state.staging_lease.release() def wait_for_layer_load(self, layer_name: str) -> None: """No-op because async loads complete before vLLM resumes requests. @@ -845,13 +1784,28 @@ def shutdown(self) -> None: for req_id in list(getattr(self, "_pending_finished_saves", {})): self.get_finished({req_id}) self._reap_save_futures(block=True) - load_executor = getattr(self, "_load_executor", None) - if load_executor is not None: - load_executor.shutdown(wait=True) - load_client = getattr(self, "_ipc_load_async", None) + load_queue = getattr(self, "_load_request_queue", None) + if load_queue is not None: + if isinstance(load_queue, asyncio.Queue): + self._load_loop.call_soon_threadsafe(load_queue.put_nowait, None) + else: + load_queue.put(None) + dispatcher_future = getattr(self, "_load_request_dispatcher_future", None) + if dispatcher_future is not None: + try: + dispatcher_future.result(timeout=120.0) + except Exception: # noqa: BLE001 + pass + load_clients = list( + dict.fromkeys( + getattr(self, "_ipc_load_async_pool", []) + or [getattr(self, "_ipc_load_async", None)] + ) + ) store_client = getattr(self, "_ipc_store_async", None) - if load_client is not None: - self._submit_load_coroutine(load_client.close()).result(timeout=10.0) + for load_client in load_clients: + if load_client is not None: + self._submit_load_coroutine(load_client.close()).result(timeout=10.0) if store_client is not None: self._submit_store_coroutine(store_client.close()).result(timeout=10.0) self._load_loop.call_soon_threadsafe(self._load_loop.stop) @@ -936,12 +1890,98 @@ def _init_server_transfer(self) -> None: and getattr(self, "_store_loop", None) is not None ): return - self._submit_load_coroutine(self._ipc_load_async.init_transfer()).result( - timeout=120.0 - ) + for load_client in self._load_ipc_clients(): + self._submit_load_coroutine(load_client.init_transfer()).result( + timeout=120.0 + ) self._submit_store_coroutine(self._ipc_store_async.init_transfer()).result( timeout=120.0 ) + self._register_load_staging_buffers() + + def _load_ipc_clients(self) -> list[Any]: + """Return fixed load IPC clients used for parallel load RPCs. + + Returns: + Load IPC clients. The list falls back to the legacy single client + when the connector was constructed by tests or older harnesses. + + Async/thread-safety: + Called on the vLLM worker thread during initialization and from the + load event loop when selecting the client for a submitted batch. + """ + clients = getattr(self, "_ipc_load_async_pool", None) + if clients: + return list(clients) + client = getattr(self, "_ipc_load_async", None) + return [client] if client is not None else [] + + def _load_ipc_client_for_buffer(self, buffer_index: int | None = None) -> Any: + """Return the load IPC client assigned to a staging buffer. + + Args: + buffer_index: Fixed staging buffer index for the transfer. + + Returns: + Async IPC client for the selected load lane. + + Async/thread-safety: + Pure selection helper. Each returned client owns its own IPC socket, + so separate fixed staging buffers can have concurrent server RPCs + instead of serializing on one client lock. + """ + clients = self._load_ipc_clients() + if not clients: + raise RuntimeError("load IPC client is not initialized") + if buffer_index is None: + return clients[0] + return clients[int(buffer_index) % len(clients)] + + def _register_load_staging_buffers(self) -> None: + """Register fixed load staging buffers with the server. + + Async/thread-safety: + Called after server transfer initialization and before request + traffic. Registration failures are logged and leave the worker on + the compatible per-load CUDA IPC payload path. + """ + if getattr(self, "_load_staging_registered", False): + return + pool = getattr(self, "_load_staging_pool", None) + if pool is None or not self._load_ipc_clients(): + return + try: + for buffer_index in range(pool.depth): + tensor = pool.buffer(buffer_index) + cp_tensor = cupy.asarray(tensor) + device_ptr = cuda_array_pointer(cp_tensor) + ipc_base_ptr, ipc_offset = _cuda_allocation_base_and_offset(device_ptr) + load_client = self._load_ipc_client_for_buffer(buffer_index) + self._submit_load_coroutine( + load_client.register_load_staging_cuda( + buffer_index=buffer_index, + cuda_ipc_handle=export_cuda_ipc_handle(cp_tensor), + allocation_bytes=int(tensor.numel()), + device_id=cuda_array_device_id(cp_tensor), + device_ptr=device_ptr, + allocation_base_ptr=ipc_base_ptr, + allocation_offset=ipc_offset, + producer_pid=os.getpid(), + ) + ).result(timeout=120.0) + except Exception as exc: # noqa: BLE001 + logger.warning( + "[CONNECTOR] registered load staging unavailable; falling back " + "to per-load CUDA IPC payloads: %s", + exc, + ) + self._load_staging_registered = False + return + self._load_staging_registered = True + logger.info( + "[CONNECTOR] registered %d fixed load staging buffers", + pool.depth, + ) def _reap_save_futures(self, block: bool) -> None: """Collect completed background save tasks. @@ -1015,6 +2055,33 @@ def _wait_for_save_staging_capacity(self, nbytes: int) -> None: record.release() self._reap_save_futures(block=False) + def _wait_for_store_staging_release(self, nbytes: int) -> None: + """Wait until one store staging lease can return to the fixed pool. + + Args: + nbytes: Size of the staging lease requested by the caller. + + Async/thread-safety: + Called from the worker thread when the fixed store staging pool is + exhausted. It first applies byte-budget backpressure, then waits + for one oldest store future if no lease was released yet. + """ + pool = getattr(self, "_store_staging_pool", None) + before = pool.available if pool is not None else 0 + self._wait_for_save_staging_capacity(nbytes) + if pool is None or pool.available > before or not self._save_futures: + return + record = self._save_futures.pop(0) + try: + record.future.result(timeout=120.0) + finally: + self._pending_save_staging_bytes = max( + 0, + self._pending_save_staging_bytes - record.staging_bytes, + ) + record.release() + self._reap_save_futures(block=False) + def _acquire_staging( self, nbytes: int, @@ -1033,19 +2100,26 @@ def _acquire_staging( Called from the worker thread. Store-path callers must retain the lease until the background server transfer completes. """ - pool = getattr(self, "_staging_pool", None) + pool = getattr(self, "_store_staging_pool", None) if pool is None: max_bytes = max( nbytes, self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES, ) - pool = CudaStagingPool( + pool = StoreCudaStagingPool( device=device, - initial_bytes=0, - max_buffer_bytes=max_bytes, + buffer_bytes=max_bytes, + depth=_store_staging_pool_depth( + max_bytes, + self._pending_store_staging_limit_bytes + or DEFAULT_PENDING_STORE_STAGING_BYTES, + ), ) - self._staging_pool = pool - return pool.acquire(nbytes) + self._store_staging_pool = pool + return pool.acquire( + nbytes, + wait_for_release=lambda: self._wait_for_store_staging_release(nbytes), + ) def _stage_store_batch( self, @@ -1244,7 +2318,27 @@ async def _transfer_load_cuda(self, **kwargs: Any) -> dict[str, Any]: Runs on the worker load event loop. A dedicated client keeps cache-hit loads from queueing behind store RPCs. """ - return await self._ipc_load_async.transfer_load_cuda(**kwargs) + buffer_index = kwargs.pop("buffer_index", None) + client = self._load_ipc_client_for_buffer(buffer_index) + return await client.transfer_load_cuda(**kwargs) + + async def _transfer_load_registered_cuda(self, **kwargs: Any) -> dict[str, Any]: + """Load through a pre-registered fixed CUDA staging buffer. + + Args: + **kwargs: forwarded registered-buffer transfer fields. + + Returns: + Server load response with timing counters. + + Async/thread-safety: + Runs on the worker load event loop. The server has already opened + the CUDA IPC mapping during initialization, so this hot-path call + only identifies the staging buffer index and logical byte range. + """ + buffer_index = int(kwargs.get("buffer_index", 0)) + client = self._load_ipc_client_for_buffer(buffer_index) + return await client.transfer_load_registered_cuda(**kwargs) async def _transfer_store_cuda(self, **kwargs: Any) -> list[str]: """Store through the dedicated worker store IPC client. diff --git a/daser/server/ipc/server.py b/daser/server/ipc/server.py index e73a6ae..95102db 100644 --- a/daser/server/ipc/server.py +++ b/daser/server/ipc/server.py @@ -126,6 +126,7 @@ def __init__( self._cuda_ipc_cache: OrderedDict[ tuple[int, int, int, int | None], "_CachedCudaArray" ] = OrderedDict() + self._load_staging_buffers: dict[int, _CachedCudaArray] = {} self._op_handlers: dict[ str, Callable[[dict[str, Any]], Awaitable[dict[str, Any]]] ] = { @@ -143,6 +144,7 @@ def __init__( "init_transfer": self._op_init_transfer, "transfer_store": self._transfer_store, "transfer_load": self._transfer_load, + "register_load_staging": self._register_load_staging, "evict_chunk": self._op_evict_chunk, "release_chunk_writer": self._op_release_chunk_writer, } @@ -484,6 +486,9 @@ async def _transfer_load(self, msg: dict[str, Any]) -> dict[str, Any]: load_ms = 0.0 sync_ms = 0.0 close_ms = 0.0 + close_one_shot_buffer = isinstance(buffer, _UncachedCudaArray) and ( + "load_staging_buffer_index" not in payload + ) try: before = asdict(transfer.stats) started = time.perf_counter() @@ -532,7 +537,7 @@ async def _transfer_load(self, msg: dict[str, Any]) -> dict[str, Any]: ) return response finally: - if isinstance(buffer, _UncachedCudaArray): + if close_one_shot_buffer: close = getattr(buffer, "close", None) close_start = time.perf_counter() close() @@ -543,6 +548,34 @@ async def _transfer_load(self, msg: dict[str, Any]) -> dict[str, Any]: close_ms, ) + async def _register_load_staging(self, msg: dict[str, Any]) -> dict[str, Any]: + """Register one fixed CUDA load staging buffer for indexed reuse. + + Args: + msg: IPC request whose payload describes a worker-owned fixed CUDA + staging allocation. + + Returns: + ``{"ok": True}`` after the CUDA IPC mapping is open and cached. + + Async/thread-safety: + Runs on the IPC event loop during worker initialization. Replacing + an existing index closes the old mapping only after the new mapping + is available. + """ + payload = msg.get("payload", {}) + buffer_index = int(payload["buffer_index"]) + opened = self._open_cuda_ipc_payload( + payload=payload, + nbytes_key="allocation_bytes", + cache_mapping=False, + ) + previous = self._load_staging_buffers.get(buffer_index) + self._load_staging_buffers[buffer_index] = opened + if previous is not None: + previous.close() + return {"ok": True} + def _record_transfer_metrics( self, op: str, @@ -669,49 +702,79 @@ def _payload_buffer(self, payload: dict[str, Any]) -> Any: """Return a byte-addressable buffer for an IPC transfer payload.""" if "data" in payload: return bytearray(payload["data"]) + if "load_staging_buffer_index" in payload: + buffer_index = int(payload["load_staging_buffer_index"]) + try: + return self._load_staging_buffers[buffer_index] + except KeyError as exc: + raise ValueError( + f"unknown load staging buffer index: {buffer_index}" + ) from exc if "cuda_ipc_handle" in payload: - local_ptr = None - producer_pid = int(payload.get("producer_pid", -1)) - device_ptr = int(payload["device_ptr"]) - nbytes = int(payload["nbytes"]) - device_id = int(payload["device_id"]) if "device_id" in payload else None - allocation_offset = int(payload.get("allocation_offset", 0)) - allocation_base_ptr = int( - payload.get("allocation_base_ptr", device_ptr - allocation_offset) + return self._open_cuda_ipc_payload(payload=payload, nbytes_key="nbytes") + raise ValueError("transfer payload requires data or cuda_ipc_handle") + + def _open_cuda_ipc_payload( + self, + *, + payload: dict[str, Any], + nbytes_key: str, + cache_mapping: bool = True, + ) -> "_CachedCudaArray": + """Open or reuse a CUDA IPC payload mapping. + + Args: + payload: CUDA IPC payload with handle, device, allocation base, and + offset fields. + nbytes_key: Payload field that identifies the mapping size. + cache_mapping: If True, reuse the general CUDA IPC LRU cache. Fixed + registered staging buffers pass False because their lifetime is + owned by ``_load_staging_buffers`` and must not be LRU-evicted. + + Returns: + Cached or one-shot CUDA array wrapper for the mapped allocation. + """ + local_ptr = None + producer_pid = int(payload.get("producer_pid", -1)) + device_ptr = int(payload["device_ptr"]) + nbytes = int(payload[nbytes_key]) + device_id = int(payload["device_id"]) if "device_id" in payload else None + allocation_offset = int(payload.get("allocation_offset", 0)) + allocation_base_ptr = int( + payload.get("allocation_base_ptr", device_ptr - allocation_offset) + ) + if producer_pid == os.getpid(): + local_ptr = allocation_base_ptr + if local_ptr is None and cache_mapping: + key = ( + producer_pid, + allocation_base_ptr, + nbytes + allocation_offset, + device_id, ) - if producer_pid == os.getpid(): - local_ptr = allocation_base_ptr - if local_ptr is None: - key = ( - producer_pid, - allocation_base_ptr, - nbytes + allocation_offset, - device_id, + cached = self._cuda_ipc_cache.get(key) + if cached is None: + self._evict_cuda_ipc_cache_if_needed() + opened = open_cuda_ipc_buffer( + handle=payload["cuda_ipc_handle"], + nbytes=nbytes, + device_id=device_id, + local_ptr=None, + allocation_offset=allocation_offset, ) - cached = self._cuda_ipc_cache.get(key) - if cached is None: - self._evict_cuda_ipc_cache_if_needed() - opened = open_cuda_ipc_buffer( - handle=payload["cuda_ipc_handle"], - nbytes=nbytes, - device_id=device_id, - local_ptr=None, - allocation_offset=allocation_offset, - ) - cached = _CachedCudaArray(opened) - self._cuda_ipc_cache[key] = cached - else: - self._cuda_ipc_cache.move_to_end(key) - return cached - opened = open_cuda_ipc_buffer( - handle=payload["cuda_ipc_handle"], - nbytes=nbytes, - device_id=device_id, - local_ptr=local_ptr, - allocation_offset=allocation_offset, - ) - return _UncachedCudaArray(opened) - raise ValueError("transfer payload requires data or cuda_ipc_handle") + cached = _CachedCudaArray(opened) + self._cuda_ipc_cache[key] = cached + else: + self._cuda_ipc_cache.move_to_end(key) + return cached + opened = open_cuda_ipc_buffer( + handle=payload["cuda_ipc_handle"], + nbytes=nbytes, + device_id=device_id, + local_ptr=local_ptr, + allocation_offset=allocation_offset, + ) + return _UncachedCudaArray(opened) def _evict_cuda_ipc_cache_if_needed(self) -> None: """Evict one cached CUDA IPC mapping when the cache is full.""" @@ -722,6 +785,9 @@ def _evict_cuda_ipc_cache_if_needed(self) -> None: def _close_cuda_ipc_cache(self) -> None: """Close all cached CUDA IPC mappings.""" + for cached in self._load_staging_buffers.values(): + cached.close() + self._load_staging_buffers.clear() for cached in self._cuda_ipc_cache.values(): cached.close() self._cuda_ipc_cache.clear() diff --git a/daser/transfer/iouring/layer.py b/daser/transfer/iouring/layer.py index 95a82a7..0d79fd4 100644 --- a/daser/transfer/iouring/layer.py +++ b/daser/transfer/iouring/layer.py @@ -532,51 +532,61 @@ async def _load_l2_miss_batch( if self._l2 is None: raise RuntimeError("L2 reads are disabled when skip_l2 is true") reads: list[tuple[dict[str, int], PinnedMemorySlice]] = [] + future_to_read: dict[ + asyncio.Future[int], + tuple[dict[str, int], PinnedMemorySlice], + ] = {} + promoted: set[int] = set() try: for span in misses: nbytes = int(span["nbytes"]) key = (int(span["file_offset"]), nbytes) pinned = await self._reserve_l1_buffer(key, nbytes) reads.append((span, pinned)) - - await asyncio.gather( - *( - loop.run_in_executor( - self._l2.executor, - self._read_l2_into, - int(span["file_offset"]), - pinned, - self._next_uring(), - ) - for span, pinned in reads + future = loop.run_in_executor( + self._l2.executor, + self._read_l2_into, + int(span["file_offset"]), + pinned, + self._next_uring(), ) - ) + future_to_read[future] = (span, pinned) - self._copy_grouped_to_dst( - dst, - [ - ( - int(span["target_offset"]), - pinned, - 0, - int(span["nbytes"]), + pending: set[asyncio.Future[int]] = set(future_to_read) + while pending: + done, pending = await asyncio.wait( + pending, + return_when=asyncio.FIRST_COMPLETED, + ) + for future in done: + future.result() + span, pinned = future_to_read[future] + self._copy_grouped_to_dst( + dst, + [ + ( + int(span["target_offset"]), + pinned, + 0, + int(span["nbytes"]), + ) + ], ) - for span, pinned in reads - ], - ) - async with self._lock: - self._raise_l2_error_locked() - for span, pinned in reads: - nbytes = int(span["nbytes"]) - key = (int(span["file_offset"]), nbytes) - self._stats.l2_reads += 1 - self._l1.put(key, pinned) + async with self._lock: + self._raise_l2_error_locked() + nbytes = int(span["nbytes"]) + key = (int(span["file_offset"]), nbytes) + self._stats.l2_reads += 1 + self._l1.put(key, pinned) + promoted.add(id(pinned)) except BaseException: + if future_to_read: + await asyncio.gather(*future_to_read, return_exceptions=True) live_buffers = set() async with self._lock: live_buffers = self._l1.resident_slice_ids() for _span, pinned in reads: - if id(pinned) not in live_buffers: + if id(pinned) not in live_buffers and id(pinned) not in promoted: pinned.close() raise diff --git a/tests/connector/test_daser_connector.py b/tests/connector/test_daser_connector.py index 79003f6..376e781 100644 --- a/tests/connector/test_daser_connector.py +++ b/tests/connector/test_daser_connector.py @@ -2,6 +2,7 @@ # Standard import asyncio +from types import SimpleNamespace # Third Party import cupy @@ -36,7 +37,7 @@ DEFAULT_ROPE_DELTA_SCALE, DEFAULT_STORE_STAGING_BYTES, MIN_STORE_STAGING_BYTES, - CudaStagingPool, + StoreCudaStagingPool, ) from daser.connector.staging import ( apply_rope_delta_to_key_block as _apply_rope_delta_to_key_block, @@ -66,9 +67,13 @@ synchronize_cuda_tensor as _synchronize_cuda_tensor, ) from daser.connector.worker import ( + LoadRequestDispatcher, WorkerConnectorMixin, _DeferredFinishedSave, + _InflightRequestLoad, + _load_staging_pool_depth, _PendingLoad, + _RequestLoadFuture, ) BLOCK_TOKENS = 4 @@ -140,24 +145,21 @@ def __init__(self, store_path: str) -> None: self._save_futures = [] self._pending_save_staging_bytes = 0 - class _InlineExecutor: - def submit(self, fn, *args): - fn(*args) - - class _DoneFuture: - def done(self) -> bool: - return True - - def result(self, timeout: float | None = None) -> None: - del timeout - - return _DoneFuture() - - self._load_executor = _InlineExecutor() - def _refresh_runtime_config(self) -> None: return + def _ensure_load_request_dispatcher(self, sample_tensor: torch.Tensor) -> None: + """Drain queued requests synchronously for worker probe tests.""" + while not self._load_request_queue.empty(): + item = self._load_request_queue.get() + if item is None: + return + self._load_kv_specs( + {item.spec_id: item.spec}, + sample_tensor, + {item.req_id: item.future}, + ) + @property def transfer_ready(self): return self._transfer_ready @@ -520,6 +522,10 @@ def result(self, timeout: float | None = None) -> None: if self._error is not None: raise self._error + def mark_done(self) -> None: + """Mark this fake future as complete.""" + self._done = True + class _AsyncLoadProbe(WorkerConnectorMixin): """Worker probe with pending async load futures.""" @@ -531,22 +537,11 @@ def __init__(self) -> None: self._save_futures = [] self._invalid_load_block_ids = set() self.released: list[object] = [] - self.load_executor_shutdown = False self.load_loop_stopped = False self.store_loop_stopped = False self.load_thread_joined = False self.store_thread_joined = False - class _Executor: - def __init__(self, probe: _AsyncLoadProbe) -> None: - self._probe = probe - - def shutdown(self, wait: bool) -> None: - del wait - self._probe.load_executor_shutdown = True - - self._load_executor = _Executor(self) - class _Loop: def __init__(self, callback) -> None: self._callback = callback @@ -804,6 +799,9 @@ def release(self) -> None: return class _Future: + def done(self) -> bool: + return True + def result(self, timeout: float): del timeout return { @@ -842,6 +840,260 @@ def fake_submit_load_coroutine(coro): assert "start_load_kv timing" not in caplog.text +def test_worker_load_pipeline_consumes_completed_batch_first(monkeypatch) -> None: + """Worker load path consumes whichever fixed-buffer batch completes first.""" + + class Probe(WorkerConnectorMixin): + def __init__(self) -> None: + self._slot_size = 4 + self._store_staging_bytes = 8 + self._kv_caches = {"layer.0": torch.empty(4, 1, 1, 4, dtype=torch.uint8)} + self._layer_names = ["layer.0"] + self._load_key_scale = 1.0 + self._load_value_scale = 1.0 + self._rope_delta_scale = 1.0 + self._rope_base = 10000.0 + self._rope_rotary_dim = 0 + self._rope_is_neox_style = True + self._load_staging_pool = None + + def load_specs(self, reqs_to_load: dict[str, ReqLoadSpec]) -> None: + """Expose load execution through a public test helper.""" + self._load_kv_specs(reqs_to_load, next(iter(self._kv_caches.values()))) + + class _FakeFuture: + def __init__(self, total_bytes: int) -> None: + self.total_bytes = total_bytes + + def done(self) -> bool: + return completed_by_bytes[self.total_bytes] + + probe = Probe() + reqs_to_load = { + "req": ReqLoadSpec( + chunk_key="hit", + start_slot=0, + num_slots=3, + block_ids=[0, 1, 2], + file_offset=0, + token_count=12, + ) + } + events: list[str] = [] + completed_by_bytes = {8: False, 4: True} + + def fake_submit(batch, buffer_index: int, sample_tensor: torch.Tensor): + del sample_tensor + total_bytes, _spans, _per_req_ranges = batch + events.append(f"submit:{total_bytes}:buf{buffer_index}") + return SimpleNamespace( + total_bytes=total_bytes, + buffer_index=buffer_index, + future=_FakeFuture(total_bytes), + ) + + def fake_consume(state, sample_tensor: torch.Tensor): + del sample_tensor + total_bytes = state.total_bytes + buffer_index = state.buffer_index + events.append(f"consume:{total_bytes}:buf{buffer_index}") + if total_bytes == 4: + completed_by_bytes[8] = True + return SimpleNamespace( + buffer_index=buffer_index, + bytes=total_bytes, + copies=1, + copy_runs=1, + ipc_ms=0.0, + copy_ms=0.0, + transfer_open_ms=0.0, + transfer_load_ms=0.0, + transfer_sync_ms=0.0, + l1_hits=0, + l1_misses=0, + l2_reads=0, + ) + + monkeypatch.setattr(probe, "_submit_load_batch", fake_submit) + monkeypatch.setattr(probe, "_consume_loaded_batch", fake_consume) + monkeypatch.setattr( + "daser.connector.worker._synchronize_cuda_tensor", + lambda _: None, + ) + + probe.load_specs(reqs_to_load) + + assert events[:3] == [ + "submit:8:buf0", + "submit:4:buf1", + "consume:4:buf1", + ] + + +def test_worker_load_batches_use_buffer_scoped_ipc_clients() -> None: + """Parallel load batches should not serialize on one async IPC client lock.""" + + class _Client: + def __init__(self, name: str) -> None: + self.name = name + self.calls: list[dict[str, object]] = [] + + async def transfer_load_registered_cuda(self, **kwargs): + self.calls.append(kwargs) + return {"client": self.name} + + class Probe(WorkerConnectorMixin): + def __init__(self) -> None: + self._ipc_load_async = _Client("fallback") + self._ipc_load_async_pool = [_Client("load0"), _Client("load1")] + + probe = Probe() + + coro0 = probe._transfer_load_registered_cuda( # noqa: SLF001 + buffer_index=0, + nbytes=4, + spans=[], + ) + coro1 = probe._transfer_load_registered_cuda( # noqa: SLF001 + buffer_index=1, + nbytes=4, + spans=[], + ) + + assert asyncio.run(coro0) == {"client": "load0"} + assert asyncio.run(coro1) == {"client": "load1"} + assert probe._ipc_load_async.calls == [] # noqa: SLF001 + + +def test_group_load_marks_each_request_done_after_its_batch(monkeypatch) -> None: + """A request should become collectable once its load batch is restored.""" + + class Probe(WorkerConnectorMixin): + def __init__(self) -> None: + self._slot_size = 4 + self._store_staging_bytes = 4 + self._kv_caches = {"layer.0": torch.empty(4, 1, 1, 4, dtype=torch.uint8)} + self._layer_names = ["layer.0"] + self._load_key_scale = 1.0 + self._load_value_scale = 1.0 + self._rope_delta_scale = 1.0 + self._rope_base = 10000.0 + self._rope_rotary_dim = 0 + self._rope_is_neox_style = True + self._load_staging_pool = None + self._pending_loads = {} + self.mid_batch_checks = 0 + + def load_specs(self, reqs_to_load: dict[str, ReqLoadSpec]) -> None: + """Expose load execution through a public test helper.""" + request_futures = {req_id: _RequestLoadFuture() for req_id in reqs_to_load} + for req_id, spec in reqs_to_load.items(): + self._pending_loads[req_id] = _PendingLoad( + future=request_futures[req_id], + block_ids=list(spec.block_ids), + lease=None, + ) + self._load_kv_specs( + reqs_to_load, + next(iter(self._kv_caches.values())), + request_futures, + ) + + def _mark_loaded_batch_requests( + self, + state, + request_futures, + remaining_batches, + ) -> None: + super()._mark_loaded_batch_requests( + state, + request_futures, + remaining_batches, + ) + if state.total_bytes == 4: + self.mid_batch_checks += 1 + assert self._pending_loads["req-b"].future.done() + assert not self._pending_loads["req-a"].future.done() + + class _FakeFuture: + def __init__(self, total_bytes: int) -> None: + self.total_bytes = total_bytes + + def done(self) -> bool: + return completed_by_bytes[self.total_bytes] + + probe = Probe() + reqs_to_load = { + "req-a": ReqLoadSpec("hit-a", 0, 1, [0], 0, 4), + "req-b": ReqLoadSpec("hit-b", 1, 1, [1], 4, 4), + } + completed_by_bytes = {4: True} + consume_count = 0 + + def fake_submit(batch, buffer_index: int, sample_tensor: torch.Tensor): + del sample_tensor + total_bytes, _spans, per_req_ranges = batch + # Make req-b complete first even though it was submitted second. + total = 8 + buffer_index if per_req_ranges[0][3].chunk_key == "hit-a" else 4 + completed_by_bytes.setdefault(total, total == 4) + return SimpleNamespace( + total_bytes=total, + buffer_index=buffer_index, + per_req_ranges=per_req_ranges, + future=_FakeFuture(total), + ) + + def fake_consume(state, sample_tensor: torch.Tensor): + nonlocal consume_count + del sample_tensor + consume_count += 1 + if state.total_bytes == 4: + completed_by_bytes[8] = True + return SimpleNamespace( + buffer_index=state.buffer_index, + bytes=4, + copies=1, + copy_runs=1, + ipc_ms=0.0, + copy_ms=0.0, + transfer_open_ms=0.0, + transfer_load_ms=0.0, + transfer_sync_ms=0.0, + l1_hits=0, + l1_misses=0, + l2_reads=0, + ) + + monkeypatch.setattr(probe, "_submit_load_batch", fake_submit) + monkeypatch.setattr(probe, "_consume_loaded_batch", fake_consume) + monkeypatch.setattr( + "daser.connector.worker._synchronize_cuda_tensor", + lambda _: None, + ) + + probe.load_specs(reqs_to_load) + + assert consume_count == 2 + assert probe.mid_batch_checks == 1 + + +def test_get_finished_releases_each_request_load_independently() -> None: + """Completed request loads should not wait for unrelated pending loads.""" + connector = _AsyncLoadProbe() + done = _AsyncLoadFuture(done=True) + pending = _AsyncLoadFuture(done=False) + connector.seed_load("req-a", done, block_ids=[1]) + connector.seed_load("req-b", pending, block_ids=[2]) + + finished_sending, finished_recving = connector.get_finished(set()) + + assert finished_sending is None + assert finished_recving == {"req-a"} + assert done.result_calls == 1 + assert pending.result_calls == 0 + assert set(connector._pending_loads) == {"req-b"} # noqa: SLF001 + + def test_start_load_kv_releases_waiting_request_when_load_cannot_start() -> None: """Worker should not leave async-hit requests waiting forever.""" connector = _WorkerProbe("") @@ -1072,6 +1324,130 @@ def test_get_finished_reports_completed_async_load() -> None: assert connector._invalid_load_block_ids == set() # noqa: SLF001 +def test_request_load_completion_handles_split_staging_ids() -> None: + """Split staging batches should complete the original request future.""" + future = _RequestLoadFuture() + connector = _AsyncLoadProbe() + state = SimpleNamespace( + per_req_ranges=[ + ( + 0, + 4, + "req#0", + ReqLoadSpec("hit-a", 0, 1, [4], 0, 4), + ) + ] + ) + + connector._mark_loaded_batch_requests( # noqa: SLF001 + state, + {"req": future}, + {"req": 1}, + ) + + assert future.done() + + +def test_load_request_dispatcher_limits_active_requests_by_staging_depth() -> None: + """Request dispatcher should cap active loads by max inflight and buffers.""" + dispatcher = LoadRequestDispatcher(max_inflight=8, staging_depth=3) + + assert dispatcher.effective_inflight == 3 + + +def test_load_staging_pool_depth_respects_available_cuda_memory(monkeypatch) -> None: + """Load staging depth should not preallocate past available CUDA headroom.""" + monkeypatch.setattr( + torch.cuda, + "mem_get_info", + lambda device=None: ((4 << 30), 80 << 30), + ) + + depth = _load_staging_pool_depth( + buffer_bytes=1536 << 20, + pending_limit_bytes=12 << 30, + device=torch.device("cuda"), + ) + + assert depth == 2 + + +def test_load_request_dispatcher_submits_without_waiting_for_batch() -> None: + """Queued requests should be submitted immediately while slots are free.""" + + class Probe(_AsyncLoadProbe): + def __init__(self) -> None: + super().__init__() + self.submitted: list[str] = [] + self.completed: list[str] = [] + + def _submit_request_load_for_dispatcher( + self, item, buffer_index, sample_tensor + ): + del sample_tensor + self.submitted.append(item.req_id) + future = _AsyncLoadFuture(done=False) + active = SimpleNamespace(future=future) + return _InflightRequestLoad( + item=item, + buffer_index=buffer_index, + batches=[], + next_batch=0, + remaining_batches=1, + active=active, + completed=[], + ) + + def _consume_dispatcher_load(self, state, sample_tensor): + del sample_tensor + self.completed.append(state.item.req_id) + state.item.future.set_result() + return 0, True + + probe = Probe() + dispatcher = LoadRequestDispatcher(max_inflight=8, staging_depth=2) + sample = torch.empty(1) + first_future = _RequestLoadFuture() + second_future = _RequestLoadFuture() + third_future = _RequestLoadFuture() + queued = [ + SimpleNamespace( + req_id="req-a", + spec_id="req-a", + spec=ReqLoadSpec("hit-a", 0, 1, [1], 0, 4), + future=first_future, + ), + SimpleNamespace( + req_id="req-b", + spec_id="req-b", + spec=ReqLoadSpec("hit-b", 1, 1, [2], 4, 4), + future=second_future, + ), + SimpleNamespace( + req_id="req-c", + spec_id="req-c", + spec=ReqLoadSpec("hit-c", 2, 1, [3], 8, 4), + future=third_future, + ), + ] + + active = dispatcher.submit_ready(probe, queued, sample) + + assert probe.submitted == ["req-a", "req-b"] + assert [item.req_id for item in queued] == ["req-c"] + assert len(active) == 2 + + active_batch = active[0].active + assert active_batch is not None + active_batch.future.mark_done() + dispatcher.consume_ready(probe, active, sample) + dispatcher.submit_ready(probe, queued, sample) + + assert probe.completed == ["req-a"] + assert probe.submitted == ["req-a", "req-b", "req-c"] + assert first_future.done() + + def test_get_finished_releases_request_after_async_load_failure() -> None: """Failed async loads should report completion and mark invalid blocks.""" connector = _AsyncLoadProbe() @@ -1096,7 +1472,6 @@ def test_shutdown_collects_failed_async_load_without_raising() -> None: assert connector._pending_loads == {} # noqa: SLF001 assert connector._invalid_load_block_ids == {4, 5} # noqa: SLF001 - assert connector.load_executor_shutdown assert connector.load_loop_stopped assert connector.store_loop_stopped assert connector.load_thread_joined @@ -2125,7 +2500,7 @@ def __init__(self) -> None: self._slot_size = self.slot_size self._store_staging_bytes = 4096 self._pending_store_staging_limit_bytes = 8192 - self._staging_pool = None + self._store_staging_pool = None self._pending_save_staging_bytes = 0 self._save_futures = [] self._rope_rotary_dim = 8 @@ -2166,7 +2541,7 @@ def __init__(self) -> None: self._slot_size = self.slot_size self._store_staging_bytes = 4096 self._pending_store_staging_limit_bytes = 8192 - self._staging_pool = None + self._store_staging_pool = None self._pending_save_staging_bytes = 0 self._save_futures = [] self._rope_rotary_dim = 8 @@ -2210,7 +2585,7 @@ def __init__(self) -> None: self._slot_size = self.slot_size self._store_staging_bytes = 4096 self._pending_store_staging_limit_bytes = 8192 - self._staging_pool = None + self._store_staging_pool = None self._pending_save_staging_bytes = 0 self._save_futures = [] self._rope_rotary_dim = 8 @@ -3648,12 +4023,12 @@ def __init__(self, total_memory: int) -> None: assert tight_pending == (8 << 30) // 5 -def test_cuda_staging_pool_reuses_preallocated_buffer(): - """Staging pool reuses its init-time allocation after release.""" - pool = CudaStagingPool( +def test_store_cuda_staging_pool_reuses_preallocated_buffer(): + """Store staging pool reuses its init-time allocation after release.""" + pool = StoreCudaStagingPool( device=torch.device("cpu"), - initial_bytes=128, - max_buffer_bytes=256, + buffer_bytes=128, + depth=1, ) lease = pool.acquire(64) diff --git a/tests/connector/test_gds_transfer.py b/tests/connector/test_gds_transfer.py index 4834fbf..9ab1110 100644 --- a/tests/connector/test_gds_transfer.py +++ b/tests/connector/test_gds_transfer.py @@ -84,3 +84,104 @@ def test_missing_file_raises(tmp_path): """FileNotFoundError raised when store file does not exist.""" with pytest.raises(FileNotFoundError): GDSTransferLayer(str(tmp_path / "nonexistent.store")) + + +def test_fixed_staging_pool_reuses_two_preallocated_buffers() -> None: + """Fixed staging pool reuses bounded preallocated buffers.""" + # First Party + from daser.connector.staging import FixedCudaStagingPool + + pool = FixedCudaStagingPool( + device=torch.device("cpu"), + buffer_bytes=16, + depth=2, + ) + + first = pool.acquire(8) + second = pool.acquire(16) + + with pytest.raises(RuntimeError, match="no fixed staging buffers available"): + pool.acquire(1) + + first.release() + third = pool.acquire(4) + + assert third.tensor.data_ptr() == first.tensor.data_ptr() + second.release() + third.release() + + +def test_fixed_staging_pool_can_block_until_buffer_release() -> None: + """Fixed staging pool can wait for a callback to release capacity.""" + # First Party + from daser.connector.staging import FixedCudaStagingPool + + pool = FixedCudaStagingPool( + device=torch.device("cpu"), + buffer_bytes=16, + depth=1, + ) + first = pool.acquire(8) + waits = 0 + + def release_first() -> None: + nonlocal waits + waits += 1 + first.release() + + second = pool.acquire(4, wait_for_release=release_first) + + assert waits == 1 + assert second.tensor.data_ptr() == first.tensor.data_ptr() + second.release() + + +def test_fixed_staging_pool_rejects_oversized_request() -> None: + """Fixed staging pool rejects requests larger than one buffer.""" + # First Party + from daser.connector.staging import FixedCudaStagingPool + + pool = FixedCudaStagingPool( + device=torch.device("cpu"), + buffer_bytes=16, + depth=2, + ) + + with pytest.raises(ValueError, match="exceeds fixed staging buffer"): + pool.acquire(17) + + +def test_depth_two_pipeline_consumes_completed_batch_first() -> None: + """Depth-two load pipeline reuses whichever buffer completes first.""" + # First Party + from daser.connector.worker import _run_fixed_depth_pipeline + + events: list[str] = [] + ready: dict[str, bool] = {"b0": False, "b1": True, "b2": True} + + def submit(batch: str, buffer_index: int) -> tuple[str, int]: + events.append(f"submit:{batch}:buf{buffer_index}") + return batch, buffer_index + + def consume(state: tuple[str, int]) -> int: + batch, buffer_index = state + events.append(f"consume:{batch}:buf{buffer_index}") + return buffer_index + + max_inflight = _run_fixed_depth_pipeline( + batches=["b0", "b1", "b2"], + submit=submit, + is_complete=lambda state: ready[state[0]], + consume=consume, + buffer_indices=[0, 1], + ) + + assert max_inflight == 2 + assert events == [ + "submit:b0:buf0", + "submit:b1:buf1", + "consume:b1:buf1", + "submit:b2:buf1", + "consume:b2:buf1", + "consume:b0:buf0", + ] diff --git a/tests/connector/test_ipc_client.py b/tests/connector/test_ipc_client.py index c76b4ba..1922677 100644 --- a/tests/connector/test_ipc_client.py +++ b/tests/connector/test_ipc_client.py @@ -370,6 +370,64 @@ async def fake_call(self, payload: dict) -> dict: assert recorded[1]["payload"]["allocation_offset"] == 2272 +@pytest.mark.asyncio +async def test_async_client_registers_and_loads_registered_staging_buffer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Async client can register fixed load staging and load by buffer index.""" + + recorded: list[dict] = [] + + async def fake_call(self, payload: dict) -> dict: + del self + recorded.append(payload) + return {"ok": True, "bytes": 16} + + monkeypatch.setattr(IPCClientAsync, "call", fake_call) + client = IPCClientAsync("/tmp/daser.sock") + + await client.register_load_staging_cuda( + buffer_index=1, + cuda_ipc_handle=b"h" * 64, + allocation_bytes=4096, + device_id=0, + device_ptr=223456, + allocation_base_ptr=221184, + allocation_offset=2272, + producer_pid=43, + ) + response = await client.transfer_load_registered_cuda( + buffer_index=1, + nbytes=16, + spans=[{"target_offset": 0, "nbytes": 16, "file_offset": 0}], + ) + + assert response == {"ok": True, "bytes": 16} + assert recorded == [ + { + "op": "register_load_staging", + "payload": { + "buffer_index": 1, + "cuda_ipc_handle": b"h" * 64, + "allocation_bytes": 4096, + "device_id": 0, + "device_ptr": 223456, + "allocation_base_ptr": 221184, + "allocation_offset": 2272, + "producer_pid": 43, + }, + }, + { + "op": "transfer_load", + "payload": { + "load_staging_buffer_index": 1, + "nbytes": 16, + }, + "spans": [{"target_offset": 0, "nbytes": 16, "file_offset": 0}], + }, + ] + + @pytest.mark.asyncio async def test_clients_transfer_drain(tmp_path): """Sync and async clients can drain server-owned transfer work.""" diff --git a/tests/server/test_ipc_server.py b/tests/server/test_ipc_server.py index 907528b..e783ecf 100644 --- a/tests/server/test_ipc_server.py +++ b/tests/server/test_ipc_server.py @@ -565,6 +565,115 @@ def fake_ensure_transfer(_server: IPCServer) -> FakeTransfer: assert opened_buffers[0].closed == 1 +@pytest.mark.asyncio +async def test_registered_load_staging_reuses_preopened_cuda_mapping( + tmp_path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Registered load staging buffers are opened once and reused by index.""" + + class FakeOpened: + def __init__(self) -> None: + self.array = bytearray(1024) + self.closed = 0 + + def close(self) -> None: + self.closed += 1 + + opened_buffers: list[FakeOpened] = [] + open_calls: list[dict[str, Any]] = [] + loaded_buffers: list[Any] = [] + + def fake_open_cuda_ipc_buffer(**kwargs: Any) -> FakeOpened: + open_calls.append(kwargs) + opened = FakeOpened() + opened_buffers.append(opened) + return opened + + class FakeTransfer(TransferLayer): + async def load_bytes(self, dst: Any, file_offset: int, nbytes: int) -> int: + return 0 + + async def store_bytes(self, src: Any, file_offset: int, nbytes: int) -> int: + return 0 + + async def load_bytes_grouped( + self, + dst: Any, + spans: list[dict[str, int]], + ) -> int: + loaded_buffers.append(dst) + return sum(int(span["nbytes"]) for span in spans) + + def close(self) -> None: + pass + + def fake_ensure_transfer(_server: IPCServer) -> FakeTransfer: + return FakeTransfer() + + monkeypatch.setattr( + "daser.server.ipc.server.open_cuda_ipc_buffer", + fake_open_cuda_ipc_buffer, + ) + monkeypatch.setattr( + "daser.server.ipc.server._CachedCudaArray.synchronize", + lambda _self: None, + ) + monkeypatch.setattr(IPCServer, "_ensure_transfer", fake_ensure_transfer) + + core = make_core() + server = IPCServer(str(tmp_path / "test.sock"), core, make_runtime_config(tmp_path)) + + await server.start() + try: + registered = await _send_recv( + str(tmp_path / "test.sock"), + { + "op": "register_load_staging", + "payload": { + "buffer_index": 1, + "cuda_ipc_handle": b"h" * 64, + "allocation_bytes": 1024, + "device_id": 0, + "device_ptr": 123456, + "allocation_base_ptr": 122880, + "allocation_offset": 576, + "producer_pid": 42, + }, + }, + ) + loaded = await _send_recv( + str(tmp_path / "test.sock"), + { + "op": "transfer_load", + "payload": { + "load_staging_buffer_index": 1, + "nbytes": 128, + }, + "spans": [{"target_offset": 0, "nbytes": 128, "file_offset": 0}], + }, + ) + finally: + await server.stop() + + assert registered == {"ok": True} + assert loaded["ok"] is True + assert loaded["bytes"] == 128 + assert len(opened_buffers) == 1 + assert len(loaded_buffers) == 1 + assert loaded_buffers[0][0:128] == opened_buffers[0].array[0:128] + assert open_calls == [ + { + "handle": b"h" * 64, + "nbytes": 1024, + "device_id": 0, + "local_ptr": None, + "allocation_offset": 576, + } + ] + assert opened_buffers[0].closed == 1 + + @pytest.mark.asyncio async def test_stop_accepting_closes_listener_before_transfer( tmp_path, diff --git a/tests/transfer/test_tiered_iouring_transfer.py b/tests/transfer/test_tiered_iouring_transfer.py index 812d0fe..adecced 100644 --- a/tests/transfer/test_tiered_iouring_transfer.py +++ b/tests/transfer/test_tiered_iouring_transfer.py @@ -94,6 +94,49 @@ def _read_l2_into( return super()._read_l2_into(file_offset, dst, uring) +class DelayedL2ReadProbe(TieredIOUringTransferLayer): + """Test transfer layer that pauses selected L2 reads.""" + + def __init__( + self, + path: str, + l1_bytes: int, + l2_bytes: int, + delayed_offsets: set[int], + ) -> None: + super().__init__( + path=path, + l1_bytes=l1_bytes, + l2_bytes=l2_bytes, + ) + self.delayed_offsets = delayed_offsets + self.release_read = threading.Event() + self.read_started = threading.Event() + self.copy_offsets: list[int] = [] + + def _read_l2_into( + self, + file_offset: int, + dst: object, + uring: NativeIOUring, + ) -> int: + """Pause configured reads before delegating to the real L2 reader.""" + if file_offset in self.delayed_offsets: + self.delayed_offsets.remove(file_offset) + self.read_started.set() + self.release_read.wait(timeout=5.0) + return super()._read_l2_into(file_offset, dst, uring) + + def _copy_grouped_to_dst( + self, + dst: object, + chunks: list[tuple[int, object, int, int]], + ) -> None: + """Record destination copy offsets before delegating.""" + self.copy_offsets.extend(int(chunk[0]) for chunk in chunks) + super()._copy_grouped_to_dst(dst, chunks) + + def _run(coro: object) -> object: """Run a coroutine on the current test event loop.""" return asyncio.get_event_loop().run_until_complete(coro) @@ -598,6 +641,64 @@ async def load_block(i: int) -> tuple[int, bytes]: _run(scenario()) +def test_iouring_l2_miss_batch_copies_each_read_as_it_completes(tmp_path) -> None: + """A completed L2 miss should copy to destination before slower reads finish.""" + + async def scenario() -> None: + path = str(tmp_path / "daser.store") + writer = TieredIOUringTransferLayer( + path=path, + l1_bytes=ALIGNMENT, + l2_bytes=ALIGNMENT * 2, + ) + try: + await writer.store_bytes(_block(b"a"), file_offset=0, nbytes=ALIGNMENT) + await writer.store_bytes( + _block(b"b"), + file_offset=ALIGNMENT, + nbytes=ALIGNMENT, + ) + await writer.drain() + finally: + writer.close() + + layer = DelayedL2ReadProbe( + path=path, + l1_bytes=ALIGNMENT * 2, + l2_bytes=ALIGNMENT * 2, + delayed_offsets={0}, + ) + try: + dst = bytearray(ALIGNMENT * 2) + task = asyncio.create_task( + layer.load_bytes_grouped( + dst, + [ + {"target_offset": 0, "file_offset": 0, "nbytes": ALIGNMENT}, + { + "target_offset": ALIGNMENT, + "file_offset": ALIGNMENT, + "nbytes": ALIGNMENT, + }, + ], + ) + ) + assert await asyncio.to_thread(layer.read_started.wait, timeout=2.0) + deadline = time.perf_counter() + 2.0 + while not layer.copy_offsets and time.perf_counter() < deadline: + await asyncio.sleep(0.01) + assert layer.copy_offsets == [ALIGNMENT] + layer.release_read.set() + loaded = await asyncio.wait_for(task, timeout=2.0) + assert loaded == ALIGNMENT * 2 + assert bytes(dst) == bytes(_block(b"a") + _block(b"b")) + finally: + layer.release_read.set() + layer.close() + + _run(scenario()) + + def test_iouring_store_returns_after_l1_before_l2_flush(tmp_path) -> None: """Store returns once L1 is readable while L2 persistence continues."""