From 4385e070124f47e4f523a81c5bcad1264fc6957f Mon Sep 17 00:00:00 2001 From: cytsoi <154662279+tcyjohn@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:50:35 +0800 Subject: [PATCH 1/2] Fence KV reuse and reloads by physical service and weight identity --- engine/cmlfq_backend.py | 98 ++++++-- engine/rollout_engine.py | 44 +++- infra/scheduling/kv_identity.py | 28 +++ infra/scheduling/vllm_cpu_offload.py | 33 ++- infra/sync/nccl_weight_sync.py | 21 +- infra/sync/prefix_cache_reset.py | 40 +++ pyproject.toml | 1 + requirements-ci.txt | 1 + scripts/restartable_vllm_server.py | 19 +- scripts/serving_protocol.py | 218 ++++++++++++++++ scripts/test_cpu_ci.sh | 5 + scripts/vllm_hot_reload_api_server.py | 104 +++++--- tests/test_checked_prefix_cache_reset.py | 52 ++++ tests/test_cmlfq_cost_scheduler.py | 11 +- tests/test_cpu_offload_backend.py | 15 ++ tests/test_nccl_store_lifecycle.py | 74 ++++++ tests/test_nccl_store_timeout.py | 54 ++++ tests/test_serving_protocol.py | 306 +++++++++++++++++++++++ tests/test_versioned_kv_compatibility.py | 223 +++++++++++++++++ tests/test_vllm_cpu_offload.py | 17 +- vllm_hot_reload.py | 32 ++- 21 files changed, 1328 insertions(+), 68 deletions(-) create mode 100644 infra/scheduling/kv_identity.py create mode 100644 infra/sync/prefix_cache_reset.py create mode 100644 scripts/serving_protocol.py create mode 100644 tests/test_checked_prefix_cache_reset.py create mode 100644 tests/test_nccl_store_lifecycle.py create mode 100644 tests/test_nccl_store_timeout.py create mode 100644 tests/test_serving_protocol.py create mode 100644 tests/test_versioned_kv_compatibility.py diff --git a/engine/cmlfq_backend.py b/engine/cmlfq_backend.py index 37f782d..ee6b470 100644 --- a/engine/cmlfq_backend.py +++ b/engine/cmlfq_backend.py @@ -14,6 +14,8 @@ from dataclasses import dataclass from typing import Any +from RL_Framework.infra.scheduling.kv_identity import same_snapshot, same_incarnation, snapshot_identity + logger = logging.getLogger(__name__) @@ -50,12 +52,28 @@ async def generate(self, source, target, prompt: str, path: str, request_id: str return BackendResult(await target.generate(prompt=prompt, **kwargs), path) if self.mode != "nixl": raise RuntimeError("transfer selected on a recompute-only backend") + leases = [] try: - # Materialize the exact next prompt on the source. Its generated - # token is discarded; the target samples the user's full output. - # Profile transfer_ms must include this preparation overhead. + # Acquisitions fail immediately when a publisher already owns a peer. + # Never wait for another engine while holding a generation lease. + for engine in dict.fromkeys((source, target)): + leases.append((engine, await engine.acquire_kv_lease())) + by_engine = {engine: lease for engine, lease in leases} + if not same_snapshot(by_engine[source], by_engine[target]): + raise RuntimeError("KV snapshot differs or is unknown") + # Strict workflows must materialize exactly the token IDs decoded on + # the target; tokenizing transcript text again can change its prefix. + preparation = {} + if kwargs.get("prompt_token_ids") is not None: + preparation["prompt_token_ids"] = kwargs["prompt_token_ids"] prepared = await source.generate( - prompt=prompt, max_new_tokens=1, temperature=0.0, top_p=1.0, n=1, + # Preparation also goes through the strict sampler hook. Its + # discarded token must use a seeded stochastic row, never the + # greedy fast path. Sampling does not change cached prompt KV. + prompt=prompt, max_new_tokens=1, + temperature=max(float(kwargs.get("temperature", 1.0)), 1e-4), + top_p=kwargs.get("top_p", 1.0), n=1, seed=kwargs.get("seed", 0), + kv_lease=by_engine[source], **preparation, kv_transfer_params={"do_remote_decode": True, "do_remote_prefill": False}, ) params = prepared.get("kv_transfer_params") @@ -67,14 +85,17 @@ async def generate(self, source, target, prompt: str, path: str, request_id: str raise RuntimeError("source exported no KV blocks") params = dict(params) params.update(do_remote_prefill=True, do_remote_decode=False) - output = await target.generate(prompt=prompt, kv_transfer_params=params, **kwargs) + output = await target.generate(prompt=prompt, kv_transfer_params=params, + kv_lease=by_engine[target], **kwargs) return BackendResult(output, "transfer") except Exception as error: - # Non-streaming requests: no partial output has reached the caller. - # Cancellation (BaseException) is deliberately not retried. logger.warning("CMLFQ KV transfer failed; recomputing: %s", error) - output = await target.generate(prompt=prompt, **kwargs) - return BackendResult(output, "recompute", type(error).__name__) + fallback = type(error).__name__ + finally: + await asyncio.gather(*(self._release_kv_lease(engine, lease) + for engine, lease in reversed(leases))) + output = await target.generate(prompt=prompt, **kwargs) + return BackendResult(output, "recompute", fallback) async def _generate_cpu(self, target, prompt, path, request_id, kwargs): if not request_id: @@ -83,32 +104,61 @@ async def _generate_cpu(self, target, prompt, path, request_id, kwargs): key = uuid.uuid4().hex params = {"libra_offload_key": key} fallback = "cache_unavailable" if path == "transfer" and previous is None else "" - wait_started = time.perf_counter() + leases = [] + lease = None status = None - if previous and path in {"stay", "transfer"}: - old_engine, descriptor = previous + wait_started = time.perf_counter() + try: try: - status = await old_engine.wait_for_cpu_offload( - descriptor["libra_offload_key"], descriptor["source_tp"], - ) - params["libra_load_key"] = descriptor["libra_offload_key"] + lease = await target.acquire_kv_lease() + leases.append((target, lease)) except Exception as error: + # Plain full-prompt generation remains available on old servers. fallback = type(error).__name__ - wait_ms = (time.perf_counter() - wait_started) * 1000 - try: + identity = snapshot_identity(lease or {}) + params["libra_kv_identity"] = identity + if previous and path in {"stay", "transfer"}: + old_engine, descriptor = previous + previous_identity = descriptor.get("libra_kv_identity") + if lease is not None and same_snapshot(previous_identity, identity): + try: + source_lease = lease if old_engine is target else await old_engine.acquire_kv_lease() + if old_engine is not target: + leases.append((old_engine, source_lease)) + if not same_incarnation(previous_identity, source_lease): + fallback = "incarnation_mismatch" + elif not same_snapshot(source_lease, identity): + fallback = "snapshot_mismatch" + else: + status = await old_engine.wait_for_cpu_offload( + descriptor["libra_offload_key"], descriptor["source_tp"], + ) + params["libra_load_key"] = descriptor["libra_offload_key"] + params["libra_kv_source_identity"] = snapshot_identity(source_lease) + except Exception as error: + fallback = type(error).__name__ + else: + fallback = "snapshot_mismatch" + wait_ms = (time.perf_counter() - wait_started) * 1000 + options = dict(kwargs, **({"kv_lease": lease} if lease is not None else {})) try: - output = await target.generate(prompt=prompt, kv_transfer_params=params, **kwargs) + output = await target.generate(prompt=prompt, kv_transfer_params=params, **options) except Exception as error: if "libra_load_key" not in params: raise fallback = type(error).__name__ params.pop("libra_load_key") - output = await target.generate(prompt=prompt, kv_transfer_params=params, **kwargs) + params.pop("libra_kv_source_identity", None) + output = await target.generate(prompt=prompt, kv_transfer_params=params, **options) except BaseException: self._release(target, key) raise + finally: + await asyncio.gather(*(self._release_kv_lease(engine, held) + for engine, held in reversed(leases))) descriptor = output.get("kv_transfer_params", {}) if descriptor.get("libra_offload_key"): + descriptor = dict(descriptor, libra_kv_identity=descriptor.get("libra_kv_identity", identity)) self._offloads[request_id] = (target, descriptor) else: self._offloads.pop(request_id, None) @@ -125,6 +175,14 @@ async def _generate_cpu(self, target, prompt, path, request_id, kwargs): fallback = "cache_prefix_miss" return BackendResult(output, actual, fallback) + async def _release_kv_lease(self, engine, lease): + try: + await engine.release_kv_lease(lease) + except Exception: + # Preserve generation failures/cancellation. The server keeps a + # failed release fenced and publication will time out fail-closed. + logger.exception("Failed to release KV weight lease %s", lease.get("lease_id")) + def _release(self, engine, key): async def cleanup(): try: diff --git a/engine/rollout_engine.py b/engine/rollout_engine.py index 91a7b4b..96d2c37 100644 --- a/engine/rollout_engine.py +++ b/engine/rollout_engine.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import uuid import logging import os import threading @@ -179,10 +180,18 @@ async def generate( ) -> dict[str, Any]: """Generate.""" await self._ensure_client() + lease = kwargs.pop("kv_lease", None) + prompt_token_ids = kwargs.pop("prompt_token_ids", None) + request_options = {} + kv_request_id = kwargs.pop("kv_request_id", None) + if kv_request_id is not None: + request_options["headers"] = {"X-Request-Id": str(kv_request_id)} + if lease is not None: + request_options.setdefault("headers", {})["X-Libra-KV-Lease"] = lease["lease_id"] payload = { "model": self.model_path, - "prompt": prompt, + "prompt": prompt_token_ids if prompt_token_ids is not None else prompt, "max_tokens": max_new_tokens, "temperature": temperature, "top_p": top_p, @@ -207,7 +216,7 @@ async def generate( try: response = await self.http_client.post( self.completions_url, - json=payload, + json=payload, **request_options, ) if response.status_code != 200: @@ -262,6 +271,37 @@ async def generate_batch( return valid_results + async def get_kv_cache_identity(self): + await self._ensure_client() + response = await self.http_client.get(f"{self.base_url}/libra/kv/identity") + response.raise_for_status() + return response.json() + + async def acquire_kv_lease(self): + await self._ensure_client() + lease_id = uuid.uuid4().hex + try: + response = await self.http_client.post( + f"{self.base_url}/libra/kv/lease", json={"lease_id": lease_id}, + ) + response.raise_for_status() + return response.json() + except BaseException: + # A server may reserve before its response is cancelled. The + # client-generated ID makes that reservation safe to withdraw. + try: + await self.http_client.delete(f"{self.base_url}/libra/kv/lease/{lease_id}") + except Exception: + logger.exception("Failed to release cancelled KV reservation %s", lease_id) + raise + + async def release_kv_lease(self, lease): + await self._ensure_client() + response = await self.http_client.delete( + f"{self.base_url}/libra/kv/lease/{lease['lease_id']}", + ) + response.raise_for_status() + async def wait_for_cpu_offload(self, key: str, source_tp: int, timeout: float = 30.0): """Await only the portion of offload not hidden by tool execution.""" await self._ensure_client() diff --git a/infra/scheduling/kv_identity.py b/infra/scheduling/kv_identity.py new file mode 100644 index 0000000..d956ca6 --- /dev/null +++ b/infra/scheduling/kv_identity.py @@ -0,0 +1,28 @@ +"""KV reuse requires known weight and service identities on both peers.""" + + +def same_snapshot(source, target): + return bool(isinstance(source, dict) and isinstance(target, dict) + and isinstance(source.get("snapshot_id"), str) and source["snapshot_id"] + and source["snapshot_id"] == target.get("snapshot_id") + and type(source.get("optimizer_revision")) is int + and type(target.get("optimizer_revision")) is int + and source["optimizer_revision"] >= 0 + and source["optimizer_revision"] == target["optimizer_revision"]) + + +def snapshot_identity(value): + return {name: value.get(name) for name in ( + "snapshot_id", "optimizer_revision", "server_instance_id", "server_generation")} + + +def same_incarnation(source, current): + """An old service generation cannot attest a cache after a reload/restart.""" + return bool(same_snapshot(source, current) + and isinstance(source.get("server_instance_id"), str) + and source["server_instance_id"] + and source["server_instance_id"] == current.get("server_instance_id") + and type(source.get("server_generation")) is int + and type(current.get("server_generation")) is int + and source["server_generation"] >= 0 + and source["server_generation"] == current["server_generation"]) diff --git a/infra/scheduling/vllm_cpu_offload.py b/infra/scheduling/vllm_cpu_offload.py index 220a3f1..491b3cc 100644 --- a/infra/scheduling/vllm_cpu_offload.py +++ b/infra/scheduling/vllm_cpu_offload.py @@ -17,6 +17,8 @@ from dataclasses import dataclass, field from pathlib import Path +from .kv_identity import same_snapshot, same_incarnation + import torch from vllm.distributed import get_tensor_model_parallel_rank from vllm.distributed.kv_transfer.kv_connector.v1.base import ( @@ -87,9 +89,11 @@ def request_finished(self, request, block_ids): "request_id": request.request_id, "key": key, "block_ids": list(block_ids)[:n // self.block_size], "tokens": list(request.all_token_ids[:n]), "source_tp": self.tp, + "libra_kv_identity": params.get("libra_kv_identity"), }) return True, {"libra_offload_key": key, "source_tp": self.tp, - "cached_tokens": n, "libra_imported_tokens": loaded} + "cached_tokens": n, "libra_imported_tokens": loaded, + "libra_kv_identity": params.get("libra_kv_identity")} def get_num_new_matched_tokens(self, request, num_computed_tokens): params = request.kv_transfer_params or {} @@ -99,11 +103,23 @@ def get_num_new_matched_tokens(self, request, num_computed_tokens): directory = cache_directory(self.root, key) try: manifest = json.loads((directory / "rank_0.json").read_text()) - if manifest.get("block_size") != self.block_size: + if (manifest.get("block_size") != self.block_size + or not same_snapshot(manifest.get("libra_kv_identity"), + params.get("libra_kv_identity")) + or not same_incarnation(manifest.get("libra_kv_identity"), + params.get("libra_kv_source_identity"))): return 0, False source_tp = int(manifest["source_tp"]) - if any(not (directory / f"rank_{rank}.json").exists() for rank in range(source_tp)): + if not 1 <= source_tp <= 64: return 0, False + for rank in range(source_tp): + shard = json.loads((directory / f"rank_{rank}.json").read_text()) + if (shard.get("tokens") != manifest["tokens"] + or shard.get("source_tp") != source_tp + or shard.get("block_size") != self.block_size + or not same_incarnation(shard.get("libra_kv_identity"), + manifest.get("libra_kv_identity"))): + return 0, False tokens = manifest["tokens"] n = min(len(tokens), len(request.prompt_token_ids) - 1) n = n // self.block_size * self.block_size @@ -114,9 +130,11 @@ def get_num_new_matched_tokens(self, request, num_computed_tokens): self.matches[request.request_id] = { "key": key, "source_tp": source_tp, "num_tokens": n, "request_id": request.request_id, + "libra_kv_identity": params.get("libra_kv_identity"), + "libra_kv_source_identity": params.get("libra_kv_source_identity"), } return n - num_computed_tokens, False - except (FileNotFoundError, json.JSONDecodeError): + except (FileNotFoundError, json.JSONDecodeError, KeyError, TypeError, ValueError): return 0, False def update_state_after_alloc(self, request, blocks, num_external_tokens): @@ -193,7 +211,7 @@ def _publish(self, host, item, started, dma_done): temporary.replace(directory / f"rank_{self.rank}.pt") atomic_json(directory / f"rank_{self.rank}.json", { "tokens": item["tokens"], "source_tp": item["source_tp"], - "block_size": self.block_size, + "block_size": self.block_size, "libra_kv_identity": item.get("libra_kv_identity"), "offload_started_at": started, "dma_done_at": dma_done, "ready_at": time.time(), "bytes": sum(t.numel() * t.element_size() for t in host.values()), }) @@ -203,6 +221,11 @@ def _publish(self, host, item, started, dma_done): def _reload(self, item): directory = cache_directory(self.root, item["key"]) + for rank in range(item["source_tp"]): + manifest = json.loads((directory / f"rank_{rank}.json").read_text()) + if (not same_incarnation(manifest.get("libra_kv_identity"), item.get("libra_kv_source_identity")) + or not same_snapshot(manifest.get("libra_kv_identity"), item.get("libra_kv_identity"))): + raise RuntimeError("CPU KV identity changed before worker reload") source_heads = self.num_heads // item["source_tp"] target_heads = self.num_heads // self.tp first_head, last_head = self.rank * target_heads, (self.rank + 1) * target_heads diff --git a/infra/sync/nccl_weight_sync.py b/infra/sync/nccl_weight_sync.py index 1844ce3..d289622 100644 --- a/infra/sync/nccl_weight_sync.py +++ b/infra/sync/nccl_weight_sync.py @@ -4,6 +4,7 @@ import os import ipaddress +import math import threading import time from dataclasses import dataclass @@ -15,6 +16,10 @@ _COMMUNICATOR_CACHE: dict[tuple[str, int, int, int, str], object] = {} _COMMUNICATOR_CACHE_LOCK = threading.Lock() _PIPELINE_DEPTH = 2 +# vLLM constructs datetime.timedelta(seconds=store_timeout), so None/inf and +# timedelta.max are unsafe. A century represents the application's unbounded +# wait while staying below signed int64 even after conversion to nanoseconds. +_UNBOUNDED_STORE_TIMEOUT_S = 100 * 365 * 24 * 60 * 60 @dataclass(frozen=True) @@ -44,13 +49,25 @@ def __post_init__(self): def _group(spec: NcclReloadSpec): from vllm.distributed.utils import StatelessProcessGroup - return StatelessProcessGroup.create( + timeout_s = float(spec.timeout_s) + if not math.isfinite(timeout_s) or timeout_s < 0: + raise ValueError("NCCL store timeout must be finite non-negative seconds") + group = StatelessProcessGroup.create( host=spec.host, port=int(spec.port), rank=int(spec.rank), world_size=int(spec.world_size), - store_timeout=max(30, int(spec.timeout_s)), + store_timeout=timeout_s if timeout_s > 0 else _UNBOUNDED_STORE_TIMEOUT_S, ) + # vLLM 0.9.2 keeps a Python socket for the master_listen_fd handed to + # TCPStore. TCPStore owns and closes that descriptor. Letting the Python + # socket also close it can abort the store's daemon during cache eviction + # or interpreter shutdown. Transfer ownership while both objects are alive; + # receiver groups have no listening socket. + listen_socket = getattr(group, "socket", None) + if spec.rank == 0 and listen_socket is not None: + listen_socket.detach() + return group def _communicator(spec: NcclReloadSpec): diff --git a/infra/sync/prefix_cache_reset.py b/infra/sync/prefix_cache_reset.py new file mode 100644 index 0000000..b7b80c8 --- /dev/null +++ b/infra/sync/prefix_cache_reset.py @@ -0,0 +1,40 @@ +"""Opt-in V1 utility preserving the scheduler's prefix-reset success signal.""" + + +def install_checked_prefix_cache_reset(core_class=None): + """Install in each spawned core process, never in installed source files. + + vLLM 0.9.2's EngineCore and AsyncLLM reset wrappers discard the scheduler's + boolean return. A distinct utility makes missing child-process installation + an explicit RPC error instead of silently accepting a discarded failure. + """ + if core_class is None: + import vllm + if vllm.__version__.split('+', 1)[0] != '0.9.2': + raise RuntimeError('Checked prefix reset requires vLLM 0.9.2') + from vllm.v1.engine.core import EngineCore + core_class = EngineCore + if getattr(core_class, '_libra_checked_prefix_reset', False): + return + + def reset(self): + if self.scheduler.reset_prefix_cache() is not True: + raise RuntimeError('Prefix cache reset failed: active KV blocks remain') + return True + + core_class.libra_reset_prefix_cache = reset + core_class._libra_checked_prefix_reset = True + + +async def reset_prefix_cache_verified(client, *, require_verified=False): + if require_verified: + core_client = getattr(client, 'engine_core', None) + if core_client is None or not hasattr(core_client, 'call_utility_async'): + raise RuntimeError('Checked prefix reset requires the V1 utility client') + if await core_client.call_utility_async('libra_reset_prefix_cache') is not True: + raise RuntimeError('Prefix cache reset utility did not attest success') + return True + result = await client.reset_prefix_cache() + if result is False: + raise RuntimeError('Prefix cache reset failed') + return True if result is True else None diff --git a/pyproject.toml b/pyproject.toml index 5b6eab7..1389452 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,7 @@ packages = [ "RL_Framework.infra.scheduling", "RL_Framework.infra.sync", "RL_Framework.launcher", + "RL_Framework.scripts", "RL_Framework.trainer", "RL_Framework.workflow", ] diff --git a/requirements-ci.txt b/requirements-ci.txt index 36bfd40..7686332 100644 --- a/requirements-ci.txt +++ b/requirements-ci.txt @@ -4,3 +4,4 @@ pytest-asyncio==0.26.0 numpy==1.26.4 PyYAML==6.0.2 httpx==0.28.1 +fastapi # Already required by the vLLM runtime; HTTP lifecycle tests need it. diff --git a/scripts/restartable_vllm_server.py b/scripts/restartable_vllm_server.py index f41af8d..a2c2f55 100644 --- a/scripts/restartable_vllm_server.py +++ b/scripts/restartable_vllm_server.py @@ -44,8 +44,8 @@ def parse_args() -> argparse.Namespace: def wait_until_ready(process: subprocess.Popen, url: str, timeout: float) -> None: - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: + deadline = time.monotonic() + timeout if timeout > 0 else None + while deadline is None or time.monotonic() < deadline: if process.poll() is not None: raise RuntimeError(f"vLLM exited with status {process.returncode}") try: @@ -70,7 +70,7 @@ def reload_inplace(health_url: str, checkpoint_path: str, timeout: float) -> dic headers={"Content-Type": "application/json"}, method="POST", ) - with urllib.request.urlopen(request, timeout=timeout + 10) as response: + with urllib.request.urlopen(request, timeout=timeout + 10 if timeout > 0 else None) as response: if response.status != 200: raise RuntimeError(f"vLLM reload returned HTTP {response.status}") return json.loads(response.read().decode("utf-8")) @@ -89,7 +89,7 @@ def reload_nccl( headers={"Content-Type": "application/json"}, method="POST", ) - with urllib.request.urlopen(request, timeout=timeout + 10) as response: + with urllib.request.urlopen(request, timeout=timeout + 10 if timeout > 0 else None) as response: if response.status != 200: raise RuntimeError(f"vLLM NCCL reload returned HTTP {response.status}") return json.loads(response.read().decode("utf-8")) @@ -133,6 +133,7 @@ def main() -> int: control_dir.mkdir(parents=True, exist_ok=True) request_path = control_dir / "reload_request.json" current_version = 0 + current_snapshot = None process: subprocess.Popen | None = None def request_shutdown(signum, _frame): @@ -178,7 +179,10 @@ def launch(model_path: str) -> subprocess.Popen: except (OSError, ValueError, KeyError, json.JSONDecodeError): time.sleep(max(0.05, args.poll_interval)) continue - if requested_version > current_version: + requested_snapshot = request.get('snapshot_id') + if (requested_version > current_version or + requested_version == current_version and requested_snapshot + and requested_snapshot != current_snapshot): if reload_mode == "checkpoint" and not Path(checkpoint_path).is_dir(): raise FileNotFoundError(checkpoint_path) current_version = requested_version @@ -205,6 +209,8 @@ def launch(model_path: str) -> subprocess.Popen: ), "timeout_seconds": float(args.ready_timeout), "chunk_bytes": int(request["nccl_chunk_bytes"]), + **{key: request[key] for key in ('snapshot_id', 'optimizer_revision', 'expected_tensor_count') + if key in request}, }, args.ready_timeout, ) @@ -226,6 +232,7 @@ def launch(model_path: str) -> subprocess.Popen: { "instance_id": args.instance_id, "version": current_version, + **{key: request[key] for key in ('snapshot_id', 'optimizer_revision', 'topology_epoch') if key in request}, "checkpoint_path": checkpoint_path, "reload_strategy": reload_strategy, "reload_method": reload_method, @@ -241,6 +248,7 @@ def launch(model_path: str) -> subprocess.Popen: { "instance_id": args.instance_id, "version": current_version, + **{key: request[key] for key in ('snapshot_id', 'optimizer_revision', 'topology_epoch') if key in request}, "checkpoint_path": checkpoint_path, "reload_strategy": reload_strategy, "reload_method": reload_method, @@ -256,6 +264,7 @@ def launch(model_path: str) -> subprocess.Popen: "server_response": response, }, ) + current_snapshot = requested_snapshot print( f"[vllm-reloader] ready version={current_version} " f"method={reload_method} strategy={reload_strategy} " diff --git a/scripts/serving_protocol.py b/scripts/serving_protocol.py new file mode 100644 index 0000000..c438c7c --- /dev/null +++ b/scripts/serving_protocol.py @@ -0,0 +1,218 @@ +"""Shared serving identity, request/KV leases and global reload receipts. + +Generation and reload share one boundary so cached KV remains tied to the +weights and service incarnation that produced it. +""" +from __future__ import annotations + +import asyncio +import hashlib +import json +import time +import uuid + + +class ServingState: + def __init__(self, tp_size, model_path=''): + if int(tp_size) < 1: + raise ValueError('Expected positive physical TP worker count') + self.tp_size = int(tp_size) + self.server_instance_id = uuid.uuid4().hex + self.server_generation = 0 + self.snapshot_id = 'initial:' + hashlib.sha256(str(model_path).encode()).hexdigest() + self.optimizer_revision = 0 + self.workers = [] + self.active_requests = 0 + self.kv_leases = set() + self.closed = False + self.failed = None + self.condition = asyncio.Condition() + + def status(self): + return dict(server_instance_id=self.server_instance_id, server_generation=self.server_generation, + snapshot_id=self.snapshot_id, optimizer_revision=self.optimizer_revision, + tp_size=self.tp_size, workers=self.workers, active_requests=self.active_requests, + kv_leases=len(self.kv_leases), closed=self.closed, failed=self.failed, + server_generation_verified=self.failed is None and not self.closed, + server_monotonic=time.monotonic()) + + async def acquire_kv_lease(self, lease_id=None): + async with self.condition: + if self.closed or self.failed: + raise RuntimeError('Rollout service is unavailable for a new KV lease') + lease_id = lease_id or uuid.uuid4().hex + self.kv_leases.add(lease_id) + return dict(self.status(), lease_id=lease_id) + + async def release_kv_lease(self, lease_id): + async with self.condition: + self.kv_leases.discard(lease_id) + self.condition.notify_all() + + async def enter(self, lease_id=None): + async with self.condition: + if lease_id is not None and lease_id not in self.kv_leases: + raise RuntimeError('Unknown or expired KV lease') + await self.condition.wait_for(lambda: not self.closed or lease_id in self.kv_leases or self.failed) + if self.failed: + raise RuntimeError('Rollout service is quarantined: ' + self.failed) + self.active_requests += 1 + return self.server_generation + + async def leave(self): + async with self.condition: + self.active_requests -= 1 + if self.active_requests < 0: + raise RuntimeError('Unbalanced serving request lifetime') + self.condition.notify_all() + + async def begin_reload(self, timeout): + async with self.condition: + if self.closed or self.failed: + raise RuntimeError('Rollout service is unavailable for reload') + self.closed = True + wait = self.condition.wait_for(lambda: self.active_requests == 0 and not self.kv_leases) + try: + if timeout > 0: + await asyncio.wait_for(wait, timeout) + else: + await wait + except BaseException: + # No mutation occurred before the drain; retry is safe. + self.closed = False + self.condition.notify_all() + raise + + def install(self, *, snapshot_id, optimizer_revision, workers, rank_offset=0, expected_tensor_count=None): + if not self.closed: + raise RuntimeError('Reload has no exclusive drained service ownership') + ranks = [worker.get('rank') for worker in workers] + if (len(workers) != self.tp_size or sorted(ranks) != list(range(self.tp_size)) + or any(worker.get('service_generation') != self.server_generation + 1 for worker in workers)): + raise RuntimeError('All TP workers must attest the next service generation') + if snapshot_id is not None: + counts = {worker.get('received_tensors') for worker in workers} + count = expected_tensor_count + if count is None and len(counts) == 1: + count = next(iter(counts)) + if (not snapshot_id or type(optimizer_revision) is not int or optimizer_revision < self.optimizer_revision + or type(count) is not int or count < 1 + or sorted(worker.get('nccl_rank', -1) for worker in workers) != list(range(1+rank_offset, 1+rank_offset+self.tp_size)) + or any(worker.get('receive_complete') is not True or worker.get('received_tensors') != count + or worker.get('loaded_parameters', 0) < 1 for worker in workers)): + raise RuntimeError('Missing complete immutable-snapshot tensor receipt on a TP worker') + self.server_generation += 1 + self.snapshot_id, self.optimizer_revision = snapshot_id, optimizer_revision if snapshot_id else 0 + self.workers = [dict(worker) for worker in workers] + + async def reopen(self): + async with self.condition: + if self.failed: + raise RuntimeError('A failed service cannot be reopened') + self.closed = False + self.condition.notify_all() + + async def fail(self, error): + async with self.condition: + self.failed = str(error) + self.closed = True + self.condition.notify_all() + + +def generation_headers(state, request_id, generation, status, streaming=False): + verified = bool(state.failed is None and generation == state.server_generation + and 200 <= status < 300 and not streaming) + values = { + 'x-libra-request-id': request_id, + 'x-libra-service-instance': state.server_instance_id, + 'x-libra-service-generation': str(state.server_generation), + 'x-libra-service-generation-start': str(generation), + 'x-libra-service-generation-end': str(state.server_generation), + 'x-libra-service-generation-verified': '1' if verified else '0', + 'x-libra-reload-overlap': '0' if generation == state.server_generation else '1', + } + return [(key.encode(), value.encode()) for key, value in values.items()] + + +def fence_kv_payload(payload, state): + """The service owns target weight identity; stale imports must recompute.""" + from RL_Framework.infra.scheduling.kv_identity import same_snapshot, snapshot_identity + params = payload.get('kv_transfer_params') if isinstance(payload, dict) else None + if isinstance(params, dict) and ('libra_offload_key' in params or 'libra_load_key' in params): + identity = snapshot_identity(state.status()) + if not same_snapshot(params.get('libra_kv_identity'), identity): + params.pop('libra_load_key', None) + params.pop('libra_kv_source_identity', None) + params['libra_kv_identity'] = identity + return payload + + +class GenerationMiddleware: + """Admit request lifetimes through the same boundary used by weight reload.""" + def __init__(self, app, state): + self.app, self.state = app, state + + async def __call__(self, scope, receive, send): + if (scope['type'] != 'http' or scope.get('method') != 'POST' + or scope.get('path') not in {'/v1/completions', '/v1/chat/completions'}): + return await self.app(scope, receive, send) + headers = dict(scope.get('headers', [])) + lease = headers.get(b'x-libra-kv-lease') + request_id = headers.get(b'x-libra-request-id', uuid.uuid4().hex.encode()).decode() + try: + generation = await self.state.enter(lease.decode() if lease else None) + except RuntimeError as exc: + body = json.dumps(dict(detail=str(exc))).encode() + await send(dict(type='http.response.start', status=409, headers=[(b'content-type', b'application/json')])) + return await send(dict(type='http.response.body', body=body)) + start, bodies, streaming = None, [], False + async def observe(message): + nonlocal start, streaming + if message['type'] == 'http.response.start': + start = dict(message) + streaming = b'text/event-stream' in dict(message.get('headers', [])).get(b'content-type', b'') + if streaming: + start['headers'] = list(start.get('headers', [])) + generation_headers( + self.state, request_id, generation, start['status'], streaming=True) + await send(start) + return + if message['type'] != 'http.response.body' or streaming: + return await send(message) + bodies.append(message) + if not message.get('more_body', False): + start['headers'] = list(start.get('headers', [])) + generation_headers( + self.state, request_id, generation, start['status']) + await send(start) + for part in bodies: + await send(part) + try: + if receive is None: + return await self.app(scope, receive, observe) + chunks = [] + while True: + message = await receive() + if message['type'] == 'http.disconnect': + return + chunks.append(message.get('body', b'')) + if not message.get('more_body', False): + break + body = b''.join(chunks) + try: + payload = json.loads(body) + except (ValueError, TypeError): + payload = None + if isinstance(payload, dict) and isinstance(payload.get('kv_transfer_params'), dict): + body = json.dumps(fence_kv_payload(payload, self.state)).encode() + scope = dict(scope, headers=[(key, value) for key, value in scope.get('headers', []) + if key.lower() != b'content-length'] + + [(b'content-length', str(len(body)).encode())]) + consumed = False + async def replay(): + nonlocal consumed + if not consumed: + consumed = True + return dict(type='http.request', body=body, more_body=False) + return await receive() + await self.app(scope, replay, observe) + finally: + await self.state.leave() diff --git a/scripts/test_cpu_ci.sh b/scripts/test_cpu_ci.sh index c3ceb95..e9726e8 100644 --- a/scripts/test_cpu_ci.sh +++ b/scripts/test_cpu_ci.sh @@ -11,6 +11,11 @@ python -m pytest -q --strict-markers --junitxml=reports/cpu-tests.xml \ tests/test_cmlfq_cost_scheduler.py \ tests/test_hetero_cmlfq_integration.py \ tests/test_rollout_engine.py \ + tests/test_checked_prefix_cache_reset.py \ + tests/test_nccl_store_lifecycle.py \ + tests/test_nccl_store_timeout.py \ + tests/test_serving_protocol.py \ + tests/test_versioned_kv_compatibility.py \ tests/test_cpu_offload_backend.py \ tests/test_r2e_patch_harness.py \ tests/test_r2e_gym_workflow.py diff --git a/scripts/vllm_hot_reload_api_server.py b/scripts/vllm_hot_reload_api_server.py index 6f5ba7d..fc5f6d8 100644 --- a/scripts/vllm_hot_reload_api_server.py +++ b/scripts/vllm_hot_reload_api_server.py @@ -18,6 +18,13 @@ from vllm.entrypoints.openai.cli_args import make_arg_parser from vllm.entrypoints.openai.cli_args import validate_parsed_serve_args from vllm.utils import FlexibleArgumentParser +from RL_Framework.infra.sync.prefix_cache_reset import ( + install_checked_prefix_cache_reset, reset_prefix_cache_verified, +) +from RL_Framework.scripts.serving_protocol import ServingState, GenerationMiddleware + +# EngineCore subprocesses import this project entrypoint before utility RPCs. +install_checked_prefix_cache_reset() class ReloadWeightsRequest(BaseModel): @@ -36,12 +43,39 @@ class ReloadWeightsNcclRequest(BaseModel): rank_offset: int timeout_seconds: float = 1200.0 chunk_bytes: int = 268435456 + snapshot_id: str | None = None + optimizer_revision: int | None = None + expected_tensor_count: int | None = None _reload_lock = asyncio.Lock() _custom_router = APIRouter() +@_custom_router.get('/publication/status') +@_custom_router.get('/libra/kv/identity') +async def publication_status(request: Request): + return request.app.state.serving_state.status() + + +@_custom_router.post('/libra/kv/lease') +async def acquire_kv_lease(request: Request): + body = await request.body() + lease_id = json.loads(body).get('lease_id') if body else None + if lease_id is not None and (not isinstance(lease_id, str) or not re.fullmatch('[a-f0-9]{32}', lease_id)): + raise HTTPException(status_code=400, detail='Invalid KV lease ID') + try: + return await request.app.state.serving_state.acquire_kv_lease(lease_id) + except RuntimeError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc + + +@_custom_router.delete('/libra/kv/lease/{lease_id}') +async def release_kv_lease(lease_id: str, request: Request): + await request.app.state.serving_state.release_kv_lease(lease_id) + return {'released': True} + + def _kv_directory(key: str) -> Path: if not re.fullmatch(r"[a-f0-9]{32}", key): raise HTTPException(status_code=400, detail="invalid KV key") @@ -88,22 +122,24 @@ async def reload_weights(payload: ReloadWeightsRequest, request: Request): async with _reload_lock: started = time.perf_counter() + state = request.app.state.serving_state + await state.begin_reload(payload.timeout_seconds) client = api_server.engine_client(request) try: results = await client.collective_rpc( - "reload_weights", - timeout=payload.timeout_seconds, + "reload_weights", timeout=payload.timeout_seconds if payload.timeout_seconds > 0 else None, args=(str(checkpoint),), ) - # Prefix-cache entries encode outputs from the previous parameters. - await client.reset_prefix_cache() - except Exception as exc: - raise HTTPException(status_code=500, detail=str(exc)) from exc - return { - "checkpoint_path": str(checkpoint), - "total_seconds": time.perf_counter() - started, - "workers": results, - } + cleared = await reset_prefix_cache_verified(client, + require_verified=os.environ.get('LIBRA_CHECK_PREFIX_CACHE') == '1') + state.install(snapshot_id=None, optimizer_revision=None, workers=results) + await state.reopen() + except BaseException as exc: + await state.fail(repr(exc)) + raise + return dict(checkpoint_path=str(checkpoint), total_seconds=time.perf_counter()-started, + workers=results, prefix_cache_cleared=cleared is True, + service_generation=state.server_generation, service_instance_id=state.server_instance_id) @_custom_router.post("/reload_weights_nccl") @@ -114,27 +150,32 @@ async def reload_weights_nccl( """Receive the next model directly from the Megatron GPU stream.""" async with _reload_lock: started = time.perf_counter() + state = request.app.state.serving_state + if ((payload.snapshot_id is None) != (payload.optimizer_revision is None) + or payload.optimizer_revision is not None and payload.optimizer_revision < state.optimizer_revision): + raise HTTPException(status_code=409, detail='Missing or regressing snapshot identity') + await state.begin_reload(payload.timeout_seconds) client = api_server.engine_client(request) try: results = await client.collective_rpc( - "reload_weights_nccl", - timeout=payload.timeout_seconds, - args=( - payload.host, - payload.port, - payload.world_size, - payload.rank_offset, - payload.timeout_seconds, - payload.chunk_bytes, - ), + "reload_weights_nccl", timeout=payload.timeout_seconds if payload.timeout_seconds > 0 else None, + args=(payload.host, payload.port, payload.world_size, payload.rank_offset, + payload.timeout_seconds, payload.chunk_bytes) + + ((payload.expected_tensor_count,) if payload.expected_tensor_count is not None else ()), ) - await client.reset_prefix_cache() - except Exception as exc: - raise HTTPException(status_code=500, detail=str(exc)) from exc - return { - "total_seconds": time.perf_counter() - started, - "workers": results, - } + cleared = await reset_prefix_cache_verified(client, + require_verified=bool(payload.snapshot_id) or os.environ.get('LIBRA_CHECK_PREFIX_CACHE') == '1') + state.install(snapshot_id=payload.snapshot_id, optimizer_revision=payload.optimizer_revision, + workers=results, rank_offset=payload.rank_offset, + expected_tensor_count=payload.expected_tensor_count) + await state.reopen() + except BaseException as exc: + await state.fail(repr(exc)) + raise + return dict(total_seconds=time.perf_counter()-started, workers=results, + prefix_cache_cleared=cleared is True, service_generation=state.server_generation, + service_instance_id=state.server_instance_id, snapshot_id=state.snapshot_id, + optimizer_revision=state.optimizer_revision) # vLLM 0.9.x constructs a fresh FastAPI app in build_app() and includes its @@ -144,9 +185,16 @@ async def reload_weights_nccl( _build_app = api_server.build_app +def make_serving_state(args): + """Build the service state shared by generation and weight reload.""" + return ServingState(args.tensor_parallel_size, args.model) + + def build_app_with_reload_routes(args): app = _build_app(args) + app.state.serving_state = make_serving_state(args) app.include_router(_custom_router) + app.add_middleware(GenerationMiddleware, state=app.state.serving_state) return app diff --git a/tests/test_checked_prefix_cache_reset.py b/tests/test_checked_prefix_cache_reset.py new file mode 100644 index 0000000..5ca224e --- /dev/null +++ b/tests/test_checked_prefix_cache_reset.py @@ -0,0 +1,52 @@ +"""A failed or missing V1 utility cannot produce a cache-clear ACK.""" + +import asyncio +from types import SimpleNamespace + +import pytest + +from RL_Framework.infra.sync.prefix_cache_reset import ( + install_checked_prefix_cache_reset, reset_prefix_cache_verified, +) + + +@pytest.mark.parametrize('success', [True, False]) +def test_checked_core_utility_preserves_scheduler_outcome(success): + class Core: + def reset_prefix_cache(self): + # Mirrors vLLM 0.9.2's original wrapper, which discards the bool. + self.scheduler.reset_prefix_cache() + + original = Core.reset_prefix_cache + install_checked_prefix_cache_reset(Core) + core = Core() + core.scheduler = SimpleNamespace(reset_prefix_cache=lambda: success) + assert Core.reset_prefix_cache is original + if success: + assert core.libra_reset_prefix_cache() is True + else: + with pytest.raises(RuntimeError, match='active KV blocks'): + core.libra_reset_prefix_cache() + + +@pytest.mark.parametrize('outcome', [True, False, 'missing']) +def test_verified_reset_accepts_only_the_real_utility_success(outcome): + async def utility(name): + assert name == 'libra_reset_prefix_cache' + if outcome == 'missing': + raise AttributeError('Unknown utility method') + return outcome + + client = SimpleNamespace(engine_core=SimpleNamespace(call_utility_async=utility)) + if outcome is True: + assert asyncio.run(reset_prefix_cache_verified(client, require_verified=True)) is True + else: + with pytest.raises((RuntimeError, AttributeError)): + asyncio.run(reset_prefix_cache_verified(client, require_verified=True)) + + +def test_legacy_none_return_is_preserved_without_claiming_verified_reset(): + async def reset(): + return None + + assert asyncio.run(reset_prefix_cache_verified(SimpleNamespace(reset_prefix_cache=reset))) is None diff --git a/tests/test_cmlfq_cost_scheduler.py b/tests/test_cmlfq_cost_scheduler.py index a30505c..461cedd 100644 --- a/tests/test_cmlfq_cost_scheduler.py +++ b/tests/test_cmlfq_cost_scheduler.py @@ -205,8 +205,17 @@ def engine(self, handler, transfer=True): "nixl" if transfer else "recompute", [(1, 4)], ) engine.engines = [VLLMRolloutEngine(host="host-a"), VLLMRolloutEngine(host="host-b")] + def versioned_handler(request): + if request.url.path == "/libra/kv/lease": + return httpx.Response(200, json=dict( + lease_id=json.loads(request.content)["lease_id"], snapshot_id="shared-snapshot", + optimizer_revision=0, server_instance_id=request.url.host, server_generation=0, + )) + if request.method == "DELETE" and request.url.path.startswith("/libra/kv/lease/"): + return httpx.Response(200, json={"released": True}) + return handler(request) for backend in engine.engines: - backend.http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + backend.http_client = httpx.AsyncClient(transport=httpx.MockTransport(versioned_handler)) backend._http_client_loop = asyncio.get_running_loop() engine._configure_cost_runtime() return engine diff --git a/tests/test_cpu_offload_backend.py b/tests/test_cpu_offload_backend.py index c7dbc7e..7e6c0fc 100644 --- a/tests/test_cpu_offload_backend.py +++ b/tests/test_cpu_offload_backend.py @@ -1,5 +1,6 @@ import asyncio import unittest +import uuid from RL_Framework.engine.cmlfq_backend import CMLFQGenerationBackend @@ -9,6 +10,20 @@ def __init__(self): self.calls = [] self.released = [] self.ready = asyncio.Event() + self.instance_id = uuid.uuid4().hex + self.generation = 0 + self.revision = 0 + self.leases = set() + + async def acquire_kv_lease(self): + lease_id = uuid.uuid4().hex + self.leases.add(lease_id) + return dict(lease_id=lease_id, snapshot_id=f"snapshot-{self.revision}", + optimizer_revision=self.revision, server_instance_id=self.instance_id, + server_generation=self.generation) + + async def release_kv_lease(self, lease): + self.leases.discard(lease["lease_id"]) async def generate(self, prompt, kv_transfer_params, **kwargs): self.calls.append(kv_transfer_params) diff --git a/tests/test_nccl_store_lifecycle.py b/tests/test_nccl_store_lifecycle.py new file mode 100644 index 0000000..3b70cc9 --- /dev/null +++ b/tests/test_nccl_store_lifecycle.py @@ -0,0 +1,74 @@ +"""The stateless rendezvous transfers its listening FD to the real TCPStore.""" +import os +from pathlib import Path +import subprocess +import sys +import textwrap +from types import ModuleType, SimpleNamespace + +import pytest + +from RL_Framework.infra.sync import nccl_weight_sync +from RL_Framework.infra.sync.nccl_weight_sync import NcclReloadSpec + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_only_master_transfers_listening_socket_ownership(monkeypatch, rank): + detached = [] + group = SimpleNamespace(socket=SimpleNamespace(detach=lambda: detached.append(True))) + utils = ModuleType("vllm.distributed.utils") + utils.StatelessProcessGroup = SimpleNamespace(create=lambda **_kwargs: group) + monkeypatch.setitem(sys.modules, "vllm.distributed.utils", utils) + + assert nccl_weight_sync._group(NcclReloadSpec("127.0.0.1", 12345, 2, rank, "cpu")) is group + assert detached == ([True] if rank == 0 else []) + + +def test_real_tcp_store_survives_python_socket_cleanup_and_cache_eviction(): + # A subprocess is essential: closing a socket that still owns TCPStore's FD + # triggers SIGABRT in a C++ daemon, bypassing Python exception handling. Use + # real CPU PyTorch but reproduce vLLM's small create() without importing GPU + # dependencies. Multiple groups also exercise independent cache lifetimes. + script = textwrap.dedent(""" + import gc + import socket + import sys + from datetime import timedelta + from types import ModuleType, SimpleNamespace + import torch.distributed as dist + from RL_Framework.infra.sync import nccl_weight_sync + + def create(host, port, rank, world_size, store_timeout): + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind((host, port)) + sock.listen() + port = sock.getsockname()[1] + store = dist.TCPStore(host, port, world_size, True, + timeout=timedelta(seconds=store_timeout), use_libuv=False, + master_listen_fd=sock.fileno()) + return SimpleNamespace(socket=sock, store=store) + + utils = ModuleType('vllm.distributed.utils') + utils.StatelessProcessGroup = SimpleNamespace(create=create) + sys.modules['vllm.distributed.utils'] = utils + for index in range(3): + spec = nccl_weight_sync.NcclReloadSpec('127.0.0.1', 0, 1, 0, 'cpu', timeout_s=2) + group = nccl_weight_sync._group(spec) + assert group.socket.fileno() == -1 + group.socket.close() + group.store.set('working', b'yes') + assert group.store.get('working') == b'yes' + nccl_weight_sync._COMMUNICATOR_CACHE[('test', index)] = SimpleNamespace(group=group) + del group + nccl_weight_sync.clear_communicator_cache() + gc.collect() + print('ORDERLY_TCPSTORE_EXIT', flush=True) + """) + root = Path(nccl_weight_sync.__file__).resolve().parents[2] + # Preserve the package alias provided by the ordinary project test runner. + environment = dict(os.environ) + completed = subprocess.run([sys.executable, "-B", "-c", script], cwd=root, + env=environment, capture_output=True, text=True, timeout=30) + assert completed.returncode == 0, completed.stdout + completed.stderr + assert "ORDERLY_TCPSTORE_EXIT" in completed.stdout + assert "TCPStoreMasterDaemon::run() failed" not in completed.stderr diff --git a/tests/test_nccl_store_timeout.py b/tests/test_nccl_store_timeout.py new file mode 100644 index 0000000..3e7cc32 --- /dev/null +++ b/tests/test_nccl_store_timeout.py @@ -0,0 +1,54 @@ +"""Shared rendezvous waits preserve application timeout=0 and subsecond limits.""" +from datetime import timedelta +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import torch.distributed as dist + +from RL_Framework.infra.sync import nccl_weight_sync +from RL_Framework.infra.sync.nccl_weight_sync import NcclReloadSpec + + +@pytest.fixture +def create_group(monkeypatch): + calls = [] + utils = ModuleType('vllm.distributed.utils') + utils.StatelessProcessGroup = SimpleNamespace(create=lambda **kwargs: calls.append(kwargs) or kwargs) + for name in ('vllm', 'vllm.distributed'): + package = ModuleType(name) + package.__path__ = [] + monkeypatch.setitem(sys.modules, name, package) + monkeypatch.setitem(sys.modules, 'vllm.distributed.utils', utils) + return calls + + +@pytest.mark.parametrize('timeout', [0.125, 1, 29.5, 1200]) +def test_positive_nccl_store_timeout_preserves_seconds_without_minimum_or_truncation(create_group, timeout): + nccl_weight_sync._group(NcclReloadSpec('127.0.0.1', 12345, 2, 0, 'cpu', timeout_s=timeout)) + assert create_group[0]['store_timeout'] == timeout + + +def test_zero_nccl_store_timeout_uses_safe_unbounded_sentinel(create_group): + nccl_weight_sync._group(NcclReloadSpec('127.0.0.1', 12345, 2, 0, 'cpu', timeout_s=0)) + timeout = create_group[0]['store_timeout'] + assert timeout > 365 * 24 * 60 * 60 + assert timeout * 1_000_000_000 < 2**63 + # Exercise the same timedelta conversion performed by vLLM. + assert timedelta(seconds=timeout).total_seconds() == timeout + + +@pytest.mark.parametrize('timeout', [-1, float('nan'), float('inf')]) +def test_invalid_nccl_store_timeout_fails_before_group_creation(create_group, timeout): + with pytest.raises(ValueError, match='non-negative seconds'): + nccl_weight_sync._group(NcclReloadSpec('127.0.0.1', 12345, 2, 0, 'cpu', timeout_s=timeout)) + assert not create_group + + +def test_unbounded_sentinel_is_accepted_by_real_torch_tcp_store(): + # No peer or wait is needed: this verifies PyTorch's C++ duration conversion + # accepts the sentinel and the resulting store remains functional. + timeout = timedelta(seconds=nccl_weight_sync._UNBOUNDED_STORE_TIMEOUT_S) + store = dist.TCPStore('127.0.0.1', 0, 1, True, timeout=timeout) + store.set('probe', b'ready') + assert store.get('probe') == b'ready' diff --git a/tests/test_serving_protocol.py b/tests/test_serving_protocol.py new file mode 100644 index 0000000..ddf26ef --- /dev/null +++ b/tests/test_serving_protocol.py @@ -0,0 +1,306 @@ +"""Shared serving correctness across arbitrary TP layouts, reload and KV leases.""" + +import asyncio +import importlib.util +import json +from pathlib import Path +from types import SimpleNamespace as NS +import sys + +import pytest +import torch + +from RL_Framework.scripts.serving_protocol import ServingState, GenerationMiddleware, fence_kv_payload +from RL_Framework.vllm_hot_reload import InplaceReloadWorkerExtension + + +def workers(tp, generation=1, tensors=7, offset=0): + return [dict(rank=rank, nccl_rank=1+offset+rank, service_generation=generation, + received_tensors=tensors, loaded_parameters=tensors, receive_complete=True) for rank in range(tp)] + + +@pytest.mark.parametrize('tp', [1, 2, 4]) +def test_initial_and_all_tp_publication_identity(tp): + async def run(): + state = ServingState(tp, '/models/test') + assert state.status()['server_generation_verified'] + assert state.optimizer_revision == state.server_generation == 0 + await state.begin_reload(0) + state.install(snapshot_id='revision100', optimizer_revision=100, workers=workers(tp, offset=4), rank_offset=4) + await state.reopen() + status = state.status() + assert status['server_generation'] == 1 and status['optimizer_revision'] == 100 + assert status['workers'] == workers(tp, offset=4) + assert status['server_generation_verified'] + asyncio.run(run()) + + +@pytest.mark.parametrize('mutate', [ + lambda rows: rows.pop(), + lambda rows: rows[1].update(rank=0), + lambda rows: rows[1].update(service_generation=0), + lambda rows: rows[1].update(receive_complete=False), + lambda rows: rows[1].update(received_tensors=6), + lambda rows: rows[1].update(nccl_rank=20), +]) +def test_incomplete_stale_or_inconsistent_tp_receipts_cannot_advance(mutate): + async def run(): + state = ServingState(2) + await state.begin_reload(0) + rows = workers(2) + mutate(rows) + with pytest.raises(RuntimeError): + state.install(snapshot_id='new', optimizer_revision=1, workers=rows) + await state.fail('invalid receipt') + assert not state.status()['server_generation_verified'] + assert state.server_generation == 0 + with pytest.raises(RuntimeError): + await state.reopen() + with pytest.raises(RuntimeError): + await state.enter() + asyncio.run(run()) + + +def test_reload_drains_inflight_requests_and_kv_continuations_before_install(): + async def run(): + state = ServingState(2) + lease = await state.acquire_kv_lease() + assert await state.enter() == 0 + installing = asyncio.create_task(state.begin_reload(0)) + await asyncio.sleep(0) + assert state.closed and not installing.done() + with pytest.raises(RuntimeError): + await state.acquire_kv_lease() + waiting = asyncio.create_task(state.enter()) + await asyncio.sleep(0) + assert not waiting.done() + assert await state.enter(lease['lease_id']) == 0 # Existing KV continuation progresses. + await state.leave() + await state.leave() + assert not installing.done() + await state.release_kv_lease(lease['lease_id']) + await installing + state.install(snapshot_id='new', optimizer_revision=1, workers=workers(2)) + await state.reopen() + assert await waiting == 1 + await state.leave() + assert state.active_requests == 0 + asyncio.run(run()) + + +def test_positive_drain_timeout_reopens_before_any_mutation(): + async def run(): + state = ServingState(1) + await state.enter() + with pytest.raises(TimeoutError): + await state.begin_reload(.001) + assert not state.closed and state.failed is None and state.server_generation == 0 + await state.leave() + asyncio.run(run()) + + +def test_server_owns_target_kv_identity_and_strips_stale_import(): + state = ServingState(2, '/models/test') + payload = dict(kv_transfer_params=dict(libra_offload_key='save', libra_load_key='old', + libra_kv_identity=dict(snapshot_id='old', optimizer_revision=-1), libra_kv_source_identity={'old': True})) + fenced = fence_kv_payload(payload, state)['kv_transfer_params'] + assert 'libra_load_key' not in fenced and 'libra_kv_source_identity' not in fenced + assert fenced['libra_kv_identity']['server_instance_id'] == state.server_instance_id + assert fenced['libra_offload_key'] == 'save' + + +def test_generation_headers_match_actual_completion_and_body_identity_is_authoritative(): + async def run(): + state, sent = ServingState(1), [] + received_payload = [] + async def app(scope, receive, send): + received_payload.append(json.loads((await receive())['body'])) + await send(dict(type='http.response.start', status=200, headers=[(b'content-type', b'application/json')])) + await send(dict(type='http.response.body', body=b'{"choices":[]}')) + middleware = GenerationMiddleware(app, state) + payload = dict(prompt=[1], kv_transfer_params=dict(libra_offload_key='save', libra_load_key='old')) + async def receive(): + return dict(type='http.request', body=json.dumps(payload).encode(), more_body=False) + async def send(message): + sent.append(message) + scope = dict(type='http', method='POST', path='/v1/completions', + headers=[(b'x-libra-request-id', b'actual-request')]) + await middleware(scope, receive, send) + headers = dict(sent[0]['headers']) + assert headers[b'x-libra-request-id'] == b'actual-request' + assert headers[b'x-libra-service-generation'] == b'0' + assert headers[b'x-libra-service-generation-verified'] == b'1' + assert headers[b'x-libra-reload-overlap'] == b'0' + assert state.active_requests == 0 + params = received_payload[0]['kv_transfer_params'] + assert 'libra_load_key' not in params + assert params['libra_kv_identity']['server_instance_id'] == state.server_instance_id + asyncio.run(run()) + + +@pytest.mark.parametrize('consume_all', [True, False]) +def test_worker_attests_actual_source_tensor_count_and_done_marker(monkeypatch, consume_all): + from RL_Framework.infra.sync import nccl_weight_sync + monkeypatch.setattr(torch.cuda, 'current_device', lambda: 0) + monkeypatch.setattr(torch.cuda, 'synchronize', lambda: None) + def source(spec): + for i in range(3): + yield f'weight{i}', torch.ones(1) + monkeypatch.setattr(nccl_weight_sync, 'receive_weights', source) + worker = InplaceReloadWorkerExtension() + worker.rank = 1 + def load(weights): + if consume_all: + return {name for name, tensor in weights} + return {next(weights)[0]} + worker.model_runner = NS(model=NS(load_weights=load)) + if not consume_all: + with pytest.raises(RuntimeError, match='done marker'): + worker.reload_weights_nccl('host', 1, 3, 0) + assert getattr(worker, '_libra_service_generation', 0) == 0 + else: + result = worker.reload_weights_nccl('host', 1, 3, 0, expected_tensor_count=3) + assert result['receive_complete'] and result['received_tensors'] == 3 + assert result['service_generation'] == 1 and result['nccl_rank'] == 2 + + +def test_supervisor_echoes_snapshot_and_revision_into_physical_ack(tmp_path, monkeypatch): + path = Path(__file__).resolve().parents[1] / 'scripts/restartable_vllm_server.py' + spec = importlib.util.spec_from_file_location('serving_supervisor_test', path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + request = dict(mode='nccl', version=100, snapshot_id='final', optimizer_revision=100, + topology_epoch=8, nccl_host='train', nccl_port=29620, nccl_world_size=5, + nccl_rank_offsets={'current': 0}, nccl_chunk_bytes=1024) + (tmp_path / 'reload_request.json').write_text(json.dumps(request)) + args = NS(instance_id='current', control_dir=str(tmp_path), health_url='http://service/health', + initial_model='/models/test', ready_timeout=0, reload_strategy='parallel', reload_method='inplace', + poll_interval=.01, command=['server']) + monkeypatch.setattr(module, 'parse_args', lambda: args) + child = NS(poll=lambda: None) + monkeypatch.setattr(module.subprocess, 'Popen', lambda *args, **kwargs: child) + monkeypatch.setattr(module, 'wait_until_ready', lambda *args: None) + monkeypatch.setattr(module, 'stop_process', lambda *args: None) + monkeypatch.setattr(module.signal, 'signal', lambda *args: None) + forwarded = [] + def reload(url, payload, timeout): + forwarded.append(payload) + return dict(workers=workers(4), service_instance_id='actual', service_generation=1) + monkeypatch.setattr(module, 'reload_nccl', reload) + calls = 0 + def sleep(seconds): + nonlocal calls + calls += 1 + if calls == 1: + request['snapshot_id'] = 'same-params-new-topology' + request['topology_epoch'] = 9 + (tmp_path / 'reload_request.json').write_text(json.dumps(request)) + else: + raise KeyboardInterrupt() + monkeypatch.setattr(module.time, 'sleep', sleep) + assert module.main() == 130 + assert [payload['snapshot_id'] for payload in forwarded] == ['final', 'same-params-new-topology'] + ack = json.loads((tmp_path / 'ack_current_100.json').read_text()) + assert ack['snapshot_id'] == 'same-params-new-topology' and ack['optimizer_revision'] == 100 + assert ack['topology_epoch'] == 9 and ack['server_response']['service_instance_id'] == 'actual' + + +def load_serving_app(monkeypatch, tp=4, invalid_ack=False): + """Import the actual FastAPI routes with only the unavailable vLLM runtime stubbed.""" + import argparse + import types + from fastapi import FastAPI, Request + from fastapi.responses import JSONResponse + modules = {} + for name in ('vllm', 'vllm.entrypoints', 'vllm.entrypoints.openai', + 'vllm.entrypoints.openai.api_server', 'vllm.entrypoints.openai.cli_args', + 'vllm.utils', 'vllm.v1', 'vllm.v1.engine', 'vllm.v1.engine.core'): + modules[name] = types.ModuleType(name) + monkeypatch.setitem(sys.modules, name, modules[name]) + modules['vllm'].__version__ = '0.9.2' + modules['vllm.entrypoints.openai'].api_server = modules['vllm.entrypoints.openai.api_server'] + modules['vllm.entrypoints.openai.cli_args'].make_arg_parser = lambda parser: parser + modules['vllm.entrypoints.openai.cli_args'].validate_parsed_serve_args = lambda args: None + modules['vllm.utils'].FlexibleArgumentParser = argparse.ArgumentParser + modules['vllm.v1.engine.core'].EngineCore = type('Core', (), {}) + monkeypatch.setitem(sys.modules, 'uvloop', NS(run=lambda coroutine: None)) + def build(args): + app = FastAPI() + @app.post('/v1/completions') + async def completion(request: Request): + return JSONResponse(dict(choices=[dict(text='ok')], received=await request.json())) + return app + async def utility(name): + assert name == 'libra_reset_prefix_cache' + return True + async def rpc(name, timeout, args=()): + assert timeout is None or timeout == 120 + if name == 'reload_weights_nccl': + rows = workers(tp, offset=args[3]) + return rows[:-1] if invalid_ack else rows + raise AssertionError(name) + client = NS(collective_rpc=rpc, engine_core=NS(call_utility_async=utility)) + api = modules['vllm.entrypoints.openai.api_server'] + api.build_app = build + api.engine_client = lambda request: request.app.state.client + path = Path(__file__).resolve().parents[1] / 'scripts/vllm_hot_reload_api_server.py' + spec = importlib.util.spec_from_file_location('libra_serving_routes_test', path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + app = module.build_app_with_reload_routes(NS(tensor_parallel_size=tp, model='/models/test')) + app.state.client = client + return app + + +def test_real_http_routes_fence_reload_by_kv_lease_and_attest_physical_workers(monkeypatch): + import httpx + app = load_serving_app(monkeypatch) + async def run(): + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url='http://service') as client: + initial = (await client.get('/publication/status')).json() + assert initial['server_generation_verified'] and initial['tp_size'] == 4 + assert (await client.get('/libra/kv/identity')).json()['snapshot_id'] == initial['snapshot_id'] + lease = (await client.post('/libra/kv/lease', json={})).json() + publishing = asyncio.create_task(client.post('/reload_weights_nccl', json=dict( + host='training', port=29620, world_size=5, rank_offset=0, + snapshot_id='final100', optimizer_revision=100, timeout_seconds=0))) + await asyncio.sleep(0) + assert not publishing.done() + # Holding an existing lease allows its continuation, even as the + # publisher has closed new admission and waits for that same lease. + continuation = await client.post('/v1/completions', headers={ + 'X-Libra-KV-Lease': lease['lease_id'], 'X-Libra-Request-Id': 'old-request'}, + json=dict(prompt=[1])) + assert continuation.headers['x-libra-service-generation'] == '0' + await client.delete('/libra/kv/lease/' + lease['lease_id']) + reply = await publishing + assert reply.status_code == 200 + assert reply.json()['prefix_cache_cleared'] + assert len(reply.json()['workers']) == 4 + live = (await client.get('/publication/status')).json() + assert live['optimizer_revision'] == 100 and live['snapshot_id'] == 'final100' + assert live['server_generation'] == 1 and live['server_generation_verified'] + assert live['workers'] == reply.json()['workers'] + response = await client.post('/v1/completions', headers={'X-Libra-Request-Id': 'new-request'}, + json=dict(prompt=[1], kv_transfer_params=dict(libra_offload_key='next', libra_load_key='old', + libra_kv_identity=lease))) + assert response.headers['x-libra-service-generation-verified'] == '1' + assert response.headers['x-libra-service-generation'] == '1' + assert 'libra_load_key' not in response.json()['received']['kv_transfer_params'] + asyncio.run(run()) + + +def test_real_http_partial_tp_failure_quarantines_service(monkeypatch): + import httpx + app = load_serving_app(monkeypatch, invalid_ack=True) + async def run(): + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app, raise_app_exceptions=False), base_url='http://service') as client: + response = await client.post('/reload_weights_nccl', json=dict(host='training', port=29620, + world_size=5, rank_offset=0, snapshot_id='broken', optimizer_revision=1, timeout_seconds=0)) + assert response.status_code == 500 + state = (await client.get('/publication/status')).json() + assert state['failed'] and not state['server_generation_verified'] + assert state['server_generation'] == 0 + assert (await client.post('/v1/completions', json=dict(prompt=[1]))).status_code == 409 + asyncio.run(run()) diff --git a/tests/test_versioned_kv_compatibility.py b/tests/test_versioned_kv_compatibility.py new file mode 100644 index 0000000..443c1ed --- /dev/null +++ b/tests/test_versioned_kv_compatibility.py @@ -0,0 +1,223 @@ +"""Version/incarnation fencing through real KV clients and the native reload gate.""" +import asyncio +import importlib.util +import json +from pathlib import Path +import sys +from types import ModuleType, SimpleNamespace + +import httpx +from fastapi import FastAPI, HTTPException, Request +import pytest +import torch + +from RL_Framework.engine.cmlfq_backend import CMLFQGenerationBackend +from RL_Framework.engine.rollout_engine import VLLMRolloutEngine +from RL_Framework.infra.scheduling.kv_identity import same_snapshot, same_incarnation +from RL_Framework.scripts.serving_protocol import ServingState, GenerationMiddleware +from test_cpu_offload_backend import FakeCPUBackend + + +def peer(revision=7, before_completion=None): + state = ServingState(1, "same-model") + state.snapshot_id, state.optimizer_revision = f"snapshot-{revision}", revision + app, calls = FastAPI(), [] + + @app.post("/libra/kv/lease") + async def acquire(request: Request): + try: + return await state.acquire_kv_lease((await request.json()).get("lease_id")) + except RuntimeError as error: + raise HTTPException(409, str(error)) + + @app.delete("/libra/kv/lease/{lease_id}") + async def release(lease_id: str): + await state.release_kv_lease(lease_id) + return {"released": True} + + @app.get("/libra/kv/identity") + async def identity(): + return state.status() + + @app.post("/v1/completions") + async def complete(request: Request): + calls.append(await request.json()) + if before_completion: + await before_completion() + return dict(choices=[dict(text="answer", logprobs={})], kv_transfer_params=dict( + remote_engine_id="source", remote_request_id="request", remote_host="host", + remote_port=5555, remote_block_ids=[[1, 2]])) + + app.add_middleware(GenerationMiddleware, state=state) + engine = VLLMRolloutEngine(host="test") + engine.http_client = httpx.AsyncClient(transport=httpx.ASGITransport(app=app)) + engine._http_client_loop = asyncio.get_running_loop() + return engine, state, calls + + +def test_identity_distinguishes_shared_weights_and_service_incarnations(): + identity = dict(snapshot_id="shared", optimizer_revision=7, server_instance_id="a", server_generation=3) + other = dict(identity, server_instance_id="b") + assert same_snapshot(identity, other) + assert not same_incarnation(identity, other) + assert not same_incarnation(identity, dict(identity, server_generation=4)) + assert not same_snapshot(identity, dict(identity, optimizer_revision=8)) + assert not same_snapshot({}, {}) + + +@pytest.mark.parametrize("target_revision", [7, 8]) +def test_nixl_same_snapshot_transfers_but_new_revision_recomputes(target_revision): + async def run(): + source, source_state, source_calls = peer() + target, target_state, target_calls = peer(target_revision) + backend = CMLFQGenerationBackend("nixl", [(1, 4)]) + output = await backend.generate(source, target, "full prompt", "transfer", prompt_token_ids=[1, 2, 3]) + assert output.path == ("transfer" if target_revision == 7 else "recompute") + assert bool(source_calls) == (target_revision == 7) + assert target_calls[-1]["prompt"] == [1, 2, 3] + assert bool(target_calls[-1].get("kv_transfer_params")) == (target_revision == 7) + assert not source_state.kv_leases and not target_state.kv_leases + await source.close() + await target.close() + asyncio.run(run()) + + +@pytest.mark.parametrize("change,reason", [("revision", "snapshot_mismatch"), ("generation", "incarnation_mismatch"), + ("restart", "incarnation_mismatch")]) +def test_cpu_old_cache_recomputes_after_reload_or_restart(change, reason): + async def run(): + source, target = FakeCPUBackend(), FakeCPUBackend() + source.ready.set() + target.ready.set() + backend = CMLFQGenerationBackend("cpu_offload", [(1, 4)]) + await backend.generate(source, source, "prompt", "stay", "request") + if change == "revision": + target.revision += 1 + elif change == "generation": + source.generation += 1 + else: + source.instance_id = "restarted" + output = await backend.generate(source, target, "prompt+tool", "transfer", "request") + assert output.path == "recompute" and output.fallback_reason == reason + assert "libra_load_key" not in target.calls[-1] + assert not source.leases and not target.leases + await backend.close() + asyncio.run(run()) + + +def test_cancelled_cpu_wait_releases_both_source_and_target_leases(): + async def run(): + source, target = FakeCPUBackend(), FakeCPUBackend() + backend = CMLFQGenerationBackend("cpu_offload", [(1, 4)]) + await backend.generate(source, source, "prompt", "stay", "request") + task = asyncio.create_task(backend.generate(source, target, "prompt+tool", "transfer", "request")) + await asyncio.sleep(0) + assert source.leases and target.leases + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert not source.leases and not target.leases + await backend.close() + asyncio.run(run()) + + +def load_connector(monkeypatch): + name = "vllm.distributed.kv_transfer.kv_connector.v1.base" + base = ModuleType(name) + base.KVConnectorBase_V1, base.KVConnectorMetadata = object, object + base.KVConnectorRole = SimpleNamespace(WORKER="worker") + distributed = ModuleType("vllm.distributed") + distributed.get_tensor_model_parallel_rank = lambda: 0 + monkeypatch.setitem(sys.modules, name, base) + monkeypatch.setitem(sys.modules, "vllm.distributed", distributed) + path = Path(__file__).parents[1] / "infra/scheduling/vllm_cpu_offload.py" + module_name = "RL_Framework.infra.scheduling.kv_compat_connector_fixture" + spec = importlib.util.spec_from_file_location(module_name, path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, module_name, module) + spec.loader.exec_module(module) + return module + + +def connector_fixture(monkeypatch, tmp_path, source_tp=2, target_tp=4, rank=0): + module = load_connector(monkeypatch) + connector = module.LibraCPUOffloadConnector.__new__(module.LibraCPUOffloadConnector) + connector.root, connector.tp, connector.rank = str(tmp_path), target_tp, rank + connector.block_size, connector.num_heads = 2, 8 + connector.matches, connector.imported = {}, {} + key = "a" * 32 + directory = tmp_path / key + directory.mkdir(exist_ok=True) + source_identity = dict(snapshot_id="snapshot", optimizer_revision=7, server_instance_id="source", server_generation=3) + target_identity = dict(source_identity, server_instance_id="target", server_generation=8) + manifest = dict(tokens=list(range(4)), source_tp=source_tp, block_size=2, libra_kv_identity=source_identity) + for source_rank in range(source_tp): + (directory / f"rank_{source_rank}.json").write_text(json.dumps(manifest)) + request = SimpleNamespace(kv_transfer_params=dict(libra_load_key=key, libra_kv_identity=target_identity, + libra_kv_source_identity=source_identity), request_id="request", prompt_token_ids=list(range(5))) + return module, connector, directory, manifest, request + + +def test_cpu_connector_checks_all_tp_source_incarnations(monkeypatch, tmp_path): + _, connector, directory, manifest, request = connector_fixture(monkeypatch, tmp_path) + assert connector.get_num_new_matched_tokens(request, 0) == (4, False) + changed = dict(manifest, libra_kv_identity=dict(manifest["libra_kv_identity"], server_generation=4)) + (directory / "rank_1.json").write_text(json.dumps(changed)) + assert connector.get_num_new_matched_tokens(request, 0) == (0, False) + + +@pytest.mark.parametrize("source_tp,target_tp", [(1, 2), (2, 1), (2, 4), (4, 2)]) +def test_cpu_real_tensor_head_resharding_keeps_identity_checks(monkeypatch, tmp_path, source_tp, target_tp): + module, connector, directory, _, request = connector_fixture(monkeypatch, tmp_path, source_tp, target_tp) + full = torch.arange(2 * 2 * 2 * 8 * 4, dtype=torch.float32).reshape(2, 2, 2, 8, 4) + for rank, shard in enumerate(full.chunk(source_tp, dim=-2)): + torch.save({"layer": shard.contiguous()}, directory / f"rank_{rank}.pt") + # Exercise the production file loading/head slicing/index-copy code on CPU; + # CUDA DMA and pinned transfers remain covered by the GPU regression. + monkeypatch.setattr(torch.Tensor, "pin_memory", lambda value: value) + module.torch = SimpleNamespace(tensor=lambda values, **kwargs: torch.tensor(values, dtype=kwargs.get("dtype")), + long=torch.long, load=torch.load, cat=torch.cat) + for rank in range(target_tp): + connector.rank = rank + destination = torch.zeros(2, 3, 2, 8 // target_tp, 4) + connector.caches = {"layer": destination} + connector._reload(dict(key=directory.name, source_tp=source_tp, num_tokens=4, block_ids=[2, 0], + **request.kv_transfer_params)) + torch.testing.assert_close(destination[:, [2, 0]], full.chunk(target_tp, dim=-2)[rank]) + + +@pytest.mark.parametrize("unknown", [None, [], "snapshot", {}, {"snapshot_id": "snapshot", "optimizer_revision": True}]) +def test_unknown_kv_identities_are_rejected_without_crashing(unknown): + identity = dict(snapshot_id="snapshot", optimizer_revision=1, server_instance_id="s", server_generation=0) + assert not same_snapshot(unknown, identity) + assert not same_snapshot(identity, unknown) + assert not same_incarnation(unknown, identity) + + +def test_worker_fences_changed_incarnation_before_touching_cuda(monkeypatch, tmp_path): + _, connector, directory, manifest, request = connector_fixture(monkeypatch, tmp_path) + assert connector.get_num_new_matched_tokens(request, 0) == (4, False) + changed = dict(manifest, libra_kv_identity=dict(manifest["libra_kv_identity"], server_instance_id="restarted")) + (directory / "rank_1.json").write_text(json.dumps(changed)) + with pytest.raises(RuntimeError, match="identity changed"): + connector._reload(dict(key=directory.name, source_tp=2, num_tokens=4, block_ids=[0, 1], + **request.kv_transfer_params)) + + +def test_cancelled_lease_response_withdraws_server_reservation(): + async def run(): + state = ServingState(1) + async def handler(request): + if request.method == "POST": + await state.acquire_kv_lease(json.loads(request.content)["lease_id"]) + raise asyncio.CancelledError() + await state.release_kv_lease(request.url.path.rsplit("/", 1)[-1]) + return httpx.Response(200, json={"released": True}) + engine = VLLMRolloutEngine(host="test") + engine.http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + engine._http_client_loop = asyncio.get_running_loop() + with pytest.raises(asyncio.CancelledError): + await engine.acquire_kv_lease() + assert not state.kv_leases + await engine.close() + asyncio.run(run()) diff --git a/tests/test_vllm_cpu_offload.py b/tests/test_vllm_cpu_offload.py index 8e08fb8..9557caa 100644 --- a/tests/test_vllm_cpu_offload.py +++ b/tests/test_vllm_cpu_offload.py @@ -13,6 +13,10 @@ from RL_Framework.infra.scheduling.vllm_cpu_offload import LibraCPUOffloadConnector, cache_directory +IDENTITY = dict(snapshot_id="test-snapshot", optimizer_revision=1, + server_instance_id="source", server_generation=0) + + def bare_connector(root, tp=1, rank=0): connector = object.__new__(LibraCPUOffloadConnector) connector.root, connector.tp, connector.rank = root, tp, rank @@ -28,9 +32,12 @@ def test_prefix_identity_and_complete_block_boundaries(): directory.mkdir() (directory / "rank_0.json").write_text(json.dumps({ "tokens": list(range(64)), "source_tp": 1, "block_size": 16, + "libra_kv_identity": IDENTITY, })) connector = bare_connector(root) - request = SimpleNamespace(kv_transfer_params={"libra_load_key": key}, + request = SimpleNamespace(kv_transfer_params={"libra_load_key": key, + "libra_kv_identity": IDENTITY, + "libra_kv_source_identity": IDENTITY}, request_id="r", prompt_token_ids=list(range(65))) assert connector.get_num_new_matched_tokens(request, 16) == (48, False) request.prompt_token_ids[1] = 999 @@ -58,12 +65,18 @@ def test_real_cuda_cpu_head_resharding(source_tp, target_tp): pinned.copy_(gpu, non_blocking=True) stream.synchronize() torch.save({"layer": pinned}, directory / f"rank_{rank}.pt") + (directory / f"rank_{rank}.json").write_text(json.dumps({ + "tokens": list(range(32)), "source_tp": source_tp, "block_size": 16, + "libra_kv_identity": IDENTITY, + })) for rank in range(target_tp): with torch.cuda.device(rank): connector = bare_connector(root, target_tp, rank) destination = torch.zeros(2, 3, 16, 8 // target_tp, 4, device="cuda") connector.caches = {"layer": destination} connector._reload({"key": key, "source_tp": source_tp, - "num_tokens": 32, "block_ids": [2, 0]}) + "num_tokens": 32, "block_ids": [2, 0], + "libra_kv_identity": IDENTITY, + "libra_kv_source_identity": IDENTITY}) observed = destination[:, [2, 0]].cpu() assert torch.equal(observed, full.chunk(target_tp, dim=-2)[rank]) diff --git a/vllm_hot_reload.py b/vllm_hot_reload.py index 5cc55f2..dc1dd87 100644 --- a/vllm_hot_reload.py +++ b/vllm_hot_reload.py @@ -48,7 +48,7 @@ def reload_weights(self, checkpoint_path: str) -> dict[str, Any]: self.vllm_config.model_config.model = previous_path raise finished_at = time.time() - return { + result = { "rank": int(getattr(self, "rank", 0)), "checkpoint_path": str(checkpoint), "previous_path": previous_path, @@ -56,6 +56,7 @@ def reload_weights(self, checkpoint_path: str) -> dict[str, Any]: "finished_at": finished_at, "load_seconds": time.perf_counter() - started, } + return self._attest_reload(result) def reload_weights_nccl( self, @@ -65,6 +66,7 @@ def reload_weights_nccl( rank_offset: int, timeout_seconds: float = 1200.0, chunk_bytes: int = 256 * 1024 * 1024, + expected_tensor_count: int | None = None, ) -> dict[str, Any]: """Load full HF-format weights from a Megatron NCCL broadcast.""" import torch @@ -85,11 +87,35 @@ def reload_weights_nccl( chunk_bytes=int(chunk_bytes), ) started = time.perf_counter() - loaded = self.model_runner.model.load_weights(receive_weights(spec)) + count, complete = 0, False + source = receive_weights(spec) + def observed(): + nonlocal count, complete + for name, tensor in source: + count += 1 + yield name, tensor + complete = True + weights = observed() + try: + loaded = self.model_runner.model.load_weights(weights) + if not complete: + raise RuntimeError('Model loader returned before the NCCL done marker') + if expected_tensor_count is not None and count != expected_tensor_count: + raise RuntimeError('Received source tensor count differs from the snapshot') + finally: + weights.close() + source.close() torch.cuda.synchronize() - return { + result = { "rank": local_rank, "nccl_rank": spec.rank, "loaded_parameters": len(loaded) if loaded is not None else 0, + "received_tensors": count, + "receive_complete": complete, "load_seconds": time.perf_counter() - started, } + return self._attest_reload(result) + + def _attest_reload(self, result): + self._libra_service_generation = getattr(self, '_libra_service_generation', 0) + 1 + return dict(result, service_generation=self._libra_service_generation) From f6f28c9ae0b8843a7313d1bbe7c53312672666ac Mon Sep 17 00:00:00 2001 From: cytsoi <154662279+tcyjohn@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:19:03 +0800 Subject: [PATCH 2/2] Limit serving fix to publication and request lifecycle --- engine/cmlfq_backend.py | 98 ++-------- infra/scheduling/kv_identity.py | 28 --- infra/scheduling/vllm_cpu_offload.py | 33 +--- scripts/serving_protocol.py | 48 +---- scripts/test_cpu_ci.sh | 1 - tests/test_cmlfq_cost_scheduler.py | 11 +- tests/test_cpu_offload_backend.py | 15 -- tests/test_serving_protocol.py | 54 +++--- tests/test_versioned_kv_compatibility.py | 223 ----------------------- tests/test_vllm_cpu_offload.py | 17 +- 10 files changed, 63 insertions(+), 465 deletions(-) delete mode 100644 infra/scheduling/kv_identity.py delete mode 100644 tests/test_versioned_kv_compatibility.py diff --git a/engine/cmlfq_backend.py b/engine/cmlfq_backend.py index ee6b470..37f782d 100644 --- a/engine/cmlfq_backend.py +++ b/engine/cmlfq_backend.py @@ -14,8 +14,6 @@ from dataclasses import dataclass from typing import Any -from RL_Framework.infra.scheduling.kv_identity import same_snapshot, same_incarnation, snapshot_identity - logger = logging.getLogger(__name__) @@ -52,28 +50,12 @@ async def generate(self, source, target, prompt: str, path: str, request_id: str return BackendResult(await target.generate(prompt=prompt, **kwargs), path) if self.mode != "nixl": raise RuntimeError("transfer selected on a recompute-only backend") - leases = [] try: - # Acquisitions fail immediately when a publisher already owns a peer. - # Never wait for another engine while holding a generation lease. - for engine in dict.fromkeys((source, target)): - leases.append((engine, await engine.acquire_kv_lease())) - by_engine = {engine: lease for engine, lease in leases} - if not same_snapshot(by_engine[source], by_engine[target]): - raise RuntimeError("KV snapshot differs or is unknown") - # Strict workflows must materialize exactly the token IDs decoded on - # the target; tokenizing transcript text again can change its prefix. - preparation = {} - if kwargs.get("prompt_token_ids") is not None: - preparation["prompt_token_ids"] = kwargs["prompt_token_ids"] + # Materialize the exact next prompt on the source. Its generated + # token is discarded; the target samples the user's full output. + # Profile transfer_ms must include this preparation overhead. prepared = await source.generate( - # Preparation also goes through the strict sampler hook. Its - # discarded token must use a seeded stochastic row, never the - # greedy fast path. Sampling does not change cached prompt KV. - prompt=prompt, max_new_tokens=1, - temperature=max(float(kwargs.get("temperature", 1.0)), 1e-4), - top_p=kwargs.get("top_p", 1.0), n=1, seed=kwargs.get("seed", 0), - kv_lease=by_engine[source], **preparation, + prompt=prompt, max_new_tokens=1, temperature=0.0, top_p=1.0, n=1, kv_transfer_params={"do_remote_decode": True, "do_remote_prefill": False}, ) params = prepared.get("kv_transfer_params") @@ -85,17 +67,14 @@ async def generate(self, source, target, prompt: str, path: str, request_id: str raise RuntimeError("source exported no KV blocks") params = dict(params) params.update(do_remote_prefill=True, do_remote_decode=False) - output = await target.generate(prompt=prompt, kv_transfer_params=params, - kv_lease=by_engine[target], **kwargs) + output = await target.generate(prompt=prompt, kv_transfer_params=params, **kwargs) return BackendResult(output, "transfer") except Exception as error: + # Non-streaming requests: no partial output has reached the caller. + # Cancellation (BaseException) is deliberately not retried. logger.warning("CMLFQ KV transfer failed; recomputing: %s", error) - fallback = type(error).__name__ - finally: - await asyncio.gather(*(self._release_kv_lease(engine, lease) - for engine, lease in reversed(leases))) - output = await target.generate(prompt=prompt, **kwargs) - return BackendResult(output, "recompute", fallback) + output = await target.generate(prompt=prompt, **kwargs) + return BackendResult(output, "recompute", type(error).__name__) async def _generate_cpu(self, target, prompt, path, request_id, kwargs): if not request_id: @@ -104,61 +83,32 @@ async def _generate_cpu(self, target, prompt, path, request_id, kwargs): key = uuid.uuid4().hex params = {"libra_offload_key": key} fallback = "cache_unavailable" if path == "transfer" and previous is None else "" - leases = [] - lease = None - status = None wait_started = time.perf_counter() - try: + status = None + if previous and path in {"stay", "transfer"}: + old_engine, descriptor = previous try: - lease = await target.acquire_kv_lease() - leases.append((target, lease)) + status = await old_engine.wait_for_cpu_offload( + descriptor["libra_offload_key"], descriptor["source_tp"], + ) + params["libra_load_key"] = descriptor["libra_offload_key"] except Exception as error: - # Plain full-prompt generation remains available on old servers. fallback = type(error).__name__ - identity = snapshot_identity(lease or {}) - params["libra_kv_identity"] = identity - if previous and path in {"stay", "transfer"}: - old_engine, descriptor = previous - previous_identity = descriptor.get("libra_kv_identity") - if lease is not None and same_snapshot(previous_identity, identity): - try: - source_lease = lease if old_engine is target else await old_engine.acquire_kv_lease() - if old_engine is not target: - leases.append((old_engine, source_lease)) - if not same_incarnation(previous_identity, source_lease): - fallback = "incarnation_mismatch" - elif not same_snapshot(source_lease, identity): - fallback = "snapshot_mismatch" - else: - status = await old_engine.wait_for_cpu_offload( - descriptor["libra_offload_key"], descriptor["source_tp"], - ) - params["libra_load_key"] = descriptor["libra_offload_key"] - params["libra_kv_source_identity"] = snapshot_identity(source_lease) - except Exception as error: - fallback = type(error).__name__ - else: - fallback = "snapshot_mismatch" - wait_ms = (time.perf_counter() - wait_started) * 1000 - options = dict(kwargs, **({"kv_lease": lease} if lease is not None else {})) + wait_ms = (time.perf_counter() - wait_started) * 1000 + try: try: - output = await target.generate(prompt=prompt, kv_transfer_params=params, **options) + output = await target.generate(prompt=prompt, kv_transfer_params=params, **kwargs) except Exception as error: if "libra_load_key" not in params: raise fallback = type(error).__name__ params.pop("libra_load_key") - params.pop("libra_kv_source_identity", None) - output = await target.generate(prompt=prompt, kv_transfer_params=params, **options) + output = await target.generate(prompt=prompt, kv_transfer_params=params, **kwargs) except BaseException: self._release(target, key) raise - finally: - await asyncio.gather(*(self._release_kv_lease(engine, held) - for engine, held in reversed(leases))) descriptor = output.get("kv_transfer_params", {}) if descriptor.get("libra_offload_key"): - descriptor = dict(descriptor, libra_kv_identity=descriptor.get("libra_kv_identity", identity)) self._offloads[request_id] = (target, descriptor) else: self._offloads.pop(request_id, None) @@ -175,14 +125,6 @@ async def _generate_cpu(self, target, prompt, path, request_id, kwargs): fallback = "cache_prefix_miss" return BackendResult(output, actual, fallback) - async def _release_kv_lease(self, engine, lease): - try: - await engine.release_kv_lease(lease) - except Exception: - # Preserve generation failures/cancellation. The server keeps a - # failed release fenced and publication will time out fail-closed. - logger.exception("Failed to release KV weight lease %s", lease.get("lease_id")) - def _release(self, engine, key): async def cleanup(): try: diff --git a/infra/scheduling/kv_identity.py b/infra/scheduling/kv_identity.py deleted file mode 100644 index d956ca6..0000000 --- a/infra/scheduling/kv_identity.py +++ /dev/null @@ -1,28 +0,0 @@ -"""KV reuse requires known weight and service identities on both peers.""" - - -def same_snapshot(source, target): - return bool(isinstance(source, dict) and isinstance(target, dict) - and isinstance(source.get("snapshot_id"), str) and source["snapshot_id"] - and source["snapshot_id"] == target.get("snapshot_id") - and type(source.get("optimizer_revision")) is int - and type(target.get("optimizer_revision")) is int - and source["optimizer_revision"] >= 0 - and source["optimizer_revision"] == target["optimizer_revision"]) - - -def snapshot_identity(value): - return {name: value.get(name) for name in ( - "snapshot_id", "optimizer_revision", "server_instance_id", "server_generation")} - - -def same_incarnation(source, current): - """An old service generation cannot attest a cache after a reload/restart.""" - return bool(same_snapshot(source, current) - and isinstance(source.get("server_instance_id"), str) - and source["server_instance_id"] - and source["server_instance_id"] == current.get("server_instance_id") - and type(source.get("server_generation")) is int - and type(current.get("server_generation")) is int - and source["server_generation"] >= 0 - and source["server_generation"] == current["server_generation"]) diff --git a/infra/scheduling/vllm_cpu_offload.py b/infra/scheduling/vllm_cpu_offload.py index 491b3cc..220a3f1 100644 --- a/infra/scheduling/vllm_cpu_offload.py +++ b/infra/scheduling/vllm_cpu_offload.py @@ -17,8 +17,6 @@ from dataclasses import dataclass, field from pathlib import Path -from .kv_identity import same_snapshot, same_incarnation - import torch from vllm.distributed import get_tensor_model_parallel_rank from vllm.distributed.kv_transfer.kv_connector.v1.base import ( @@ -89,11 +87,9 @@ def request_finished(self, request, block_ids): "request_id": request.request_id, "key": key, "block_ids": list(block_ids)[:n // self.block_size], "tokens": list(request.all_token_ids[:n]), "source_tp": self.tp, - "libra_kv_identity": params.get("libra_kv_identity"), }) return True, {"libra_offload_key": key, "source_tp": self.tp, - "cached_tokens": n, "libra_imported_tokens": loaded, - "libra_kv_identity": params.get("libra_kv_identity")} + "cached_tokens": n, "libra_imported_tokens": loaded} def get_num_new_matched_tokens(self, request, num_computed_tokens): params = request.kv_transfer_params or {} @@ -103,23 +99,11 @@ def get_num_new_matched_tokens(self, request, num_computed_tokens): directory = cache_directory(self.root, key) try: manifest = json.loads((directory / "rank_0.json").read_text()) - if (manifest.get("block_size") != self.block_size - or not same_snapshot(manifest.get("libra_kv_identity"), - params.get("libra_kv_identity")) - or not same_incarnation(manifest.get("libra_kv_identity"), - params.get("libra_kv_source_identity"))): + if manifest.get("block_size") != self.block_size: return 0, False source_tp = int(manifest["source_tp"]) - if not 1 <= source_tp <= 64: + if any(not (directory / f"rank_{rank}.json").exists() for rank in range(source_tp)): return 0, False - for rank in range(source_tp): - shard = json.loads((directory / f"rank_{rank}.json").read_text()) - if (shard.get("tokens") != manifest["tokens"] - or shard.get("source_tp") != source_tp - or shard.get("block_size") != self.block_size - or not same_incarnation(shard.get("libra_kv_identity"), - manifest.get("libra_kv_identity"))): - return 0, False tokens = manifest["tokens"] n = min(len(tokens), len(request.prompt_token_ids) - 1) n = n // self.block_size * self.block_size @@ -130,11 +114,9 @@ def get_num_new_matched_tokens(self, request, num_computed_tokens): self.matches[request.request_id] = { "key": key, "source_tp": source_tp, "num_tokens": n, "request_id": request.request_id, - "libra_kv_identity": params.get("libra_kv_identity"), - "libra_kv_source_identity": params.get("libra_kv_source_identity"), } return n - num_computed_tokens, False - except (FileNotFoundError, json.JSONDecodeError, KeyError, TypeError, ValueError): + except (FileNotFoundError, json.JSONDecodeError): return 0, False def update_state_after_alloc(self, request, blocks, num_external_tokens): @@ -211,7 +193,7 @@ def _publish(self, host, item, started, dma_done): temporary.replace(directory / f"rank_{self.rank}.pt") atomic_json(directory / f"rank_{self.rank}.json", { "tokens": item["tokens"], "source_tp": item["source_tp"], - "block_size": self.block_size, "libra_kv_identity": item.get("libra_kv_identity"), + "block_size": self.block_size, "offload_started_at": started, "dma_done_at": dma_done, "ready_at": time.time(), "bytes": sum(t.numel() * t.element_size() for t in host.values()), }) @@ -221,11 +203,6 @@ def _publish(self, host, item, started, dma_done): def _reload(self, item): directory = cache_directory(self.root, item["key"]) - for rank in range(item["source_tp"]): - manifest = json.loads((directory / f"rank_{rank}.json").read_text()) - if (not same_incarnation(manifest.get("libra_kv_identity"), item.get("libra_kv_source_identity")) - or not same_snapshot(manifest.get("libra_kv_identity"), item.get("libra_kv_identity"))): - raise RuntimeError("CPU KV identity changed before worker reload") source_heads = self.num_heads // item["source_tp"] target_heads = self.num_heads // self.tp first_head, last_head = self.rank * target_heads, (self.rank + 1) * target_heads diff --git a/scripts/serving_protocol.py b/scripts/serving_protocol.py index c438c7c..b4f87f6 100644 --- a/scripts/serving_protocol.py +++ b/scripts/serving_protocol.py @@ -1,7 +1,7 @@ -"""Shared serving identity, request/KV leases and global reload receipts. +"""Shared serving identity, request leases and global reload receipts. -Generation and reload share one boundary so cached KV remains tied to the -weights and service incarnation that produced it. +Generation and reload share one boundary. Request reservations keep verified +service identity stable until the reserved request has completed. """ from __future__ import annotations @@ -134,19 +134,6 @@ def generation_headers(state, request_id, generation, status, streaming=False): return [(key.encode(), value.encode()) for key, value in values.items()] -def fence_kv_payload(payload, state): - """The service owns target weight identity; stale imports must recompute.""" - from RL_Framework.infra.scheduling.kv_identity import same_snapshot, snapshot_identity - params = payload.get('kv_transfer_params') if isinstance(payload, dict) else None - if isinstance(params, dict) and ('libra_offload_key' in params or 'libra_load_key' in params): - identity = snapshot_identity(state.status()) - if not same_snapshot(params.get('libra_kv_identity'), identity): - params.pop('libra_load_key', None) - params.pop('libra_kv_source_identity', None) - params['libra_kv_identity'] = identity - return payload - - class GenerationMiddleware: """Admit request lifetimes through the same boundary used by weight reload.""" def __init__(self, app, state): @@ -186,33 +173,6 @@ async def observe(message): for part in bodies: await send(part) try: - if receive is None: - return await self.app(scope, receive, observe) - chunks = [] - while True: - message = await receive() - if message['type'] == 'http.disconnect': - return - chunks.append(message.get('body', b'')) - if not message.get('more_body', False): - break - body = b''.join(chunks) - try: - payload = json.loads(body) - except (ValueError, TypeError): - payload = None - if isinstance(payload, dict) and isinstance(payload.get('kv_transfer_params'), dict): - body = json.dumps(fence_kv_payload(payload, self.state)).encode() - scope = dict(scope, headers=[(key, value) for key, value in scope.get('headers', []) - if key.lower() != b'content-length'] - + [(b'content-length', str(len(body)).encode())]) - consumed = False - async def replay(): - nonlocal consumed - if not consumed: - consumed = True - return dict(type='http.request', body=body, more_body=False) - return await receive() - await self.app(scope, replay, observe) + await self.app(scope, receive, observe) finally: await self.state.leave() diff --git a/scripts/test_cpu_ci.sh b/scripts/test_cpu_ci.sh index e9726e8..efec092 100644 --- a/scripts/test_cpu_ci.sh +++ b/scripts/test_cpu_ci.sh @@ -15,7 +15,6 @@ python -m pytest -q --strict-markers --junitxml=reports/cpu-tests.xml \ tests/test_nccl_store_lifecycle.py \ tests/test_nccl_store_timeout.py \ tests/test_serving_protocol.py \ - tests/test_versioned_kv_compatibility.py \ tests/test_cpu_offload_backend.py \ tests/test_r2e_patch_harness.py \ tests/test_r2e_gym_workflow.py diff --git a/tests/test_cmlfq_cost_scheduler.py b/tests/test_cmlfq_cost_scheduler.py index 461cedd..a30505c 100644 --- a/tests/test_cmlfq_cost_scheduler.py +++ b/tests/test_cmlfq_cost_scheduler.py @@ -205,17 +205,8 @@ def engine(self, handler, transfer=True): "nixl" if transfer else "recompute", [(1, 4)], ) engine.engines = [VLLMRolloutEngine(host="host-a"), VLLMRolloutEngine(host="host-b")] - def versioned_handler(request): - if request.url.path == "/libra/kv/lease": - return httpx.Response(200, json=dict( - lease_id=json.loads(request.content)["lease_id"], snapshot_id="shared-snapshot", - optimizer_revision=0, server_instance_id=request.url.host, server_generation=0, - )) - if request.method == "DELETE" and request.url.path.startswith("/libra/kv/lease/"): - return httpx.Response(200, json={"released": True}) - return handler(request) for backend in engine.engines: - backend.http_client = httpx.AsyncClient(transport=httpx.MockTransport(versioned_handler)) + backend.http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) backend._http_client_loop = asyncio.get_running_loop() engine._configure_cost_runtime() return engine diff --git a/tests/test_cpu_offload_backend.py b/tests/test_cpu_offload_backend.py index 7e6c0fc..c7dbc7e 100644 --- a/tests/test_cpu_offload_backend.py +++ b/tests/test_cpu_offload_backend.py @@ -1,6 +1,5 @@ import asyncio import unittest -import uuid from RL_Framework.engine.cmlfq_backend import CMLFQGenerationBackend @@ -10,20 +9,6 @@ def __init__(self): self.calls = [] self.released = [] self.ready = asyncio.Event() - self.instance_id = uuid.uuid4().hex - self.generation = 0 - self.revision = 0 - self.leases = set() - - async def acquire_kv_lease(self): - lease_id = uuid.uuid4().hex - self.leases.add(lease_id) - return dict(lease_id=lease_id, snapshot_id=f"snapshot-{self.revision}", - optimizer_revision=self.revision, server_instance_id=self.instance_id, - server_generation=self.generation) - - async def release_kv_lease(self, lease): - self.leases.discard(lease["lease_id"]) async def generate(self, prompt, kv_transfer_params, **kwargs): self.calls.append(kv_transfer_params) diff --git a/tests/test_serving_protocol.py b/tests/test_serving_protocol.py index ddf26ef..9009720 100644 --- a/tests/test_serving_protocol.py +++ b/tests/test_serving_protocol.py @@ -1,4 +1,4 @@ -"""Shared serving correctness across arbitrary TP layouts, reload and KV leases.""" +"""Shared serving correctness across arbitrary TP layouts, reload and request leases.""" import asyncio import importlib.util @@ -10,7 +10,7 @@ import pytest import torch -from RL_Framework.scripts.serving_protocol import ServingState, GenerationMiddleware, fence_kv_payload +from RL_Framework.scripts.serving_protocol import ServingState, GenerationMiddleware from RL_Framework.vllm_hot_reload import InplaceReloadWorkerExtension @@ -61,7 +61,7 @@ async def run(): asyncio.run(run()) -def test_reload_drains_inflight_requests_and_kv_continuations_before_install(): +def test_reload_drains_inflight_requests_and_reserved_dispatch_before_install(): async def run(): state = ServingState(2) lease = await state.acquire_kv_lease() @@ -74,7 +74,7 @@ async def run(): waiting = asyncio.create_task(state.enter()) await asyncio.sleep(0) assert not waiting.done() - assert await state.enter(lease['lease_id']) == 0 # Existing KV continuation progresses. + assert await state.enter(lease['lease_id']) == 0 # A previously reserved request progresses. await state.leave() await state.leave() assert not installing.done() @@ -99,17 +99,7 @@ async def run(): asyncio.run(run()) -def test_server_owns_target_kv_identity_and_strips_stale_import(): - state = ServingState(2, '/models/test') - payload = dict(kv_transfer_params=dict(libra_offload_key='save', libra_load_key='old', - libra_kv_identity=dict(snapshot_id='old', optimizer_revision=-1), libra_kv_source_identity={'old': True})) - fenced = fence_kv_payload(payload, state)['kv_transfer_params'] - assert 'libra_load_key' not in fenced and 'libra_kv_source_identity' not in fenced - assert fenced['libra_kv_identity']['server_instance_id'] == state.server_instance_id - assert fenced['libra_offload_key'] == 'save' - - -def test_generation_headers_match_actual_completion_and_body_identity_is_authoritative(): +def test_generation_headers_attest_completion_without_changing_request_body(): async def run(): state, sent = ServingState(1), [] received_payload = [] @@ -118,7 +108,7 @@ async def app(scope, receive, send): await send(dict(type='http.response.start', status=200, headers=[(b'content-type', b'application/json')])) await send(dict(type='http.response.body', body=b'{"choices":[]}')) middleware = GenerationMiddleware(app, state) - payload = dict(prompt=[1], kv_transfer_params=dict(libra_offload_key='save', libra_load_key='old')) + payload = dict(prompt=[1, 2, 3], seed=42, temperature=.8, top_p=.95, custom_field='unchanged') async def receive(): return dict(type='http.request', body=json.dumps(payload).encode(), more_body=False) async def send(message): @@ -132,9 +122,7 @@ async def send(message): assert headers[b'x-libra-service-generation-verified'] == b'1' assert headers[b'x-libra-reload-overlap'] == b'0' assert state.active_requests == 0 - params = received_payload[0]['kv_transfer_params'] - assert 'libra_load_key' not in params - assert params['libra_kv_identity']['server_instance_id'] == state.server_instance_id + assert received_payload == [payload] asyncio.run(run()) @@ -253,7 +241,7 @@ async def rpc(name, timeout, args=()): return app -def test_real_http_routes_fence_reload_by_kv_lease_and_attest_physical_workers(monkeypatch): +def test_real_http_routes_fence_reload_by_request_lease_and_attest_physical_workers(monkeypatch): import httpx app = load_serving_app(monkeypatch) async def run(): @@ -283,11 +271,10 @@ async def run(): assert live['server_generation'] == 1 and live['server_generation_verified'] assert live['workers'] == reply.json()['workers'] response = await client.post('/v1/completions', headers={'X-Libra-Request-Id': 'new-request'}, - json=dict(prompt=[1], kv_transfer_params=dict(libra_offload_key='next', libra_load_key='old', - libra_kv_identity=lease))) + json=dict(prompt=[1, 2], seed=42)) assert response.headers['x-libra-service-generation-verified'] == '1' assert response.headers['x-libra-service-generation'] == '1' - assert 'libra_load_key' not in response.json()['received']['kv_transfer_params'] + assert response.json()['received'] == dict(prompt=[1, 2], seed=42) asyncio.run(run()) @@ -304,3 +291,24 @@ async def run(): assert state['server_generation'] == 0 assert (await client.post('/v1/completions', json=dict(prompt=[1]))).status_code == 409 asyncio.run(run()) + + +def test_cancelled_lease_response_withdraws_server_reservation(): + import httpx + from RL_Framework.engine.rollout_engine import VLLMRolloutEngine + async def run(): + state = ServingState(1) + async def handler(request): + if request.method == "POST": + await state.acquire_kv_lease(json.loads(request.content)["lease_id"]) + raise asyncio.CancelledError() + await state.release_kv_lease(request.url.path.rsplit("/", 1)[-1]) + return httpx.Response(200, json={"released": True}) + engine = VLLMRolloutEngine(host="test") + engine.http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + engine._http_client_loop = asyncio.get_running_loop() + with pytest.raises(asyncio.CancelledError): + await engine.acquire_kv_lease() + assert not state.kv_leases + await engine.close() + asyncio.run(run()) diff --git a/tests/test_versioned_kv_compatibility.py b/tests/test_versioned_kv_compatibility.py deleted file mode 100644 index 443c1ed..0000000 --- a/tests/test_versioned_kv_compatibility.py +++ /dev/null @@ -1,223 +0,0 @@ -"""Version/incarnation fencing through real KV clients and the native reload gate.""" -import asyncio -import importlib.util -import json -from pathlib import Path -import sys -from types import ModuleType, SimpleNamespace - -import httpx -from fastapi import FastAPI, HTTPException, Request -import pytest -import torch - -from RL_Framework.engine.cmlfq_backend import CMLFQGenerationBackend -from RL_Framework.engine.rollout_engine import VLLMRolloutEngine -from RL_Framework.infra.scheduling.kv_identity import same_snapshot, same_incarnation -from RL_Framework.scripts.serving_protocol import ServingState, GenerationMiddleware -from test_cpu_offload_backend import FakeCPUBackend - - -def peer(revision=7, before_completion=None): - state = ServingState(1, "same-model") - state.snapshot_id, state.optimizer_revision = f"snapshot-{revision}", revision - app, calls = FastAPI(), [] - - @app.post("/libra/kv/lease") - async def acquire(request: Request): - try: - return await state.acquire_kv_lease((await request.json()).get("lease_id")) - except RuntimeError as error: - raise HTTPException(409, str(error)) - - @app.delete("/libra/kv/lease/{lease_id}") - async def release(lease_id: str): - await state.release_kv_lease(lease_id) - return {"released": True} - - @app.get("/libra/kv/identity") - async def identity(): - return state.status() - - @app.post("/v1/completions") - async def complete(request: Request): - calls.append(await request.json()) - if before_completion: - await before_completion() - return dict(choices=[dict(text="answer", logprobs={})], kv_transfer_params=dict( - remote_engine_id="source", remote_request_id="request", remote_host="host", - remote_port=5555, remote_block_ids=[[1, 2]])) - - app.add_middleware(GenerationMiddleware, state=state) - engine = VLLMRolloutEngine(host="test") - engine.http_client = httpx.AsyncClient(transport=httpx.ASGITransport(app=app)) - engine._http_client_loop = asyncio.get_running_loop() - return engine, state, calls - - -def test_identity_distinguishes_shared_weights_and_service_incarnations(): - identity = dict(snapshot_id="shared", optimizer_revision=7, server_instance_id="a", server_generation=3) - other = dict(identity, server_instance_id="b") - assert same_snapshot(identity, other) - assert not same_incarnation(identity, other) - assert not same_incarnation(identity, dict(identity, server_generation=4)) - assert not same_snapshot(identity, dict(identity, optimizer_revision=8)) - assert not same_snapshot({}, {}) - - -@pytest.mark.parametrize("target_revision", [7, 8]) -def test_nixl_same_snapshot_transfers_but_new_revision_recomputes(target_revision): - async def run(): - source, source_state, source_calls = peer() - target, target_state, target_calls = peer(target_revision) - backend = CMLFQGenerationBackend("nixl", [(1, 4)]) - output = await backend.generate(source, target, "full prompt", "transfer", prompt_token_ids=[1, 2, 3]) - assert output.path == ("transfer" if target_revision == 7 else "recompute") - assert bool(source_calls) == (target_revision == 7) - assert target_calls[-1]["prompt"] == [1, 2, 3] - assert bool(target_calls[-1].get("kv_transfer_params")) == (target_revision == 7) - assert not source_state.kv_leases and not target_state.kv_leases - await source.close() - await target.close() - asyncio.run(run()) - - -@pytest.mark.parametrize("change,reason", [("revision", "snapshot_mismatch"), ("generation", "incarnation_mismatch"), - ("restart", "incarnation_mismatch")]) -def test_cpu_old_cache_recomputes_after_reload_or_restart(change, reason): - async def run(): - source, target = FakeCPUBackend(), FakeCPUBackend() - source.ready.set() - target.ready.set() - backend = CMLFQGenerationBackend("cpu_offload", [(1, 4)]) - await backend.generate(source, source, "prompt", "stay", "request") - if change == "revision": - target.revision += 1 - elif change == "generation": - source.generation += 1 - else: - source.instance_id = "restarted" - output = await backend.generate(source, target, "prompt+tool", "transfer", "request") - assert output.path == "recompute" and output.fallback_reason == reason - assert "libra_load_key" not in target.calls[-1] - assert not source.leases and not target.leases - await backend.close() - asyncio.run(run()) - - -def test_cancelled_cpu_wait_releases_both_source_and_target_leases(): - async def run(): - source, target = FakeCPUBackend(), FakeCPUBackend() - backend = CMLFQGenerationBackend("cpu_offload", [(1, 4)]) - await backend.generate(source, source, "prompt", "stay", "request") - task = asyncio.create_task(backend.generate(source, target, "prompt+tool", "transfer", "request")) - await asyncio.sleep(0) - assert source.leases and target.leases - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - assert not source.leases and not target.leases - await backend.close() - asyncio.run(run()) - - -def load_connector(monkeypatch): - name = "vllm.distributed.kv_transfer.kv_connector.v1.base" - base = ModuleType(name) - base.KVConnectorBase_V1, base.KVConnectorMetadata = object, object - base.KVConnectorRole = SimpleNamespace(WORKER="worker") - distributed = ModuleType("vllm.distributed") - distributed.get_tensor_model_parallel_rank = lambda: 0 - monkeypatch.setitem(sys.modules, name, base) - monkeypatch.setitem(sys.modules, "vllm.distributed", distributed) - path = Path(__file__).parents[1] / "infra/scheduling/vllm_cpu_offload.py" - module_name = "RL_Framework.infra.scheduling.kv_compat_connector_fixture" - spec = importlib.util.spec_from_file_location(module_name, path) - module = importlib.util.module_from_spec(spec) - monkeypatch.setitem(sys.modules, module_name, module) - spec.loader.exec_module(module) - return module - - -def connector_fixture(monkeypatch, tmp_path, source_tp=2, target_tp=4, rank=0): - module = load_connector(monkeypatch) - connector = module.LibraCPUOffloadConnector.__new__(module.LibraCPUOffloadConnector) - connector.root, connector.tp, connector.rank = str(tmp_path), target_tp, rank - connector.block_size, connector.num_heads = 2, 8 - connector.matches, connector.imported = {}, {} - key = "a" * 32 - directory = tmp_path / key - directory.mkdir(exist_ok=True) - source_identity = dict(snapshot_id="snapshot", optimizer_revision=7, server_instance_id="source", server_generation=3) - target_identity = dict(source_identity, server_instance_id="target", server_generation=8) - manifest = dict(tokens=list(range(4)), source_tp=source_tp, block_size=2, libra_kv_identity=source_identity) - for source_rank in range(source_tp): - (directory / f"rank_{source_rank}.json").write_text(json.dumps(manifest)) - request = SimpleNamespace(kv_transfer_params=dict(libra_load_key=key, libra_kv_identity=target_identity, - libra_kv_source_identity=source_identity), request_id="request", prompt_token_ids=list(range(5))) - return module, connector, directory, manifest, request - - -def test_cpu_connector_checks_all_tp_source_incarnations(monkeypatch, tmp_path): - _, connector, directory, manifest, request = connector_fixture(monkeypatch, tmp_path) - assert connector.get_num_new_matched_tokens(request, 0) == (4, False) - changed = dict(manifest, libra_kv_identity=dict(manifest["libra_kv_identity"], server_generation=4)) - (directory / "rank_1.json").write_text(json.dumps(changed)) - assert connector.get_num_new_matched_tokens(request, 0) == (0, False) - - -@pytest.mark.parametrize("source_tp,target_tp", [(1, 2), (2, 1), (2, 4), (4, 2)]) -def test_cpu_real_tensor_head_resharding_keeps_identity_checks(monkeypatch, tmp_path, source_tp, target_tp): - module, connector, directory, _, request = connector_fixture(monkeypatch, tmp_path, source_tp, target_tp) - full = torch.arange(2 * 2 * 2 * 8 * 4, dtype=torch.float32).reshape(2, 2, 2, 8, 4) - for rank, shard in enumerate(full.chunk(source_tp, dim=-2)): - torch.save({"layer": shard.contiguous()}, directory / f"rank_{rank}.pt") - # Exercise the production file loading/head slicing/index-copy code on CPU; - # CUDA DMA and pinned transfers remain covered by the GPU regression. - monkeypatch.setattr(torch.Tensor, "pin_memory", lambda value: value) - module.torch = SimpleNamespace(tensor=lambda values, **kwargs: torch.tensor(values, dtype=kwargs.get("dtype")), - long=torch.long, load=torch.load, cat=torch.cat) - for rank in range(target_tp): - connector.rank = rank - destination = torch.zeros(2, 3, 2, 8 // target_tp, 4) - connector.caches = {"layer": destination} - connector._reload(dict(key=directory.name, source_tp=source_tp, num_tokens=4, block_ids=[2, 0], - **request.kv_transfer_params)) - torch.testing.assert_close(destination[:, [2, 0]], full.chunk(target_tp, dim=-2)[rank]) - - -@pytest.mark.parametrize("unknown", [None, [], "snapshot", {}, {"snapshot_id": "snapshot", "optimizer_revision": True}]) -def test_unknown_kv_identities_are_rejected_without_crashing(unknown): - identity = dict(snapshot_id="snapshot", optimizer_revision=1, server_instance_id="s", server_generation=0) - assert not same_snapshot(unknown, identity) - assert not same_snapshot(identity, unknown) - assert not same_incarnation(unknown, identity) - - -def test_worker_fences_changed_incarnation_before_touching_cuda(monkeypatch, tmp_path): - _, connector, directory, manifest, request = connector_fixture(monkeypatch, tmp_path) - assert connector.get_num_new_matched_tokens(request, 0) == (4, False) - changed = dict(manifest, libra_kv_identity=dict(manifest["libra_kv_identity"], server_instance_id="restarted")) - (directory / "rank_1.json").write_text(json.dumps(changed)) - with pytest.raises(RuntimeError, match="identity changed"): - connector._reload(dict(key=directory.name, source_tp=2, num_tokens=4, block_ids=[0, 1], - **request.kv_transfer_params)) - - -def test_cancelled_lease_response_withdraws_server_reservation(): - async def run(): - state = ServingState(1) - async def handler(request): - if request.method == "POST": - await state.acquire_kv_lease(json.loads(request.content)["lease_id"]) - raise asyncio.CancelledError() - await state.release_kv_lease(request.url.path.rsplit("/", 1)[-1]) - return httpx.Response(200, json={"released": True}) - engine = VLLMRolloutEngine(host="test") - engine.http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) - engine._http_client_loop = asyncio.get_running_loop() - with pytest.raises(asyncio.CancelledError): - await engine.acquire_kv_lease() - assert not state.kv_leases - await engine.close() - asyncio.run(run()) diff --git a/tests/test_vllm_cpu_offload.py b/tests/test_vllm_cpu_offload.py index 9557caa..8e08fb8 100644 --- a/tests/test_vllm_cpu_offload.py +++ b/tests/test_vllm_cpu_offload.py @@ -13,10 +13,6 @@ from RL_Framework.infra.scheduling.vllm_cpu_offload import LibraCPUOffloadConnector, cache_directory -IDENTITY = dict(snapshot_id="test-snapshot", optimizer_revision=1, - server_instance_id="source", server_generation=0) - - def bare_connector(root, tp=1, rank=0): connector = object.__new__(LibraCPUOffloadConnector) connector.root, connector.tp, connector.rank = root, tp, rank @@ -32,12 +28,9 @@ def test_prefix_identity_and_complete_block_boundaries(): directory.mkdir() (directory / "rank_0.json").write_text(json.dumps({ "tokens": list(range(64)), "source_tp": 1, "block_size": 16, - "libra_kv_identity": IDENTITY, })) connector = bare_connector(root) - request = SimpleNamespace(kv_transfer_params={"libra_load_key": key, - "libra_kv_identity": IDENTITY, - "libra_kv_source_identity": IDENTITY}, + request = SimpleNamespace(kv_transfer_params={"libra_load_key": key}, request_id="r", prompt_token_ids=list(range(65))) assert connector.get_num_new_matched_tokens(request, 16) == (48, False) request.prompt_token_ids[1] = 999 @@ -65,18 +58,12 @@ def test_real_cuda_cpu_head_resharding(source_tp, target_tp): pinned.copy_(gpu, non_blocking=True) stream.synchronize() torch.save({"layer": pinned}, directory / f"rank_{rank}.pt") - (directory / f"rank_{rank}.json").write_text(json.dumps({ - "tokens": list(range(32)), "source_tp": source_tp, "block_size": 16, - "libra_kv_identity": IDENTITY, - })) for rank in range(target_tp): with torch.cuda.device(rank): connector = bare_connector(root, target_tp, rank) destination = torch.zeros(2, 3, 16, 8 // target_tp, 4, device="cuda") connector.caches = {"layer": destination} connector._reload({"key": key, "source_tp": source_tp, - "num_tokens": 32, "block_ids": [2, 0], - "libra_kv_identity": IDENTITY, - "libra_kv_source_identity": IDENTITY}) + "num_tokens": 32, "block_ids": [2, 0]}) observed = destination[:, [2, 0]].cpu() assert torch.equal(observed, full.chunk(target_tp, dim=-2)[rank])