Verify complete TP publication and physical service lifecycle

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()

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):

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

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",
 ]

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.

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} "