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/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..b4f87f6 --- /dev/null +++ b/scripts/serving_protocol.py @@ -0,0 +1,178 @@ +"""Shared serving identity, request leases and global reload receipts. + +Generation and reload share one boundary. Request reservations keep verified +service identity stable until the reserved request has completed. +""" +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()] + + +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: + 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 c3ceb95..efec092 100644 --- a/scripts/test_cpu_ci.sh +++ b/scripts/test_cpu_ci.sh @@ -11,6 +11,10 @@ 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_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_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..9009720 --- /dev/null +++ b/tests/test_serving_protocol.py @@ -0,0 +1,314 @@ +"""Shared serving correctness across arbitrary TP layouts, reload and request 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 +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_reserved_dispatch_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 # A previously reserved request 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_generation_headers_attest_completion_without_changing_request_body(): + 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, 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): + 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 + assert received_payload == [payload] + 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_request_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, 2], seed=42)) + assert response.headers['x-libra-service-generation-verified'] == '1' + assert response.headers['x-libra-service-generation'] == '1' + assert response.json()['received'] == dict(prompt=[1, 2], seed=42) + 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()) + + +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/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)