diff --git a/benchmarks/utils/metrics.py b/benchmarks/utils/metrics.py index 540e172..6b277c1 100644 --- a/benchmarks/utils/metrics.py +++ b/benchmarks/utils/metrics.py @@ -7,110 +7,6 @@ from typing import Any -def correctness_check( - name: str, - cold_outputs: list[Any], - warm_outputs: list[Any], - prompts: list[list[int]], - max_num_seqs: int, -) -> dict[str, Any]: - """Compare cold vs warm generated output exactly. - - Args: - name: System label used in diagnostics. - cold_outputs: Outputs from the cold pass. - warm_outputs: Outputs from the warm pass. - prompts: Tokenized benchmark inputs in output order. - max_num_seqs: vLLM admission limit used for mismatch diagnostics. - - Returns: - Correctness counters. Only exact generated text and token-ID matches - are accepted. - - Thread-safety: - Pure function apart from reading output objects. - """ - del name - mismatches = 0 - mismatch_indices: list[int] = [] - mismatch_details: list[dict[str, Any]] = [] - prompt_alignment_mismatches = 0 - total = len(cold_outputs) - for i, (cold, warm) in enumerate(zip(cold_outputs, warm_outputs, strict=False)): - cold_prompt = list(getattr(cold, "prompt_token_ids", prompts[i])) - warm_prompt = list(getattr(warm, "prompt_token_ids", prompts[i])) - if cold_prompt != warm_prompt or cold_prompt != list(prompts[i]): - prompt_alignment_mismatches += 1 - if _generated_token_ids(cold) == _generated_token_ids(warm) and _output_text( - cold - ) == _output_text(warm): - continue - mismatches += 1 - mismatch_indices.append(i) - mismatch_details.append( - { - "index": i, - "wave": i // max(1, max_num_seqs), - "position": i % max(1, max_num_seqs), - "prompt_tokens": len(prompts[i]), - "cold_token_ids": _generated_token_ids(cold), - "warm_token_ids": _generated_token_ids(warm), - "cold_text": _output_text(cold), - "warm_text": _output_text(warm), - } - ) - return { - "mismatches": mismatches, - "total": total, - "indices": mismatch_indices, - "mismatch_details": mismatch_details, - "prompt_alignment_mismatches": prompt_alignment_mismatches, - } - - -def correctness_check_with_visibility( - name: str, - cold_outputs: list[Any], - warm_outputs: list[Any], - prompts: list[list[int]], - max_num_seqs: int, - visible_mask: list[bool], -) -> dict[str, Any]: - """Compare cold/warm outputs and split exact mismatches by visible hits. - - Args: - name: System label used in diagnostics. - cold_outputs: Outputs from the cold pass. - warm_outputs: Outputs from the warm pass. - prompts: Tokenized benchmark inputs in output order. - max_num_seqs: vLLM admission limit used for mismatch diagnostics. - visible_mask: Per-prompt cache-visibility mask. - - Returns: - Correctness counters including visible-hit mismatch counters. - - Thread-safety: - Pure function apart from reading output objects. - """ - result = correctness_check(name, cold_outputs, warm_outputs, prompts, max_num_seqs) - visible_total = 0 - visible_mismatches = 0 - for cold, warm, visible in zip( - cold_outputs, warm_outputs, visible_mask, strict=False - ): - if not visible: - continue - visible_total += 1 - if _generated_token_ids(cold) == _generated_token_ids(warm) and _output_text( - cold - ) == _output_text(warm): - continue - visible_mismatches += 1 - result["visible_mismatches"] = visible_mismatches - result["visible_total"] = visible_total - return result - - def contains_accuracy( results: list[Any], answers_by_id: dict[int, list[str]] ) -> float | None: @@ -363,17 +259,3 @@ def _first_number(value: dict[str, Any], keys: tuple[str, ...]) -> float | None: if isinstance(raw, int | float): return float(raw) return None - - -def _generated_token_ids(output: Any) -> list[int]: - """Return generated token IDs from a vLLM RequestOutput-like object.""" - if not getattr(output, "outputs", None): - return [] - return [int(token_id) for token_id in getattr(output.outputs[0], "token_ids", [])] - - -def _output_text(output: Any) -> str: - """Return generated text from a vLLM RequestOutput-like object.""" - if not getattr(output, "outputs", None): - return "" - return str(getattr(output.outputs[0], "text", "")) diff --git a/benchmarks/utils/servers.py b/benchmarks/utils/servers.py index 5baaa53..6e7df2f 100644 --- a/benchmarks/utils/servers.py +++ b/benchmarks/utils/servers.py @@ -399,25 +399,6 @@ def _daser_server_command(self) -> list[str]: cmd.extend(["--l2-size", str(self.l2_size_bytes)]) return cmd - async def stop_all(self) -> None: - """Terminate all child processes.""" - for proc in reversed(self._procs): - if proc.poll() is not None: - continue - proc.terminate() - deadline = time.monotonic() + 15.0 - for proc in self._procs: - if proc.poll() is not None: - continue - try: - proc.wait(timeout=max(0.1, deadline - time.monotonic())) - except subprocess.TimeoutExpired: - pass - for proc in self._procs: - if proc.poll() is None: - proc.kill() - self._procs.clear() - async def _start_vllm( self, log_name: str, diff --git a/daser/connector/daser_connector.py b/daser/connector/daser_connector.py index 659780b..1e2527b 100644 --- a/daser/connector/daser_connector.py +++ b/daser/connector/daser_connector.py @@ -145,7 +145,6 @@ def __init__( self._pending_loads: dict[str, Any] = {} self._invalid_load_block_ids: set[int] = set() 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() diff --git a/daser/connector/scheduler.py b/daser/connector/scheduler.py index 70f4057..f8ef8f9 100644 --- a/daser/connector/scheduler.py +++ b/daser/connector/scheduler.py @@ -815,18 +815,6 @@ def has_pending_store(self, req_id: str) -> bool: """ return req_id in self._pending_stores - def count_pending_stores_for_request(self, req_id: str) -> int: - """Return number of synthetic slot stores pending for a request. - - Args: - req_id: base vLLM request ID. - - Returns: - Count of pending slot-store entries. - """ - prefix = f"{req_id}:store:" - return len([key for key in self._pending_stores if key.startswith(prefix)]) - def drop_pending_alloc(self, req_id: str) -> None: """Remove pending allocation state for a request. diff --git a/daser/connector/staging.py b/daser/connector/staging.py index 4394394..4e20cbd 100644 --- a/daser/connector/staging.py +++ b/daser/connector/staging.py @@ -122,7 +122,6 @@ def __init__( raise ValueError("buffer_bytes must be positive") if depth <= 0: raise ValueError("depth must be positive") - self._device = device self._buffer_bytes = buffer_bytes self._buffers: list[torch.Tensor] = [ torch.empty(buffer_bytes, dtype=torch.uint8, device=device) @@ -438,41 +437,6 @@ def apply_rope_delta_to_kv_key_block( ) -def _transform_loaded_kv_batch( - layer_batch: torch.Tensor, - load_key_scale: float, - load_value_scale: float, - pos_offset: int, - rope_delta_scale: float, - rope_base: float, - rope_rotary_dim: int, - rope_is_neox_style: bool, -) -> None: - """Apply load-time scaling and RoPE relocation to a batch of KV blocks.""" - if layer_batch.dim() < 2 or layer_batch.shape[1] < 2: - return - if load_key_scale != 1.0: - layer_batch[:, 0].mul_(load_key_scale) - if load_value_scale != 1.0: - layer_batch[:, 1].mul_(load_value_scale) - if ( - not pos_offset - or layer_batch.dim() != 5 - or rope_rotary_dim <= 0 - or layer_batch.shape[-1] < rope_rotary_dim - ): - return - key_batch = layer_batch[:, 0].contiguous() - apply_rope_delta_to_key_block( - key_batch, - delta=round(pos_offset * rope_delta_scale), - rope_base=rope_base, - rotary_dim=rope_rotary_dim, - is_neox_style=rope_is_neox_style, - ) - layer_batch[:, 0].copy_(key_batch) - - def _transform_loaded_staging_batch( staging_by_layer: torch.Tensor, layer_sample: torch.Tensor, diff --git a/daser/connector/worker.py b/daser/connector/worker.py index 42e420a..78143e1 100644 --- a/daser/connector/worker.py +++ b/daser/connector/worker.py @@ -5,9 +5,6 @@ # 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 @@ -86,7 +83,6 @@ 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 @@ -94,98 +90,6 @@ _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. @@ -240,21 +144,6 @@ def _load_staging_pool_depth( 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 class _DeferredFinishedSave: """Store work held until vLLM reports a request as finished.""" @@ -1093,7 +982,6 @@ async def _run_load_request_dispatcher(self, sample_tensor: torch.Tensor) -> Non max_inflight=_LOAD_REQUEST_MAX_INFLIGHT, staging_depth=load_staging_pool.depth, ) - self._load_request_dispatcher = dispatcher queued: list[_QueuedLoadRequest] = [] active: list[_InflightRequestLoad] = [] try: @@ -1163,139 +1051,6 @@ async def _wait_for_dispatcher_completion( 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. - - Args: - 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, @@ -1334,69 +1089,6 @@ def _ensure_load_staging_pool( 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, diff --git a/tests/connector/test_daser_connector.py b/tests/connector/test_daser_connector.py index 376e781..d347c37 100644 --- a/tests/connector/test_daser_connector.py +++ b/tests/connector/test_daser_connector.py @@ -150,15 +150,19 @@ def _refresh_runtime_config(self) -> None: def _ensure_load_request_dispatcher(self, sample_tensor: torch.Tensor) -> None: """Drain queued requests synchronously for worker probe tests.""" + dispatcher = LoadRequestDispatcher(max_inflight=8, staging_depth=1) + queued = [] 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}, - ) + queued.append(item) + active = dispatcher.submit_ready(self, queued, sample_tensor) + while active: + consumed = dispatcher.consume_ready(self, active, sample_tensor) + if not consumed: + return + active.extend(dispatcher.submit_ready(self, queued, sample_tensor)) @property def transfer_ready(self): @@ -840,96 +844,6 @@ 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.""" @@ -965,118 +879,6 @@ def __init__(self) -> None: 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() @@ -1324,30 +1126,6 @@ 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) diff --git a/tests/connector/test_gds_transfer.py b/tests/connector/test_gds_transfer.py index 9ab1110..6627521 100644 --- a/tests/connector/test_gds_transfer.py +++ b/tests/connector/test_gds_transfer.py @@ -149,39 +149,3 @@ def test_fixed_staging_pool_rejects_oversized_request() -> None: 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/unit/test_benchmark_correctness.py b/tests/unit/test_benchmark_correctness.py index 75f79c5..45cbb14 100644 --- a/tests/unit/test_benchmark_correctness.py +++ b/tests/unit/test_benchmark_correctness.py @@ -7,89 +7,10 @@ # First Party from benchmarks.utils.metrics import ( contains_accuracy, - correctness_check, - correctness_check_with_visibility, request_text_exact_match, ) -def _output(token_id: int, text: str) -> SimpleNamespace: - return SimpleNamespace( - prompt_token_ids=[1, 2, 3], - outputs=[ - SimpleNamespace( - token_ids=[token_id], - text=text, - ) - ], - ) - - -def test_correctness_accepts_exact_text_and_tokens() -> None: - """Outputs pass when generated text and token IDs are identical.""" - result = correctness_check( - "test", - [_output(10, " yes")], - [_output(10, " yes")], - [[1, 2, 3]], - 64, - ) - - assert result["mismatches"] == 0 - assert result["total"] == 1 - assert result["indices"] == [] - - -def test_correctness_rejects_text_mismatch() -> None: - """Outputs fail when text differs even if token IDs match.""" - result = correctness_check( - "test", - [_output(10, " yes")], - [_output(10, " no")], - [[1, 2, 3]], - 64, - ) - - assert result["mismatches"] == 1 - assert result["indices"] == [0] - assert result["mismatch_details"][0]["cold_token_ids"] == [10] - assert result["mismatch_details"][0]["warm_token_ids"] == [10] - assert result["mismatch_details"][0]["cold_text"] == " yes" - assert result["mismatch_details"][0]["warm_text"] == " no" - - -def test_correctness_rejects_token_mismatch() -> None: - """Outputs fail when token IDs differ even if decoded text matches.""" - result = correctness_check( - "test", - [_output(10, " yes")], - [_output(11, " yes")], - [[1, 2, 3]], - 64, - ) - - assert result["mismatches"] == 1 - assert result["indices"] == [0] - assert result["mismatch_details"][0]["cold_token_ids"] == [10] - assert result["mismatch_details"][0]["warm_token_ids"] == [11] - - -def test_correctness_splits_visible_hit_mismatches() -> None: - """Visible-hit counters count exact mismatches only for visible prompts.""" - result = correctness_check_with_visibility( - "test", - [_output(10, " yes"), _output(20, " maybe")], - [_output(11, " yes"), _output(21, " no")], - [[1, 2, 3], [4, 5, 6]], - 64, - [True, False], - ) - - assert result["mismatches"] == 2 - assert result["visible_total"] == 1 - assert result["visible_mismatches"] == 1 - - def test_contains_accuracy_ignores_errors() -> None: """Answer containment ignores failed requests.""" results = [