From 536cbefb8d5c04f0b7d67889bef735ae8bfa8c3f Mon Sep 17 00:00:00 2001 From: Vivek Kalyan Date: Thu, 23 Jul 2026 18:42:58 -0700 Subject: [PATCH 1/9] feat: Record LoRA weight update phase metrics --- src/art/megatron/service.py | 92 ++++++++++++++++--- src/art/megatron/train.py | 10 +- src/art/megatron/weights/lora_publish.py | 55 ++++++++++- .../src/art_vllm_runtime/dedicated_server.py | 24 ++++- vllm_runtime/src/art_vllm_runtime/metrics.py | 22 +++++ 5 files changed, 186 insertions(+), 17 deletions(-) diff --git a/src/art/megatron/service.py b/src/art/megatron/service.py index fe04050f9..49f8f1f58 100644 --- a/src/art/megatron/service.py +++ b/src/art/megatron/service.py @@ -8,6 +8,7 @@ import socket import subprocess import sys +import time from typing import Any, AsyncIterator, Literal, TypedDict, cast from urllib.parse import urlparse import uuid @@ -767,7 +768,9 @@ async def _reload_adapter(self, checkpoint_path: str, step: int) -> None: self._latest_step = step self._loaded_adapter_steps.add(step) - async def _update_in_flight_adapter(self, checkpoint_path: str, step: int) -> None: + async def _update_in_flight_adapter( + self, checkpoint_path: str, step: int + ) -> dict[str, float]: import httpx self._raise_if_child_failed() @@ -777,6 +780,7 @@ async def _update_in_flight_adapter(self, checkpoint_path: str, step: int) -> No self.serving_capabilities.require( "policy_token_spans", operation="In-flight LoRA updates" ) + rpc_started = time.perf_counter() async with httpx.AsyncClient() as client: response = await client.post( f"{self._vllm_base_url}/art/in_flight_lora_update", @@ -790,16 +794,37 @@ async def _update_in_flight_adapter(self, checkpoint_path: str, step: int) -> No timeout=60.0, ) response.raise_for_status() + rpc_s = time.perf_counter() - rpc_started + json_response = getattr(response, "json", None) + response_body = json_response() if callable(json_response) else {} + runtime_timing = response_body.get("timing_s", {}) + if not isinstance(runtime_timing, dict): + runtime_timing = {} self._latest_step = step self._loaded_adapter_steps.add(step) + metrics = { + "time/weight_update_service_rpc_s": rpc_s, + "weight_update/policy_version": float(step), + } + for phase in ( + "begin_update", + "load_adapter", + "update_waiting_cache", + "commit_update", + "total", + ): + value = runtime_timing.get(phase) + if isinstance(value, (int, float)): + metrics[f"time/weight_update_vllm_{phase}_s"] = float(value) + return metrics async def _load_rollout_lora_for_step( self, checkpoint_path: str, step: int - ) -> None: + ) -> dict[str, float]: if self.rollout_weight_update_mode == "in_flight_lora": - await self._update_in_flight_adapter(checkpoint_path, step) - else: - await self._reload_adapter(checkpoint_path, step) + return await self._update_in_flight_adapter(checkpoint_path, step) + await self._reload_adapter(checkpoint_path, step) + return {} async def acquire_exact_adapter(self, step: int, checkpoint_path: str) -> str: if self.rollout_weights_mode != "lora": @@ -1156,16 +1181,27 @@ async def _handle_training_lora_ready( checkpoint_dir: str | None, staging_lora_path: str, step: int, - ) -> str: + ) -> tuple[str, dict[str, float]]: + metrics: dict[str, float] = {} if checkpoint_dir is None: checkpoint_dir = self._publish_staged_training_checkpoint( staging_lora_path=staging_lora_path, step=step, ) if self.is_dedicated and self.rollout_weights_mode == "lora": - await self._load_rollout_lora_for_step(checkpoint_dir, step) + metrics = await self._load_rollout_lora_for_step(checkpoint_dir, step) self._status(f"Loaded checkpoint {step} into vLLM") - return checkpoint_dir + return checkpoint_dir, metrics + + @staticmethod + def _weight_update_metrics_from_event( + result: dict[str, Any], + ) -> dict[str, float]: + return { + key: float(value) + for key, value in result.items() + if key not in {"event", "step"} and isinstance(value, (int, float)) + } async def _finish_training_checkpoint( self, @@ -1326,6 +1362,7 @@ async def train( write_megatron_job(job, job_path=job_path) checkpoint_dir: str | None = None optimizer_world_size: int | None = None + pending_metrics: dict[str, float] | None = None async for result in stream_megatron_job( job, job_path=job_path, @@ -1333,11 +1370,20 @@ async def train( process_log_path=self._megatron_log_path, ): if result.get("event") == LORA_READY_EVENT: - checkpoint_dir = await self._handle_training_lora_ready( + ( + checkpoint_dir, + install_metrics, + ) = await self._handle_training_lora_ready( checkpoint_dir=checkpoint_dir, staging_lora_path=staging_lora_path, step=next_step, ) + if pending_metrics is None: + pending_metrics = {} + pending_metrics.update( + self._weight_update_metrics_from_event(result) + ) + pending_metrics.update(install_metrics) continue if ( world_size := self._optimizer_ready_world_size( @@ -1346,7 +1392,14 @@ async def train( ) is not None: optimizer_world_size = world_size continue - yield {key: float(value) for key, value in result.items()} + if pending_metrics is not None: + yield pending_metrics + pending_metrics = { + key: float(value) for key, value in result.items() + } + + if pending_metrics is not None: + yield pending_metrics await self._finish_training_checkpoint( checkpoint_dir=checkpoint_dir, @@ -1386,6 +1439,7 @@ async def train( checkpoint_dir = None optimizer_world_size = None + pending_metrics = None async for result in stream_megatron_job( job, job_path=job_path, @@ -1393,11 +1447,20 @@ async def train( process_log_path=self._megatron_log_path, ): if result.get("event") == LORA_READY_EVENT: - checkpoint_dir = await self._handle_training_lora_ready( + ( + checkpoint_dir, + install_metrics, + ) = await self._handle_training_lora_ready( checkpoint_dir=checkpoint_dir, staging_lora_path=staging_lora_path, step=next_step, ) + if pending_metrics is None: + pending_metrics = {} + pending_metrics.update( + self._weight_update_metrics_from_event(result) + ) + pending_metrics.update(install_metrics) continue if ( world_size := self._optimizer_ready_world_size( @@ -1406,7 +1469,12 @@ async def train( ) is not None: optimizer_world_size = world_size continue - yield {key: float(value) for key, value in result.items()} + if pending_metrics is not None: + yield pending_metrics + pending_metrics = {key: float(value) for key, value in result.items()} + + if pending_metrics is not None: + yield pending_metrics await self._finish_training_checkpoint( checkpoint_dir=checkpoint_dir, diff --git a/src/art/megatron/train.py b/src/art/megatron/train.py index 284620115..6868eee33 100644 --- a/src/art/megatron/train.py +++ b/src/art/megatron/train.py @@ -1026,7 +1026,7 @@ def _save_lora_and_optimizer( optimizer_ready_log_path: str | None = None, ) -> None: assert runtime.optimizer is not None - save_vllm_lora_from_model( + publish_metrics = save_vllm_lora_from_model( model=runtime.model, adapter_dtypes=adapter_dtypes, handler=runtime.model_support_handler, @@ -1036,7 +1036,13 @@ def _save_lora_and_optimizer( world_size=runtime.world_size, ) if lora_ready_log_path is not None and runtime.rank == 0: - _write_job_event(lora_ready_log_path, LORA_READY_EVENT, step=step) + assert publish_metrics is not None + _write_job_event( + lora_ready_log_path, + LORA_READY_EVENT, + step=step, + **publish_metrics.as_training_metrics(), + ) if _should_save_optimizer( runtime, step=step, diff --git a/src/art/megatron/weights/lora_publish.py b/src/art/megatron/weights/lora_publish.py index e9d8b4b08..715521e7d 100644 --- a/src/art/megatron/weights/lora_publish.py +++ b/src/art/megatron/weights/lora_publish.py @@ -1,4 +1,7 @@ from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from pathlib import Path +import time from typing import Any, NamedTuple import torch @@ -19,6 +22,34 @@ from art.megatron.training.model_chunks import ModelChunks +@dataclass(frozen=True) +class LoraPublishMetrics: + """Phase-level measurements for one trainer-side LoRA publication.""" + + gather_pack_s: float + stage_to_cpu_s: float + write_s: float + total_s: float + logical_bytes: int + transported_bytes: int + tensor_count: int + + def as_training_metrics(self) -> dict[str, int | float]: + return { + "time/weight_update_trainer_gather_pack_s": self.gather_pack_s, + "time/weight_update_trainer_stage_to_cpu_s": self.stage_to_cpu_s, + "time/weight_update_trainer_write_s": self.write_s, + "time/weight_update_trainer_publish_s": self.total_s, + "weight_update/logical_bytes": self.logical_bytes, + "weight_update/transported_bytes": self.transported_bytes, + "weight_update/tensor_count": self.tensor_count, + } + + +def _tensor_bytes(tensor: torch.Tensor) -> int: + return tensor.numel() * tensor.element_size() + + class PackedExpertShardMeta(NamedTuple): key: str owner_rank: int @@ -757,7 +788,9 @@ def save_vllm_lora_from_model( rank: int, world_size: int, slot_ref: LoRASlotRef | None = None, -) -> None: +) -> LoraPublishMetrics | None: + total_started = time.perf_counter() + gather_pack_started = total_started result = build_vllm_lora_tensors_from_model( model=model, adapter_dtypes=adapter_dtypes, @@ -767,10 +800,28 @@ def save_vllm_lora_from_model( world_size=world_size, slot_ref=slot_ref, ) + gather_pack_s = time.perf_counter() - gather_pack_started if result is None: - return + return None vllm_tensors, published_config = result + logical_bytes = sum(_tensor_bytes(tensor) for tensor in vllm_tensors.values()) + + stage_started = time.perf_counter() stager = _PinnedCpuStager() published_tensors = _stage_published_tensors(vllm_tensors, stager) stager.finish() + stage_to_cpu_s = time.perf_counter() - stage_started + + write_started = time.perf_counter() save_vllm_lora_tensors(output_dir, published_tensors, published_config) + write_s = time.perf_counter() - write_started + adapter_path = Path(output_dir) / "adapter_model.safetensors" + return LoraPublishMetrics( + gather_pack_s=gather_pack_s, + stage_to_cpu_s=stage_to_cpu_s, + write_s=write_s, + total_s=time.perf_counter() - total_started, + logical_bytes=logical_bytes, + transported_bytes=adapter_path.stat().st_size, + tensor_count=len(vllm_tensors), + ) diff --git a/vllm_runtime/src/art_vllm_runtime/dedicated_server.py b/vllm_runtime/src/art_vllm_runtime/dedicated_server.py index 36b8a0ffd..96c29c2ef 100644 --- a/vllm_runtime/src/art_vllm_runtime/dedicated_server.py +++ b/vllm_runtime/src/art_vllm_runtime/dedicated_server.py @@ -5,6 +5,7 @@ from http import HTTPStatus import json import os +import time from fastapi.responses import JSONResponse from pydantic import BaseModel, Field @@ -227,8 +228,12 @@ async def in_flight_lora_update( models = raw_request.app.state.openai_serving_models engine_client = engine(raw_request) coordinator = lora_update_coordinator(models, engine_client) + update_started = time.perf_counter() + begin_started = update_started await coordinator.begin_update(lora_slot) + begin_s = time.perf_counter() - begin_started try: + load_started = time.perf_counter() load_result = await models.load_lora_adapter( LoadLoRAAdapterRequest( lora_name=lora_slot, @@ -238,23 +243,38 @@ async def in_flight_lora_update( ), base_model_name=body.base_model_name, ) + load_s = time.perf_counter() - load_started if isinstance(load_result, ErrorResponse): await coordinator.fail_update(lora_slot) return JSONResponse( content=load_result.model_dump(mode="python"), status_code=load_result.error.code, ) + cache_started = time.perf_counter() waiting_cache_salt = await engine_client.engine_core.call_utility_async( "art_update_waiting_lora_cache_salt", lora_slot, policy_version, ) + update_waiting_cache_s = time.perf_counter() - cache_started + commit_started = time.perf_counter() await coordinator.commit_update( lora_slot, policy_version, models.lora_requests[lora_slot], ) - from art_vllm_runtime.metrics import record_policy_cache_waiting_update + commit_s = time.perf_counter() - commit_started + timing_s = { + "begin_update": begin_s, + "load_adapter": load_s, + "update_waiting_cache": update_waiting_cache_s, + "commit_update": commit_s, + "total": time.perf_counter() - update_started, + } + from art_vllm_runtime.metrics import ( + record_policy_cache_waiting_update, + record_weight_update, + ) record_policy_cache_waiting_update( updated=int(waiting_cache_salt["updated_waiting_requests"]), @@ -262,6 +282,7 @@ async def in_flight_lora_update( waiting_cache_salt["skipped_started_waiting_requests"] ), ) + record_weight_update(timing_s) except BaseException: await coordinator.fail_update(lora_slot) raise @@ -272,6 +293,7 @@ async def in_flight_lora_update( "lora_slot": lora_slot, "policy_version": policy_version, "waiting_cache_salt": waiting_cache_salt, + "timing_s": timing_s, } ) diff --git a/vllm_runtime/src/art_vllm_runtime/metrics.py b/vllm_runtime/src/art_vllm_runtime/metrics.py index 0c8be3d8f..0f1db59c0 100644 --- a/vllm_runtime/src/art_vllm_runtime/metrics.py +++ b/vllm_runtime/src/art_vllm_runtime/metrics.py @@ -32,6 +32,12 @@ def __init__(self) -> None: "policy_cache_unsalted_lora_requests_total": 0.0, "policy_cache_waiting_requests_updated_total": 0.0, "policy_cache_started_waiting_requests_skipped_total": 0.0, + "weight_update_count_total": 0.0, + "weight_update_begin_s_total": 0.0, + "weight_update_load_adapter_s_total": 0.0, + "weight_update_waiting_cache_s_total": 0.0, + "weight_update_commit_s_total": 0.0, + "weight_update_s_total": 0.0, } def configure(self, vllm_config: Any, *, engine_idx: int) -> None: @@ -206,6 +212,18 @@ def record_policy_cache_waiting_update( float(skipped_started) ) + def record_weight_update(self, timing_s: dict[str, float]) -> None: + with self._lock: + self._counters["weight_update_count_total"] += 1.0 + for phase, metric in ( + ("begin_update", "weight_update_begin_s_total"), + ("load_adapter", "weight_update_load_adapter_s_total"), + ("update_waiting_cache", "weight_update_waiting_cache_s_total"), + ("commit_update", "weight_update_commit_s_total"), + ("total", "weight_update_s_total"), + ): + self._counters[metric] += float(timing_s.get(phase, 0.0)) + _STATE = _ArtRuntimeMetricsState() @@ -246,3 +264,7 @@ def record_policy_cache_waiting_update(*, updated: int, skipped_started: int) -> updated=updated, skipped_started=skipped_started, ) + + +def record_weight_update(timing_s: dict[str, float]) -> None: + _STATE.record_weight_update(timing_s) From 458cfb4afc2fcafc5617c9c7ebe8610153bcd78b Mon Sep 17 00:00:00 2001 From: Vivek Kalyan Date: Thu, 23 Jul 2026 18:43:05 -0700 Subject: [PATCH 2/9] feat: Add representative LoRA update replay harness --- scripts/benchmarks/replay_lora_updates.py | 6 + src/art/megatron/weights/update_replay.py | 394 ++++++++++++++++++++++ tests/unit/test_weight_update_replay.py | 131 +++++++ 3 files changed, 531 insertions(+) create mode 100644 scripts/benchmarks/replay_lora_updates.py create mode 100644 src/art/megatron/weights/update_replay.py create mode 100644 tests/unit/test_weight_update_replay.py diff --git a/scripts/benchmarks/replay_lora_updates.py b/scripts/benchmarks/replay_lora_updates.py new file mode 100644 index 000000000..f3f0b3abc --- /dev/null +++ b/scripts/benchmarks/replay_lora_updates.py @@ -0,0 +1,6 @@ +#!/usr/bin/env python3 + +from art.megatron.weights.update_replay import main + +if __name__ == "__main__": + main() diff --git a/src/art/megatron/weights/update_replay.py b/src/art/megatron/weights/update_replay.py new file mode 100644 index 000000000..cfbcd5a55 --- /dev/null +++ b/src/art/megatron/weights/update_replay.py @@ -0,0 +1,394 @@ +"""Replay in-flight LoRA updates against an ART vLLM runtime. + +The replay input is a sequence of real trainer-produced checkpoints. Keeping the +captured safetensors files intact preserves the serving tensor names, shapes, +dtypes, and byte volume without rerunning Megatron, reward, or judging. +""" + +from __future__ import annotations + +import argparse +import asyncio +from dataclasses import asdict, dataclass +import hashlib +import json +from pathlib import Path +import statistics +import time +from typing import Any + +import httpx +from safetensors import safe_open + +from art.megatron.model_support.lora_disk import load_adapter_config + + +@dataclass(frozen=True) +class TensorSpec: + name: str + shape: tuple[int, ...] + dtype: str + bytes: int + + +@dataclass(frozen=True) +class AdapterSnapshot: + path: str + file_sha256: str + manifest_sha256: str + logical_bytes: int + transported_bytes: int + tensor_count: int + rank: int + base_model: str + tensors: tuple[TensorSpec, ...] + + +@dataclass +class LoadCounters: + prompt_tokens: int = 0 + completion_tokens: int = 0 + requests: int = 0 + errors: int = 0 + + def __sub__(self, other: LoadCounters) -> LoadCounters: + return LoadCounters( + prompt_tokens=self.prompt_tokens - other.prompt_tokens, + completion_tokens=self.completion_tokens - other.completion_tokens, + requests=self.requests - other.requests, + errors=self.errors - other.errors, + ) + + +def _file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def inspect_adapter(path: str | Path) -> AdapterSnapshot: + adapter_dir = Path(path).resolve() + tensor_path = adapter_dir / "adapter_model.safetensors" + config = load_adapter_config(adapter_dir) + rank = int(config.get("r", config.get("rank", 0))) + base_model = str(config.get("base_model_name_or_path", "")) + tensors: list[TensorSpec] = [] + logical_bytes = 0 + with safe_open(tensor_path, framework="pt", device="cpu") as handle: + for name in sorted(handle.keys()): + tensor = handle.get_tensor(name) + tensor_bytes = tensor.numel() * tensor.element_size() + logical_bytes += tensor_bytes + tensors.append( + TensorSpec( + name=name, + shape=tuple(int(dim) for dim in tensor.shape), + dtype=str(tensor.dtype), + bytes=tensor_bytes, + ) + ) + manifest_payload = json.dumps( + [asdict(tensor) for tensor in tensors], + sort_keys=True, + separators=(",", ":"), + ).encode() + return AdapterSnapshot( + path=str(adapter_dir), + file_sha256=_file_sha256(tensor_path), + manifest_sha256=hashlib.sha256(manifest_payload).hexdigest(), + logical_bytes=logical_bytes, + transported_bytes=tensor_path.stat().st_size, + tensor_count=len(tensors), + rank=rank, + base_model=base_model, + tensors=tuple(tensors), + ) + + +def validate_replay_snapshots( + snapshots: list[AdapterSnapshot], + *, + expected_rank: int, + expected_experts: int, + expected_model_substring: str, +) -> None: + if not snapshots: + raise ValueError("At least one adapter snapshot is required") + reference = snapshots[0] + for snapshot in snapshots: + if snapshot.rank != expected_rank: + raise ValueError( + f"{snapshot.path} has rank {snapshot.rank}; expected {expected_rank}" + ) + if expected_model_substring not in snapshot.base_model: + raise ValueError( + f"{snapshot.path} targets {snapshot.base_model!r}; expected a model " + f"containing {expected_model_substring!r}" + ) + if snapshot.manifest_sha256 != reference.manifest_sha256: + raise ValueError( + f"{snapshot.path} does not have the reference serving tensor layout" + ) + expert_tensors = [ + tensor + for tensor in reference.tensors + if ".mlp.experts." in tensor.name or ".mlp.experts.base_layer." in tensor.name + ] + if not expert_tensors: + raise ValueError("Adapter has no packed Qwen MoE expert tensors") + mismatched = [ + tensor.name for tensor in expert_tensors if expected_experts not in tensor.shape + ] + if mismatched: + raise ValueError( + "Packed expert tensors do not preserve the expected expert dimension: " + + ", ".join(mismatched[:3]) + ) + + +def validate_committed_policy_version( + response_body: dict[str, Any], *, expected: int +) -> None: + committed = int(response_body.get("policy_version", -1)) + if committed != expected: + raise RuntimeError(f"Runtime committed policy {committed}; expected {expected}") + + +class FixedLoad: + def __init__( + self, + client: httpx.AsyncClient, + *, + request: dict[str, Any], + concurrency: int, + ) -> None: + self._client = client + self._request = request + self._concurrency = concurrency + self._running = False + self._tasks: list[asyncio.Task[None]] = [] + self.counters = LoadCounters() + + async def _worker(self) -> None: + while self._running: + try: + response = await self._client.post( + "/v1/chat/completions", json=self._request + ) + response.raise_for_status() + usage = response.json().get("usage", {}) + self.counters.prompt_tokens += int(usage.get("prompt_tokens", 0)) + self.counters.completion_tokens += int( + usage.get("completion_tokens", 0) + ) + self.counters.requests += 1 + except BaseException: + self.counters.errors += 1 + if not self._running: + return + + def snapshot(self) -> LoadCounters: + return LoadCounters(**asdict(self.counters)) + + async def start(self) -> None: + self._running = True + self._tasks = [ + asyncio.create_task(self._worker()) for _ in range(self._concurrency) + ] + + async def stop(self) -> None: + self._running = False + await asyncio.gather(*self._tasks, return_exceptions=True) + + +def _percentile(values: list[float], percentile: float) -> float: + if not values: + return 0.0 + ordered = sorted(values) + index = min(round((len(ordered) - 1) * percentile), len(ordered) - 1) + return ordered[index] + + +async def replay_updates(args: argparse.Namespace) -> dict[str, Any]: + snapshots = [inspect_adapter(path) for path in args.adapter_path] + validate_replay_snapshots( + snapshots, + expected_rank=args.expected_rank, + expected_experts=args.expected_experts, + expected_model_substring=args.expected_model_substring, + ) + server_paths = args.server_adapter_path or [item.path for item in snapshots] + if len(server_paths) != len(snapshots): + raise ValueError("--server-adapter-path must match --adapter-path count") + + headers = {"Authorization": f"Bearer {args.api_key}"} if args.api_key else {} + timeout = httpx.Timeout(args.timeout_s) + async with httpx.AsyncClient( + base_url=args.server_url.rstrip("/"), + headers=headers, + timeout=timeout, + ) as client: + metrics_response = await client.get("/art/metrics") + metrics_response.raise_for_status() + runtime_metrics = metrics_response.json() + runtime_world_size = int( + runtime_metrics.get("metrics", {}).get("world_size", 0) + ) + if runtime_world_size != args.expected_inference_world_size: + raise ValueError( + f"Runtime world size is {runtime_world_size}; expected " + f"{args.expected_inference_world_size}" + ) + + request: dict[str, Any] | None = None + if args.mode == "fixed-load": + if args.request_json is None: + raise ValueError("--request-json is required in fixed-load mode") + request = json.loads(Path(args.request_json).read_text()) + fixed_load = ( + FixedLoad(client, request=request, concurrency=args.load_concurrency) + if request is not None + else None + ) + if fixed_load is not None: + await fixed_load.start() + + update_records: list[dict[str, Any]] = [] + measured_update_s = 0.0 + try: + total_updates = args.warmups + args.updates + for update_index in range(total_updates): + snapshot_index = update_index % len(snapshots) + policy_version = args.first_policy_version + update_index + before_load = fixed_load.snapshot() if fixed_load else LoadCounters() + started = time.perf_counter() + response = await client.post( + "/art/in_flight_lora_update", + json={ + "model_name": args.model_name, + "base_model_name": args.base_model_name, + "lora_slot": args.lora_slot, + "lora_path": server_paths[snapshot_index], + "policy_version": policy_version, + }, + ) + response.raise_for_status() + wall_s = time.perf_counter() - started + body = response.json() + validate_committed_policy_version(body, expected=policy_version) + after_load = fixed_load.snapshot() if fixed_load else LoadCounters() + if update_index >= args.warmups: + measured_update_s += wall_s + update_records.append( + { + "update_index": update_index - args.warmups, + "policy_version": policy_version, + "snapshot_index": snapshot_index, + "snapshot_sha256": snapshots[snapshot_index].file_sha256, + "wall_s": wall_s, + "runtime_timing_s": body.get("timing_s", {}), + "load_during_update": asdict(after_load - before_load), + } + ) + + control_load = LoadCounters() + control_s = 0.0 + if fixed_load is not None and measured_update_s > 0: + before_control = fixed_load.snapshot() + control_started = time.perf_counter() + await asyncio.sleep(measured_update_s) + control_s = time.perf_counter() - control_started + control_load = fixed_load.snapshot() - before_control + finally: + if fixed_load is not None: + await fixed_load.stop() + + post_snapshots = [inspect_adapter(path) for path in args.adapter_path] + if [item.file_sha256 for item in snapshots] != [ + item.file_sha256 for item in post_snapshots + ]: + raise RuntimeError("Replay mutated a captured adapter snapshot") + + update_wall = [float(record["wall_s"]) for record in update_records] + update_completion_tokens = sum( + int(record["load_during_update"]["completion_tokens"]) + for record in update_records + ) + control_completion_rate = ( + control_load.completion_tokens / control_s if control_s > 0 else 0.0 + ) + lost_completion_tokens = max( + control_completion_rate * measured_update_s - update_completion_tokens, + 0.0, + ) + return { + "schema_version": 1, + "mode": args.mode, + "warmups": args.warmups, + "measured_updates": args.updates, + "runtime_world_size": runtime_world_size, + "snapshots": [ + {key: value for key, value in asdict(snapshot).items() if key != "tensors"} + for snapshot in snapshots + ], + "updates": update_records, + "control": { + "wall_s": control_s, + "load": asdict(control_load), + "completion_tok_per_s": control_completion_rate, + }, + "summary": { + "update_wall_mean_s": statistics.fmean(update_wall), + "update_wall_p50_s": _percentile(update_wall, 0.50), + "update_wall_p95_s": _percentile(update_wall, 0.95), + "measured_update_wall_s": measured_update_s, + "completion_tokens_during_updates": update_completion_tokens, + "lost_completion_tokens": lost_completion_tokens, + "lost_completion_tokens_per_update": lost_completion_tokens + / max(args.updates, 1), + }, + } + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Replay captured Qwen3.6 LoRA updates against ART vLLM" + ) + parser.add_argument("--server-url", required=True) + parser.add_argument("--adapter-path", action="append", required=True) + parser.add_argument("--server-adapter-path", action="append") + parser.add_argument("--model-name", required=True) + parser.add_argument("--base-model-name", required=True) + parser.add_argument("--lora-slot", required=True) + parser.add_argument("--api-key") + parser.add_argument("--mode", choices=("idle", "fixed-load"), default="idle") + parser.add_argument("--request-json") + parser.add_argument("--load-concurrency", type=int, default=8) + parser.add_argument("--warmups", type=int, default=5) + parser.add_argument("--updates", type=int, default=30) + parser.add_argument("--first-policy-version", type=int, default=1) + parser.add_argument("--timeout-s", type=float, default=120.0) + parser.add_argument("--expected-rank", type=int, default=8) + parser.add_argument("--expected-experts", type=int, default=256) + parser.add_argument("--expected-model-substring", default="Qwen3.6-35B-A3B") + parser.add_argument("--expected-inference-world-size", type=int, default=2) + parser.add_argument("--output", type=Path, required=True) + return parser + + +def main(argv: list[str] | None = None) -> None: + args = build_parser().parse_args(argv) + if args.warmups < 0 or args.updates < 1: + raise SystemExit( + "--warmups must be non-negative and --updates must be positive" + ) + result = asyncio.run(replay_updates(args)) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n") + print(json.dumps(result["summary"], sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tests/unit/test_weight_update_replay.py b/tests/unit/test_weight_update_replay.py new file mode 100644 index 000000000..f37a62ff3 --- /dev/null +++ b/tests/unit/test_weight_update_replay.py @@ -0,0 +1,131 @@ +import json +from pathlib import Path + +import pytest +import torch + +from art.megatron.model_support.lora_disk import save_vllm_lora_tensors +from art.megatron.service import MegatronService +from art.megatron.weights.lora_publish import LoraPublishMetrics +from art.megatron.weights.update_replay import ( + inspect_adapter, + validate_committed_policy_version, + validate_replay_snapshots, +) + + +def _write_snapshot( + path: Path, + *, + rank: int = 8, + experts: int = 256, + value: float = 1.0, +) -> None: + save_vllm_lora_tensors( + path, + { + ( + "base_model.model.model.language_model.layers.0." + "mlp.experts.lora_A.weight" + ): torch.full((experts, rank, 2), value, dtype=torch.bfloat16), + ( + "base_model.model.model.language_model.layers.0." + "mlp.experts.lora_B.weight" + ): torch.full((experts, 4, rank), value, dtype=torch.bfloat16), + }, + { + "base_model_name_or_path": "Qwen/Qwen3.6-35B-A3B", + "r": rank, + "lora_alpha": rank, + }, + ) + + +def test_publish_metrics_use_stable_training_metric_names() -> None: + metrics = LoraPublishMetrics( + gather_pack_s=1.0, + stage_to_cpu_s=2.0, + write_s=3.0, + total_s=6.0, + logical_bytes=10, + transported_bytes=12, + tensor_count=4, + ) + + assert metrics.as_training_metrics() == { + "time/weight_update_trainer_gather_pack_s": 1.0, + "time/weight_update_trainer_stage_to_cpu_s": 2.0, + "time/weight_update_trainer_write_s": 3.0, + "time/weight_update_trainer_publish_s": 6.0, + "weight_update/logical_bytes": 10, + "weight_update/transported_bytes": 12, + "weight_update/tensor_count": 4, + } + + +def test_weight_update_event_metrics_exclude_control_fields() -> None: + assert MegatronService._weight_update_metrics_from_event( + { + "event": "lora_ready", + "step": 7, + "time/weight_update_trainer_publish_s": 5.0, + "weight_update/logical_bytes": 1024, + "diagnostic": "ignored", + } + ) == { + "time/weight_update_trainer_publish_s": 5.0, + "weight_update/logical_bytes": 1024.0, + } + + +def test_replay_snapshot_preserves_layout_bytes_and_integrity(tmp_path: Path) -> None: + first_path = tmp_path / "first" + second_path = tmp_path / "second" + _write_snapshot(first_path, value=1.0) + _write_snapshot(second_path, value=2.0) + + first = inspect_adapter(first_path) + second = inspect_adapter(second_path) + validate_replay_snapshots( + [first, second], + expected_rank=8, + expected_experts=256, + expected_model_substring="Qwen3.6-35B-A3B", + ) + + assert first.manifest_sha256 == second.manifest_sha256 + assert first.file_sha256 != second.file_sha256 + assert first.logical_bytes == 256 * 8 * (2 + 4) * 2 + assert ( + first.transported_bytes + == (first_path / "adapter_model.safetensors").stat().st_size + ) + assert first.tensor_count == 2 + + +def test_replay_snapshot_rejects_nonrepresentative_rank(tmp_path: Path) -> None: + path = tmp_path / "rank-one" + _write_snapshot(path, rank=1) + + with pytest.raises(ValueError, match="has rank 1; expected 8"): + validate_replay_snapshots( + [inspect_adapter(path)], + expected_rank=8, + expected_experts=256, + expected_model_substring="Qwen3.6-35B-A3B", + ) + + +def test_policy_version_must_match_committed_runtime_version() -> None: + validate_committed_policy_version({"policy_version": 12}, expected=12) + + with pytest.raises(RuntimeError, match="committed policy 11; expected 12"): + validate_committed_policy_version({"policy_version": 11}, expected=12) + + +def test_snapshot_config_is_machine_readable(tmp_path: Path) -> None: + path = tmp_path / "snapshot" + _write_snapshot(path) + + config = json.loads((path / "adapter_config.json").read_text()) + assert config["art_lora_format"] == "vllm" From 393b6865b878e2be9c06b74e33346f7d948d77b4 Mon Sep 17 00:00:00 2001 From: Vivek Kalyan Date: Thu, 23 Jul 2026 19:16:56 -0700 Subject: [PATCH 3/9] feat: Replay updates through the Megatron publisher --- src/art/megatron/runtime/jobs.py | 11 +++++ src/art/megatron/service.py | 49 +++++++++++++++++++++++ src/art/megatron/train.py | 48 +++++++++++++++++++++- src/art/megatron/weights/update_replay.py | 23 ++++++++--- tests/unit/test_weight_update_replay.py | 33 +++++++++++++++ 5 files changed, 157 insertions(+), 7 deletions(-) diff --git a/src/art/megatron/runtime/jobs.py b/src/art/megatron/runtime/jobs.py index 4285c88bc..022be1a72 100644 --- a/src/art/megatron/runtime/jobs.py +++ b/src/art/megatron/runtime/jobs.py @@ -59,6 +59,16 @@ class MegatronSyncJob(BaseModel): log_path: str = DEFAULT_TRAINING_LOG_PATH +class MegatronLoraPublishJob(BaseModel): + kind: Literal["publish_lora"] = "publish_lora" + step: int = Field(ge=0) + source_lora_path: str + output_lora_path: str + content_version: int = Field(ge=0) + allow_unvalidated_arch: bool = False + log_path: str = DEFAULT_TRAINING_LOG_PATH + + class MegatronOptimizerSaveJob(BaseModel): kind: Literal["save_optimizer"] = "save_optimizer" step: int = Field(ge=0) @@ -89,6 +99,7 @@ class MegatronSFTTrainingJob(BaseModel): MegatronTrainingJob | MegatronMergedTrainingJob | MegatronSyncJob + | MegatronLoraPublishJob | MegatronOptimizerSaveJob | MegatronSFTTrainingJob, Field(discriminator="kind"), diff --git a/src/art/megatron/service.py b/src/art/megatron/service.py index 49f8f1f58..ed8dd0ae7 100644 --- a/src/art/megatron/service.py +++ b/src/art/megatron/service.py @@ -73,6 +73,7 @@ from .runtime.jobs import ( LORA_READY_EVENT, OPTIMIZER_READY_EVENT, + MegatronLoraPublishJob, MegatronMergedTrainingJob, MegatronOptimizerSaveJob, MegatronSFTTrainingJob, @@ -935,6 +936,54 @@ async def _sync_dedicated_merged_weights( pass self._latest_step = step + async def replay_lora_update( + self, + *, + source_lora_path: str, + output_lora_path: str, + policy_version: int, + content_version: int, + ) -> dict[str, float]: + """Publish real resident Megatron tensors, then install that exact output.""" + if not self.is_dedicated: + raise RuntimeError("LoRA update replay requires dedicated trainer GPUs") + if self.rollout_weight_update_mode != "in_flight_lora": + raise RuntimeError("LoRA update replay requires in_flight_lora mode") + if os.path.exists(output_lora_path): + raise FileExistsError(f"Replay output already exists: {output_lora_path}") + await self._ensure_megatron_running() + self._clear_pending_jobs() + job_path, log_path = self._create_megatron_job_paths() + job = MegatronLoraPublishJob( + step=policy_version, + source_lora_path=source_lora_path, + output_lora_path=output_lora_path, + content_version=content_version, + allow_unvalidated_arch=self._allow_unvalidated_arch, + log_path=log_path, + ) + write_megatron_job(job, job_path=job_path) + publish_metrics: dict[str, float] | None = None + async for result in stream_megatron_job( + job, + job_path=job_path, + process=self._megatron_process, + process_log_path=self._megatron_log_path, + ): + if result.get("event") != LORA_READY_EVENT: + raise RuntimeError(f"Unexpected replay worker event: {result!r}") + if int(result.get("step", -1)) != policy_version: + raise RuntimeError(f"Replay worker published wrong step: {result!r}") + if int(result.get("content_version", -1)) != content_version: + raise RuntimeError(f"Replay worker published wrong content: {result!r}") + publish_metrics = self._weight_update_metrics_from_event(result) + if publish_metrics is None: + raise RuntimeError("Replay worker produced no LoRA-ready metrics") + install_metrics = await self._update_in_flight_adapter( + output_lora_path, policy_version + ) + return {**publish_metrics, **install_metrics} + async def _sleep_runtime(self) -> None: import httpx diff --git a/src/art/megatron/train.py b/src/art/megatron/train.py index 6868eee33..7bbfdf57b 100644 --- a/src/art/megatron/train.py +++ b/src/art/megatron/train.py @@ -70,6 +70,7 @@ LORA_READY_EVENT, OPTIMIZER_READY_EVENT, MegatronJob, + MegatronLoraPublishJob, MegatronMergedTrainingJob, MegatronOptimizerSaveJob, MegatronSFTTrainingJob, @@ -889,6 +890,49 @@ def _run_megatron_job(runtime: TrainingRuntime, job: MegatronJob) -> None: pause_generation=False, ) return + if isinstance(job, MegatronLoraPublishJob): + adapter_model = _load_adapter_into_model( + runtime.model, + job.source_lora_path, + runtime.rank, + handler=runtime.model_support_handler, + ) + adapter_dtypes = {key: tensor.dtype for key, tensor in adapter_model.items()} + del adapter_model + from art.megatron.weights.lora_publish import iter_lora_modules + + content_value = float(job.content_version + 1) / 1024.0 + with torch.no_grad(): + for module in iter_lora_modules(runtime.model): + module.B_T.view(-1)[0] = content_value + publish_metrics = save_vllm_lora_from_model( + model=runtime.model, + adapter_dtypes=adapter_dtypes, + handler=runtime.model_support_handler, + adapter_config=load_adapter_config(job.source_lora_path), + output_dir=job.output_lora_path, + rank=runtime.rank, + world_size=runtime.world_size, + ) + if runtime.rank == 0: + assert publish_metrics is not None + _write_job_event( + job.log_path, + LORA_READY_EVENT, + step=job.step, + content_version=job.content_version, + **{ + "topology/trainer_tp": ps.get_tensor_model_parallel_world_size(), + "topology/trainer_pp": ps.get_pipeline_model_parallel_world_size(), + "topology/trainer_cp": ps.get_context_parallel_world_size(), + "topology/trainer_ep": ps.get_expert_model_parallel_world_size(), + "topology/trainer_etp": ( + ps.get_expert_tensor_parallel_world_size() + ), + }, + **publish_metrics.as_training_metrics(), + ) + return if isinstance(job, MegatronSFTTrainingJob): run_megatron_sft_job(runtime, job) return @@ -903,7 +947,9 @@ def _run_megatron_job(runtime: TrainingRuntime, job: MegatronJob) -> None: def _job_cleanup_path(job: MegatronJob) -> str | None: - if isinstance(job, (MegatronOptimizerSaveJob, MegatronSyncJob)): + if isinstance( + job, (MegatronOptimizerSaveJob, MegatronSyncJob, MegatronLoraPublishJob) + ): return None if isinstance(job, MegatronSFTTrainingJob): return job.sft_data_dir diff --git a/src/art/megatron/weights/update_replay.py b/src/art/megatron/weights/update_replay.py index cfbcd5a55..cbae0bed2 100644 --- a/src/art/megatron/weights/update_replay.py +++ b/src/art/megatron/weights/update_replay.py @@ -1,8 +1,7 @@ -"""Replay in-flight LoRA updates against an ART vLLM runtime. +"""Receiver-only replay of pre-published LoRA updates. -The replay input is a sequence of real trainer-produced checkpoints. Keeping the -captured safetensors files intact preserves the serving tensor names, shapes, -dtypes, and byte volume without rerunning Megatron, reward, or judging. +This diagnostic deliberately excludes trainer publication. Use +``MegatronService.replay_lora_update`` for end-to-end measurements. """ from __future__ import annotations @@ -116,6 +115,12 @@ def validate_replay_snapshots( ) -> None: if not snapshots: raise ValueError("At least one adapter snapshot is required") + paths = [snapshot.path for snapshot in snapshots] + hashes = [snapshot.file_sha256 for snapshot in snapshots] + if len(set(paths)) != len(paths): + raise ValueError("Receiver-only replay requires unique snapshot paths") + if len(set(hashes)) != len(hashes): + raise ValueError("Receiver-only replay requires unique snapshot contents") reference = snapshots[0] for snapshot in snapshots: if snapshot.rank != expected_rank: @@ -222,6 +227,12 @@ async def replay_updates(args: argparse.Namespace) -> dict[str, Any]: server_paths = args.server_adapter_path or [item.path for item in snapshots] if len(server_paths) != len(snapshots): raise ValueError("--server-adapter-path must match --adapter-path count") + total_updates = args.warmups + args.updates + if len(snapshots) != total_updates: + raise ValueError( + "Receiver-only replay requires exactly one unique snapshot per update: " + f"{len(snapshots)} != {total_updates}" + ) headers = {"Authorization": f"Bearer {args.api_key}"} if args.api_key else {} timeout = httpx.Timeout(args.timeout_s) @@ -258,9 +269,8 @@ async def replay_updates(args: argparse.Namespace) -> dict[str, Any]: update_records: list[dict[str, Any]] = [] measured_update_s = 0.0 try: - total_updates = args.warmups + args.updates for update_index in range(total_updates): - snapshot_index = update_index % len(snapshots) + snapshot_index = update_index policy_version = args.first_policy_version + update_index before_load = fixed_load.snapshot() if fixed_load else LoadCounters() started = time.perf_counter() @@ -325,6 +335,7 @@ async def replay_updates(args: argparse.Namespace) -> dict[str, Any]: ) return { "schema_version": 1, + "measurement_scope": "receiver_only_pre_published_snapshots", "mode": args.mode, "warmups": args.warmups, "measured_updates": args.updates, diff --git a/tests/unit/test_weight_update_replay.py b/tests/unit/test_weight_update_replay.py index f37a62ff3..ad564a133 100644 --- a/tests/unit/test_weight_update_replay.py +++ b/tests/unit/test_weight_update_replay.py @@ -5,6 +5,11 @@ import torch from art.megatron.model_support.lora_disk import save_vllm_lora_tensors +from art.megatron.runtime.jobs import ( + MegatronLoraPublishJob, + dump_megatron_job, + load_megatron_job, +) from art.megatron.service import MegatronService from art.megatron.weights.lora_publish import LoraPublishMetrics from art.megatron.weights.update_replay import ( @@ -116,6 +121,34 @@ def test_replay_snapshot_rejects_nonrepresentative_rank(tmp_path: Path) -> None: ) +def test_receiver_replay_rejects_duplicate_snapshot_contents(tmp_path: Path) -> None: + first_path = tmp_path / "first" + second_path = tmp_path / "second" + _write_snapshot(first_path, value=1.0) + _write_snapshot(second_path, value=1.0) + + with pytest.raises(ValueError, match="unique snapshot contents"): + validate_replay_snapshots( + [inspect_adapter(first_path), inspect_adapter(second_path)], + expected_rank=8, + expected_experts=256, + expected_model_substring="Qwen3.6-35B-A3B", + ) + + +def test_lora_publish_job_roundtrips_through_worker_protocol() -> None: + job = MegatronLoraPublishJob( + step=9, + source_lora_path="/checkpoints/seed", + output_lora_path="/scratch/replay/0009", + content_version=8, + allow_unvalidated_arch=True, + log_path="/scratch/replay/0009.jsonl", + ) + + assert load_megatron_job(dump_megatron_job(job)) == job + + def test_policy_version_must_match_committed_runtime_version() -> None: validate_committed_policy_version({"policy_version": 12}, expected=12) From e947f5a6458e5bc40689a813a4bb133d2b1560f0 Mon Sep 17 00:00:00 2001 From: Vivek Kalyan Date: Thu, 23 Jul 2026 19:17:08 -0700 Subject: [PATCH 4/9] fix: Reject stale in-flight policy updates --- .../test_runtime_project_isolation.py | 7 +++++++ .../src/art_vllm_runtime/dedicated_server.py | 2 +- .../src/art_vllm_runtime/policy_spans.py | 21 ++++++++++++++++++- 3 files changed, 28 insertions(+), 2 deletions(-) diff --git a/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py b/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py index eb7501217..6ffc9a7c5 100644 --- a/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py +++ b/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py @@ -173,6 +173,11 @@ async def admit(): blocked = not admission.done() await coordinator.commit_update(slot, 5, new) version, admitted_lora = await admission + stale_rejected = False + try: + await coordinator.begin_update(slot, 5) + except ValueError: + stale_rejected = True async with coordinator.admission(slot): cancelled_update = asyncio.create_task(coordinator.begin_update(slot)) @@ -191,6 +196,7 @@ async def admit(): "policy_version": version, "lora_path": admitted_lora.lora_path, "recovered_after_cancel": recovered, + "stale_rejected": stale_rejected, }, sort_keys=True)) asyncio.run(main()) @@ -204,6 +210,7 @@ async def admit(): "lora_path": "new", "policy_version": 5, "recovered_after_cancel": True, + "stale_rejected": True, } diff --git a/vllm_runtime/src/art_vllm_runtime/dedicated_server.py b/vllm_runtime/src/art_vllm_runtime/dedicated_server.py index 96c29c2ef..1cf16b943 100644 --- a/vllm_runtime/src/art_vllm_runtime/dedicated_server.py +++ b/vllm_runtime/src/art_vllm_runtime/dedicated_server.py @@ -230,7 +230,7 @@ async def in_flight_lora_update( coordinator = lora_update_coordinator(models, engine_client) update_started = time.perf_counter() begin_started = update_started - await coordinator.begin_update(lora_slot) + await coordinator.begin_update(lora_slot, policy_version) begin_s = time.perf_counter() - begin_started try: load_started = time.perf_counter() diff --git a/vllm_runtime/src/art_vllm_runtime/policy_spans.py b/vllm_runtime/src/art_vllm_runtime/policy_spans.py index da90c7192..4a1459018 100644 --- a/vllm_runtime/src/art_vllm_runtime/policy_spans.py +++ b/vllm_runtime/src/art_vllm_runtime/policy_spans.py @@ -101,12 +101,23 @@ async def admission( state.active_admissions -= 1 state.condition.notify_all() - async def begin_update(self, lora_slot: str) -> None: + async def begin_update( + self, lora_slot: str, policy_version: int | None = None + ) -> None: state = self._state(lora_slot) async with state.condition: acquired = False try: await state.condition.wait_for(lambda: not state.update_active) + if ( + policy_version is not None + and state.policy_version is not None + and int(policy_version) <= state.policy_version + ): + raise ValueError( + f"LoRA slot {lora_slot!r} policy version must increase: " + f"{policy_version} <= {state.policy_version}" + ) state.update_active = True state.blocked = True acquired = True @@ -128,6 +139,14 @@ async def commit_update( async with state.condition: if not state.update_active: raise RuntimeError(f"No active LoRA update for slot {lora_slot!r}") + if ( + state.policy_version is not None + and int(policy_version) <= state.policy_version + ): + raise ValueError( + f"LoRA slot {lora_slot!r} policy version must increase: " + f"{policy_version} <= {state.policy_version}" + ) state.policy_version = int(policy_version) state.lora_request = lora_request state.update_active = False From 1fc2cdbbe7b60965a7e47cd198e39e2d65ad26af Mon Sep 17 00:00:00 2001 From: Vivek Kalyan Date: Thu, 23 Jul 2026 19:17:08 -0700 Subject: [PATCH 5/9] ops: Add four-H200 LoRA replay launcher --- .../lora-update-replay-h200.sky.yaml | 60 ++++ .../benchmarks/prepare_lora_update_replay.py | 83 +++++ .../benchmarks/service_lora_update_replay.py | 284 ++++++++++++++++++ 3 files changed, 427 insertions(+) create mode 100644 scripts/benchmarks/lora-update-replay-h200.sky.yaml create mode 100644 scripts/benchmarks/prepare_lora_update_replay.py create mode 100644 scripts/benchmarks/service_lora_update_replay.py diff --git a/scripts/benchmarks/lora-update-replay-h200.sky.yaml b/scripts/benchmarks/lora-update-replay-h200.sky.yaml new file mode 100644 index 000000000..3cbc372ed --- /dev/null +++ b/scripts/benchmarks/lora-update-replay-h200.sky.yaml @@ -0,0 +1,60 @@ +name: art-lora-update-replay-h200 + +workdir: . + +resources: + cloud: kubernetes + accelerators: H200:4 + cpus: "32+" + memory: "256+" + +envs: + CUDA_HOME: /usr/local/cuda-12.8 + TORCH_CUDA_ARCH_LIST: "9.0" + UV_CONCURRENT_BUILDS: "1" + UV_LINK_MODE: copy + WANDB_ENTITY: wb-training + WANDB_PROJECT: bench + +setup: | + set -euo pipefail + curl -LsSf https://astral.sh/uv/install.sh | sh + export PATH="$HOME/.local/bin:$PATH" + uv sync --extra megatron + uv run --no-sync bash src/art/megatron/setup.sh + cd vllm_runtime + uv sync + +run: | + set -euo pipefail + export PATH="$HOME/.local/bin:$PATH" + mkdir -p scratch/lora-update-replay + uv run --no-sync python scripts/benchmarks/prepare_lora_update_replay.py \ + --output scratch/lora-update-replay/inputs + + request_jsonl="$PWD/scratch/lora-update-replay/inputs/requests.jsonl" + uv run --no-sync python scripts/benchmarks/service_lora_update_replay.py \ + --request-jsonl "$request_jsonl" \ + --mode fixed-load \ + --warmups 5 \ + --updates 30 \ + --output scratch/lora-update-replay/fixed-load + + uv run --no-sync python scripts/benchmarks/service_lora_update_replay.py \ + --request-jsonl "$request_jsonl" \ + --mode idle \ + --warmups 5 \ + --updates 30 \ + --output scratch/lora-update-replay/idle + +config: + kubernetes: + pod_config: + spec: + schedulerName: binpack-scheduler + activeDeadlineSeconds: 21600 + containers: + - name: ray-node + envFrom: + - secretRef: + name: vivek-dev-api-keys diff --git a/scripts/benchmarks/prepare_lora_update_replay.py b/scripts/benchmarks/prepare_lora_update_replay.py new file mode 100644 index 000000000..049f91b7c --- /dev/null +++ b/scripts/benchmarks/prepare_lora_update_replay.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 + +import argparse +import hashlib +import json +from pathlib import Path + +BONNIE_ARTIFACT = "wb-training/bench/bonnie:v1" +BONNIE_TRAIN_SHA256 = "fd26c97423de94a158d5a54e38485ec55b3fd3504f5f0346850fcfa79ad3ce66" + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + artifact_dir = args.output / "bonnie-v1" + artifact_dir.mkdir(parents=True, exist_ok=True) + + import wandb + + wandb.Api().artifact(BONNIE_ARTIFACT).download(root=str(artifact_dir)) + train_path = artifact_dir / "train.jsonl" + digest = _sha256(train_path) + if digest != BONNIE_TRAIN_SHA256: + raise RuntimeError( + f"Bonnie train artifact checksum is {digest}; " + f"expected {BONNIE_TRAIN_SHA256}" + ) + request_path = args.output / "requests.jsonl" + requests: list[dict[str, object]] = [] + with train_path.open(encoding="utf-8") as handle: + for source_index, line in enumerate(handle): + if not line.strip(): + continue + row = json.loads(line) + messages = [dict(message) for message in row["messages"]] + for index, message in enumerate(messages): + if index > 0 and message.get("role") == "system": + message["role"] = "user" + message["content"] = "[System note]\n" + str( + message.get("content") or "" + ) + request: dict[str, object] = { + "messages": messages, + "n": 8, + "temperature": 0.8, + "max_tokens": 256, + "logprobs": True, + "chat_template_kwargs": {"enable_thinking": False}, + } + if row.get("tools"): + request["tools"] = row["tools"] + requests.append( + { + "id": str( + row.get("id") + or ( + f"mar31_v2:train:{row.get('trace_id', source_index)}:" + f"assistant-{row.get('assistant_turn_index', 0)}" + ) + ), + "request": request, + } + ) + if len(requests) == 64: + break + request_path.write_text( + "".join(json.dumps(row, sort_keys=True) + "\n" for row in requests) + ) + print(f"BONNIE_TRAIN_JSONL={train_path}") + print(f"REQUEST_JSONL={request_path}") + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmarks/service_lora_update_replay.py b/scripts/benchmarks/service_lora_update_replay.py new file mode 100644 index 000000000..c67a6f104 --- /dev/null +++ b/scripts/benchmarks/service_lora_update_replay.py @@ -0,0 +1,284 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import argparse +import asyncio +import hashlib +import json +from pathlib import Path +import random +import statistics +import subprocess +import time +from typing import Any + +import httpx + +import art +from art.megatron.service import MegatronService +from art.megatron.weights.update_replay import inspect_adapter + +BASE_MODEL = "Qwen/Qwen3.6-35B-A3B" + + +def _hash_request(request: dict[str, Any]) -> str: + encoded = json.dumps(request, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(encoded).hexdigest() + + +def _load_requests(path: Path) -> list[dict[str, Any]]: + rows = [json.loads(line) for line in path.read_text().splitlines() if line] + if len(rows) != 64: + raise ValueError(f"Expected 64 pinned request rows, got {len(rows)}") + ids: set[str] = set() + hashes: set[str] = set() + for row in rows: + request = row["request"] + request["model"] = "bonnie-replay:active" + expected = { + "n": 8, + "temperature": 0.8, + "max_tokens": 256, + "logprobs": True, + } + for key, value in expected.items(): + if request.get(key) != value: + raise ValueError( + f"Request {row.get('id')} has {key}={request.get(key)}" + ) + if request.get("chat_template_kwargs") != {"enable_thinking": False}: + raise ValueError(f"Request {row.get('id')} enables thinking") + row["request_hash"] = _hash_request(request) + ids.add(str(row["id"])) + hashes.add(row["request_hash"]) + if len(ids) != 64 or len(hashes) != 64: + raise ValueError("Pinned request trace contains duplicate IDs or requests") + return rows + + +async def _load_worker( + *, + client: httpx.AsyncClient, + requests: list[dict[str, Any]], + worker: int, + stop: asyncio.Event, + samples: list[dict[str, Any]], +) -> None: + index = worker + while not stop.is_set(): + row = requests[index % len(requests)] + index += 8 + started = time.perf_counter() + error = None + prompt_tokens = completion_tokens = 0 + try: + response = await client.post("/v1/chat/completions", json=row["request"]) + response.raise_for_status() + usage = response.json().get("usage", {}) + prompt_tokens = int(usage.get("prompt_tokens", 0)) + completion_tokens = int(usage.get("completion_tokens", 0)) + except Exception as exc: + error = f"{type(exc).__name__}: {exc}" + samples.append( + { + "worker": worker, + "id": row["id"], + "request_hash": row["request_hash"], + "latency_s": time.perf_counter() - started, + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "error": error, + } + ) + + +def _bootstrap_ci(values: list[float]) -> list[float]: + rng = random.Random(20260723) + means = [statistics.fmean(rng.choices(values, k=len(values))) for _ in range(2000)] + means.sort() + return [means[50], means[1949]] + + +def _gpu_memory_snapshot() -> list[dict[str, int]]: + output = subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=index,memory.used,memory.total", + "--format=csv,noheader,nounits", + ], + text=True, + ) + rows = [] + for line in output.splitlines(): + index, used, total = (int(part.strip()) for part in line.split(",")) + rows.append({"gpu_index": index, "used_mib": used, "total_mib": total}) + return rows + + +async def run(args: argparse.Namespace) -> None: + requests = _load_requests(args.request_jsonl) + output = args.output.resolve() + output.mkdir(parents=True, exist_ok=True) + art.init_megatron_runtime_config( + topology=art.MegatronTopologyConfig(tp=1, pp=1, cp=2, ep=2, etp=1), + packed_sequence_length=122880, + streaming_weight_offload=True, + ) + service = MegatronService( + model_name="bonnie-replay", + base_model=BASE_MODEL, + config={ + "trainer_gpu_ids": [0, 1], + "inference_gpu_ids": [2, 3], + "rollout_weights_mode": "lora", + "rollout_weight_update_mode": "in_flight_lora", + "allow_unvalidated_arch": True, + "chat_template_kwargs": {"enable_thinking": False}, + "lora_config": {"rank": 8}, + "engine_args": { + "tensor_parallel_size": 2, + "max_model_len": 16384, + "max_num_batched_tokens": 131072, + "max_num_seqs": 256, + "enable_prefix_caching": True, + "gpu_memory_utilization": 0.9, + }, + }, + output_dir=str(output / "service"), + ) + seed = Path(service._resolve_active_lora_path()) + await service.start_openai_server(None) + base_url = service._vllm_base_url + headers = service._runtime_headers() + if service._in_flight_lora_slot != "bonnie-replay:active": + raise RuntimeError("Replay slot name changed unexpectedly") + samples: list[dict[str, Any]] = [] + stop = asyncio.Event() + async with httpx.AsyncClient( + base_url=base_url, headers=headers, timeout=300.0 + ) as client: + runtime_before = (await client.get("/art/metrics")).json() + load_tasks = ( + [ + asyncio.create_task( + _load_worker( + client=client, + requests=requests, + worker=worker, + stop=stop, + samples=samples, + ) + ) + for worker in range(8) + ] + if args.mode == "fixed-load" + else [] + ) + updates: list[dict[str, Any]] = [] + hashes: set[str] = set() + try: + for index in range(args.warmups + args.updates): + policy_version = index + 1 + checkpoint = output / "checkpoints" / f"{policy_version:04d}" + sample_start = len(samples) + memory_before = _gpu_memory_snapshot() + metrics = await service.replay_lora_update( + source_lora_path=str(seed), + output_lora_path=str(checkpoint), + policy_version=policy_version, + content_version=index, + ) + snapshot = inspect_adapter(checkpoint) + if snapshot.file_sha256 in hashes: + raise RuntimeError("Replay produced duplicate update contents") + hashes.add(snapshot.file_sha256) + updates.append( + { + "index": index, + "measured": index >= args.warmups, + "policy_version": policy_version, + "checkpoint": str(checkpoint), + "sha256": snapshot.file_sha256, + "metrics": metrics, + "gpu_memory_before": memory_before, + "gpu_memory_after": _gpu_memory_snapshot(), + "sample_start": sample_start, + "sample_end": len(samples), + } + ) + finally: + stop.set() + await asyncio.gather(*load_tasks) + runtime_after = (await client.get("/art/metrics")).json() + await service.aclose() + + measured = [row for row in updates if row["measured"]] + totals = [ + float(row["metrics"]["time/weight_update_trainer_publish_s"]) + + float(row["metrics"]["time/weight_update_service_rpc_s"]) + for row in measured + ] + (output / "updates.jsonl").write_text( + "".join(json.dumps(row, sort_keys=True) + "\n" for row in updates) + ) + (output / "samples.jsonl").write_text( + "".join(json.dumps(row, sort_keys=True) + "\n" for row in samples) + ) + (output / "runtime-before.json").write_text(json.dumps(runtime_before, indent=2)) + (output / "runtime-after.json").write_text(json.dumps(runtime_after, indent=2)) + topology = { + "trainer_gpu_ids": [0, 1], + "inference_gpu_ids": [2, 3], + "inference_tp": 2, + "trainer": {"tp": 1, "pp": 1, "cp": 2, "ep": 2, "etp": 1}, + } + (output / "topology.json").write_text(json.dumps(topology, indent=2)) + summary = { + "update_total_mean_s": statistics.fmean(totals), + "update_total_p50_s": statistics.median(totals), + "update_total_p95_s": sorted(totals)[round(0.95 * (len(totals) - 1))], + "update_total_mean_95ci_s": _bootstrap_ci(totals), + "prompt_tokens": sum(row["prompt_tokens"] for row in samples), + "completion_tokens": sum(row["completion_tokens"] for row in samples), + "request_latency_mean_s": ( + statistics.fmean(row["latency_s"] for row in samples) if samples else None + ), + "request_errors": sum(row["error"] is not None for row in samples), + } + (output / "summary.json").write_text(json.dumps(summary, indent=2)) + manifest = { + "schema_version": 1, + "mode": args.mode, + "request_trace": str(args.request_jsonl), + "request_ids_and_hashes": [ + {"id": row["id"], "request_hash": row["request_hash"]} for row in requests + ], + "routed_expert_collection": { + "supported": False, + "reason": "The replay uses vLLM's OpenAI JSON endpoint.", + }, + "artifacts": [ + "updates.jsonl", + "samples.jsonl", + "runtime-before.json", + "runtime-after.json", + "topology.json", + "summary.json", + ], + } + (output / "manifest.json").write_text(json.dumps(manifest, indent=2)) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--request-jsonl", type=Path, required=True) + parser.add_argument("--mode", choices=("idle", "fixed-load"), required=True) + parser.add_argument("--warmups", type=int, default=5) + parser.add_argument("--updates", type=int, default=30) + parser.add_argument("--output", type=Path, required=True) + asyncio.run(run(parser.parse_args())) + + +if __name__ == "__main__": + main() From 2617183812c1aaef1eab59fac1cd20e210a300ed Mon Sep 17 00:00:00 2001 From: Vivek Kalyan Date: Thu, 23 Jul 2026 19:19:08 -0700 Subject: [PATCH 6/9] fix: Validate replay GPU topology --- .../benchmarks/service_lora_update_replay.py | 38 ++++++++++++++++--- 1 file changed, 33 insertions(+), 5 deletions(-) diff --git a/scripts/benchmarks/service_lora_update_replay.py b/scripts/benchmarks/service_lora_update_replay.py index c67a6f104..75ee61749 100644 --- a/scripts/benchmarks/service_lora_update_replay.py +++ b/scripts/benchmarks/service_lora_update_replay.py @@ -100,19 +100,28 @@ def _bootstrap_ci(values: list[float]) -> list[float]: return [means[50], means[1949]] -def _gpu_memory_snapshot() -> list[dict[str, int]]: +def _gpu_memory_snapshot() -> list[dict[str, int | str]]: output = subprocess.check_output( [ "nvidia-smi", - "--query-gpu=index,memory.used,memory.total", + "--query-gpu=index,name,memory.used,memory.total", "--format=csv,noheader,nounits", ], text=True, ) - rows = [] + rows: list[dict[str, int | str]] = [] for line in output.splitlines(): - index, used, total = (int(part.strip()) for part in line.split(",")) - rows.append({"gpu_index": index, "used_mib": used, "total_mib": total}) + raw_index, name, raw_used, raw_total = ( + part.strip() for part in line.split(",", maxsplit=3) + ) + rows.append( + { + "gpu_index": int(raw_index), + "name": name, + "used_mib": int(raw_used), + "total_mib": int(raw_total), + } + ) return rows @@ -159,6 +168,8 @@ async def run(args: argparse.Namespace) -> None: base_url=base_url, headers=headers, timeout=300.0 ) as client: runtime_before = (await client.get("/art/metrics")).json() + if int(runtime_before.get("metrics", {}).get("world_size", 0)) != 2: + raise RuntimeError(f"Expected vLLM TP2 runtime: {runtime_before!r}") load_tasks = ( [ asyncio.create_task( @@ -189,6 +200,23 @@ async def run(args: argparse.Namespace) -> None: policy_version=policy_version, content_version=index, ) + expected_topology = { + "topology/trainer_tp": 1.0, + "topology/trainer_pp": 1.0, + "topology/trainer_cp": 2.0, + "topology/trainer_ep": 2.0, + "topology/trainer_etp": 1.0, + } + if { + key: metrics.get(key) for key in expected_topology + } != expected_topology: + raise RuntimeError(f"Unexpected trainer topology: {metrics!r}") + if len(memory_before) != 4: + raise RuntimeError( + f"Expected four visible H200s, got {memory_before!r}" + ) + if any("H200" not in str(row["name"]) for row in memory_before): + raise RuntimeError(f"Replay requires H200 GPUs: {memory_before!r}") snapshot = inspect_adapter(checkpoint) if snapshot.file_sha256 in hashes: raise RuntimeError("Replay produced duplicate update contents") From 4f98768a8c44759202ca86fcacd596d261dbc3bb Mon Sep 17 00:00:00 2001 From: Vivek Kalyan Date: Thu, 23 Jul 2026 19:25:13 -0700 Subject: [PATCH 7/9] fix: Complete LoRA replay evidence contract --- .../lora-update-replay-h200.sky.yaml | 16 +- .../benchmarks/prepare_lora_update_replay.py | 2 + .../benchmarks/service_lora_update_replay.py | 205 +++++++++++++++++- .../benchmarks/upload_lora_update_replay.py | 59 +++++ tests/unit/test_weight_update_replay.py | 26 +++ 5 files changed, 300 insertions(+), 8 deletions(-) create mode 100644 scripts/benchmarks/upload_lora_update_replay.py diff --git a/scripts/benchmarks/lora-update-replay-h200.sky.yaml b/scripts/benchmarks/lora-update-replay-h200.sky.yaml index 3cbc372ed..c62d54cbe 100644 --- a/scripts/benchmarks/lora-update-replay-h200.sky.yaml +++ b/scripts/benchmarks/lora-update-replay-h200.sky.yaml @@ -33,19 +33,31 @@ run: | --output scratch/lora-update-replay/inputs request_jsonl="$PWD/scratch/lora-update-replay/inputs/requests.jsonl" + mkdir -p scratch/lora-update-replay/fixed-load uv run --no-sync python scripts/benchmarks/service_lora_update_replay.py \ --request-jsonl "$request_jsonl" \ --mode fixed-load \ --warmups 5 \ --updates 30 \ - --output scratch/lora-update-replay/fixed-load + --output scratch/lora-update-replay/fixed-load \ + > scratch/lora-update-replay/fixed-load/stdout.log \ + 2> scratch/lora-update-replay/fixed-load/stderr.log + uv run --no-sync python scripts/benchmarks/upload_lora_update_replay.py \ + --input scratch/lora-update-replay/fixed-load \ + --mode fixed-load + mkdir -p scratch/lora-update-replay/idle uv run --no-sync python scripts/benchmarks/service_lora_update_replay.py \ --request-jsonl "$request_jsonl" \ --mode idle \ --warmups 5 \ --updates 30 \ - --output scratch/lora-update-replay/idle + --output scratch/lora-update-replay/idle \ + > scratch/lora-update-replay/idle/stdout.log \ + 2> scratch/lora-update-replay/idle/stderr.log + uv run --no-sync python scripts/benchmarks/upload_lora_update_replay.py \ + --input scratch/lora-update-replay/idle \ + --mode idle config: kubernetes: diff --git a/scripts/benchmarks/prepare_lora_update_replay.py b/scripts/benchmarks/prepare_lora_update_replay.py index 049f91b7c..7626c592c 100644 --- a/scripts/benchmarks/prepare_lora_update_replay.py +++ b/scripts/benchmarks/prepare_lora_update_replay.py @@ -54,6 +54,8 @@ def main() -> None: "temperature": 0.8, "max_tokens": 256, "logprobs": True, + "stream": True, + "stream_options": {"include_usage": True}, "chat_template_kwargs": {"enable_thinking": False}, } if row.get("tools"): diff --git a/scripts/benchmarks/service_lora_update_replay.py b/scripts/benchmarks/service_lora_update_replay.py index 75ee61749..dc09091b2 100644 --- a/scripts/benchmarks/service_lora_update_replay.py +++ b/scripts/benchmarks/service_lora_update_replay.py @@ -4,10 +4,12 @@ import argparse import asyncio +from dataclasses import asdict import hashlib import json from pathlib import Path import random +import shutil import statistics import subprocess import time @@ -41,6 +43,7 @@ def _load_requests(path: Path) -> list[dict[str, Any]]: "temperature": 0.8, "max_tokens": 256, "logprobs": True, + "stream": True, } for key, value in expected.items(): if request.get(key) != value: @@ -64,6 +67,7 @@ async def _load_worker( worker: int, stop: asyncio.Event, samples: list[dict[str, Any]], + phase: list[str], ) -> None: index = worker while not stop.is_set(): @@ -72,20 +76,53 @@ async def _load_worker( started = time.perf_counter() error = None prompt_tokens = completion_tokens = 0 + ttft_s: float | None = None + last_token_time_by_choice: dict[int, float] = {} + inter_token_intervals_s: list[float] = [] try: - response = await client.post("/v1/chat/completions", json=row["request"]) - response.raise_for_status() - usage = response.json().get("usage", {}) - prompt_tokens = int(usage.get("prompt_tokens", 0)) - completion_tokens = int(usage.get("completion_tokens", 0)) + async with client.stream( + "POST", "/v1/chat/completions", json=row["request"] + ) as response: + response.raise_for_status() + async for line in response.aiter_lines(): + if not line.startswith("data: ") or line == "data: [DONE]": + continue + chunk = json.loads(line.removeprefix("data: ")) + usage = chunk.get("usage") or {} + prompt_tokens = int(usage.get("prompt_tokens", prompt_tokens)) + completion_tokens = int( + usage.get("completion_tokens", completion_tokens) + ) + for choice in chunk.get("choices") or []: + choice_index = int(choice.get("index", 0)) + logprobs = choice.get("logprobs") or {} + emitted = len(logprobs.get("content") or []) + if emitted == 0: + delta = choice.get("delta") or {} + emitted = int( + bool(delta.get("content") or delta.get("tool_calls")) + ) + if emitted: + now = time.perf_counter() + if ttft_s is None: + ttft_s = now - started + previous = last_token_time_by_choice.get(choice_index) + if previous is not None: + inter_token_intervals_s.append(now - previous) + if emitted > 1: + inter_token_intervals_s.extend([0.0] * (emitted - 1)) + last_token_time_by_choice[choice_index] = now except Exception as exc: error = f"{type(exc).__name__}: {exc}" samples.append( { "worker": worker, + "phase": phase[0], "id": row["id"], "request_hash": row["request_hash"], "latency_s": time.perf_counter() - started, + "ttft_s": ttft_s, + "inter_token_intervals_s": inter_token_intervals_s, "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "error": error, @@ -93,6 +130,34 @@ async def _load_worker( ) +async def _record_control( + *, + name: str, + duration_s: float, + phase: list[str], + samples: list[dict[str, Any]], +) -> dict[str, Any]: + phase[0] = name + sample_start = len(samples) + started = time.perf_counter() + await asyncio.sleep(duration_s) + elapsed_s = time.perf_counter() - started + sample_end = len(samples) + lane = samples[sample_start:sample_end] + completion_tokens = sum(int(row["completion_tokens"]) for row in lane) + return { + "name": name, + "duration_s": elapsed_s, + "sample_start": sample_start, + "sample_end": sample_end, + "request_count": len(lane), + "request_errors": sum(row["error"] is not None for row in lane), + "prompt_tokens": sum(int(row["prompt_tokens"]) for row in lane), + "completion_tokens": completion_tokens, + "completion_tokens_per_s": completion_tokens / elapsed_s, + } + + def _bootstrap_ci(values: list[float]) -> list[float]: rng = random.Random(20260723) means = [statistics.fmean(rng.choices(values, k=len(values))) for _ in range(2000)] @@ -147,7 +212,7 @@ async def run(args: argparse.Namespace) -> None: "lora_config": {"rank": 8}, "engine_args": { "tensor_parallel_size": 2, - "max_model_len": 16384, + "max_model_len": 32769, "max_num_batched_tokens": 131072, "max_num_seqs": 256, "enable_prefix_caching": True, @@ -163,6 +228,8 @@ async def run(args: argparse.Namespace) -> None: if service._in_flight_lora_slot != "bonnie-replay:active": raise RuntimeError("Replay slot name changed unexpectedly") samples: list[dict[str, Any]] = [] + controls: list[dict[str, Any]] = [] + phase = ["idle"] stop = asyncio.Event() async with httpx.AsyncClient( base_url=base_url, headers=headers, timeout=300.0 @@ -179,6 +246,7 @@ async def run(args: argparse.Namespace) -> None: worker=worker, stop=stop, samples=samples, + phase=phase, ) ) for worker in range(8) @@ -188,18 +256,31 @@ async def run(args: argparse.Namespace) -> None: ) updates: list[dict[str, Any]] = [] hashes: set[str] = set() + manifest_sha256: str | None = None try: + if load_tasks: + controls.append( + await _record_control( + name="control-before", + duration_s=args.control_seconds, + phase=phase, + samples=samples, + ) + ) + phase[0] = "updates" for index in range(args.warmups + args.updates): policy_version = index + 1 checkpoint = output / "checkpoints" / f"{policy_version:04d}" sample_start = len(samples) memory_before = _gpu_memory_snapshot() + update_started = time.perf_counter() metrics = await service.replay_lora_update( source_lora_path=str(seed), output_lora_path=str(checkpoint), policy_version=policy_version, content_version=index, ) + update_wall_s = time.perf_counter() - update_started expected_topology = { "topology/trainer_tp": 1.0, "topology/trainer_pp": 1.0, @@ -220,6 +301,10 @@ async def run(args: argparse.Namespace) -> None: snapshot = inspect_adapter(checkpoint) if snapshot.file_sha256 in hashes: raise RuntimeError("Replay produced duplicate update contents") + if manifest_sha256 is None: + manifest_sha256 = snapshot.manifest_sha256 + elif snapshot.manifest_sha256 != manifest_sha256: + raise RuntimeError("Replay changed the serving tensor manifest") hashes.add(snapshot.file_sha256) updates.append( { @@ -228,13 +313,42 @@ async def run(args: argparse.Namespace) -> None: "policy_version": policy_version, "checkpoint": str(checkpoint), "sha256": snapshot.file_sha256, + "manifest_sha256": snapshot.manifest_sha256, "metrics": metrics, + "update_wall_s": update_wall_s, "gpu_memory_before": memory_before, "gpu_memory_after": _gpu_memory_snapshot(), "sample_start": sample_start, "sample_end": len(samples), } ) + if len(updates) > 2: + shutil.rmtree(Path(updates[-3]["checkpoint"])) + if load_tasks: + controls.append( + await _record_control( + name="control-after", + duration_s=args.control_seconds, + phase=phase, + samples=samples, + ) + ) + before_rate = float(controls[0]["completion_tokens_per_s"]) + after_rate = float(controls[-1]["completion_tokens_per_s"]) + drift = ( + abs(after_rate - before_rate) / before_rate + if before_rate > 0 + else float("inf") + ) + if drift > 0.05: + controls.append( + await _record_control( + name="control-repeat", + duration_s=args.control_seconds, + phase=phase, + samples=samples, + ) + ) finally: stop.set() await asyncio.gather(*load_tasks) @@ -247,6 +361,11 @@ async def run(args: argparse.Namespace) -> None: + float(row["metrics"]["time/weight_update_service_rpc_s"]) for row in measured ] + measured_sample_rows = [ + sample + for update in measured + for sample in samples[update["sample_start"] : update["sample_end"]] + ] (output / "updates.jsonl").write_text( "".join(json.dumps(row, sort_keys=True) + "\n" for row in updates) ) @@ -262,22 +381,87 @@ async def run(args: argparse.Namespace) -> None: "trainer": {"tp": 1, "pp": 1, "cp": 2, "ep": 2, "etp": 1}, } (output / "topology.json").write_text(json.dumps(topology, indent=2)) + reference_snapshot = inspect_adapter(Path(updates[-1]["checkpoint"])) + tensor_manifest = { + "schema_version": 1, + "manifest_sha256": reference_snapshot.manifest_sha256, + "rank": reference_snapshot.rank, + "base_model": reference_snapshot.base_model, + "logical_bytes": reference_snapshot.logical_bytes, + "transported_bytes": reference_snapshot.transported_bytes, + "tensor_count": reference_snapshot.tensor_count, + "tensors": [asdict(tensor) for tensor in reference_snapshot.tensors], + } + (output / "tensor-manifest.json").write_text(json.dumps(tensor_manifest, indent=2)) + request_trace = { + "schema_version": 1, + "source_artifact": "wb-training/bench/bonnie:v1", + "source_train_jsonl_sha256": ( + "fd26c97423de94a158d5a54e38485ec55b3fd3504f5f0346850fcfa79ad3ce66" + ), + "selection": "first 64 non-empty train.jsonl rows in artifact order", + "requests": requests, + } + (output / "request-trace.json").write_text( + json.dumps(request_trace, indent=2, sort_keys=True) + ) summary = { "update_total_mean_s": statistics.fmean(totals), "update_total_p50_s": statistics.median(totals), "update_total_p95_s": sorted(totals)[round(0.95 * (len(totals) - 1))], "update_total_mean_95ci_s": _bootstrap_ci(totals), + "update_total_stddev_s": statistics.stdev(totals), "prompt_tokens": sum(row["prompt_tokens"] for row in samples), "completion_tokens": sum(row["completion_tokens"] for row in samples), "request_latency_mean_s": ( statistics.fmean(row["latency_s"] for row in samples) if samples else None ), "request_errors": sum(row["error"] is not None for row in samples), + "controls": controls, } + ttfts = [ + float(row["ttft_s"]) + for row in samples + if row["ttft_s"] is not None and row["error"] is None + ] + itls = [ + float(interval) + for row in samples + if row["error"] is None + for interval in row["inter_token_intervals_s"] + ] + if samples: + summary["ttft_p50_s"] = statistics.median(ttfts) if ttfts else None + summary["ttft_p95_s"] = ( + sorted(ttfts)[round(0.95 * (len(ttfts) - 1))] if ttfts else None + ) + summary["inter_token_latency_p50_s"] = statistics.median(itls) if itls else None + summary["inter_token_latency_p95_s"] = ( + sorted(itls)[round(0.95 * (len(itls) - 1))] if itls else None + ) + if controls: + invalid_controls = [ + row for row in controls if not row["request_count"] or row["request_errors"] + ] + if invalid_controls: + raise RuntimeError(f"Invalid no-update control lanes: {invalid_controls}") + control_rate = statistics.fmean( + float(row["completion_tokens_per_s"]) for row in controls + ) + measured_lane_s = sum(float(row["update_wall_s"]) for row in measured) + measured_completion_tokens = sum( + int(row["completion_tokens"]) for row in measured_sample_rows + ) + summary["lost_tokens_per_update"] = ( + control_rate * measured_lane_s - measured_completion_tokens + ) / len(measured) (output / "summary.json").write_text(json.dumps(summary, indent=2)) manifest = { "schema_version": 1, "mode": args.mode, + "art_commit": subprocess.check_output( + ["git", "rev-parse", "HEAD"], text=True + ).strip(), "request_trace": str(args.request_jsonl), "request_ids_and_hashes": [ {"id": row["id"], "request_hash": row["request_hash"]} for row in requests @@ -292,10 +476,18 @@ async def run(args: argparse.Namespace) -> None: "runtime-before.json", "runtime-after.json", "topology.json", + "tensor-manifest.json", + "request-trace.json", "summary.json", + "stdout.log", + "stderr.log", ], } (output / "manifest.json").write_text(json.dumps(manifest, indent=2)) + if summary["request_errors"]: + raise RuntimeError( + f"Evidence run observed {summary['request_errors']} request errors" + ) def main() -> None: @@ -304,6 +496,7 @@ def main() -> None: parser.add_argument("--mode", choices=("idle", "fixed-load"), required=True) parser.add_argument("--warmups", type=int, default=5) parser.add_argument("--updates", type=int, default=30) + parser.add_argument("--control-seconds", type=float, default=60.0) parser.add_argument("--output", type=Path, required=True) asyncio.run(run(parser.parse_args())) diff --git a/scripts/benchmarks/upload_lora_update_replay.py b/scripts/benchmarks/upload_lora_update_replay.py new file mode 100644 index 000000000..969730c8b --- /dev/null +++ b/scripts/benchmarks/upload_lora_update_replay.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path + +REQUIRED_FILES = { + "manifest.json", + "request-trace.json", + "tensor-manifest.json", + "samples.jsonl", + "summary.json", + "stdout.log", + "stderr.log", +} + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--input", type=Path, required=True) + parser.add_argument("--mode", choices=("idle", "fixed-load"), required=True) + args = parser.parse_args() + + missing = sorted( + name for name in REQUIRED_FILES if not (args.input / name).is_file() + ) + if missing: + raise RuntimeError(f"Refusing to upload incomplete evidence: {missing}") + + import wandb + + run = wandb.init( + entity=os.environ.get("WANDB_ENTITY", "wb-training"), + project=os.environ.get("WANDB_PROJECT", "bench"), + job_type="lora-update-replay", + config={"mode": args.mode, "artifact_schema_version": 1}, + ) + assert run is not None + artifact = wandb.Artifact( + name=f"art-lora-update-replay-{args.mode}", + type="lora-update-replay-result", + metadata=json.loads((args.input / "summary.json").read_text()), + ) + artifact.add_dir(str(args.input)) + logged = run.log_artifact(artifact) + logged.wait() + artifact_ref = f"{logged.entity}/{logged.project}/{logged.name}" + run.summary["result_artifact"] = artifact_ref + run.summary["result_artifact_digest"] = logged.digest + run.finish() + print(f"WANDB_ARTIFACT={artifact_ref}") + print(f"WANDB_ARTIFACT_DIGEST={logged.digest}") + + +if __name__ == "__main__": + main() diff --git a/tests/unit/test_weight_update_replay.py b/tests/unit/test_weight_update_replay.py index ad564a133..4b9fb763c 100644 --- a/tests/unit/test_weight_update_replay.py +++ b/tests/unit/test_weight_update_replay.py @@ -162,3 +162,29 @@ def test_snapshot_config_is_machine_readable(tmp_path: Path) -> None: config = json.loads((path / "adapter_config.json").read_text()) assert config["art_lora_format"] == "vllm" + + +def test_h200_replay_launcher_preserves_evidence_contract() -> None: + root = Path(__file__).parents[2] + launcher = ( + root / "scripts/benchmarks/lora-update-replay-h200.sky.yaml" + ).read_text() + service = (root / "scripts/benchmarks/service_lora_update_replay.py").read_text() + uploader = (root / "scripts/benchmarks/upload_lora_update_replay.py").read_text() + + assert "accelerators: H200:4" in launcher + assert "name: vivek-dev-api-keys" in launcher + assert "--mode fixed-load" in launcher + assert "--mode idle" in launcher + assert 'max_model_len": 32769' in service + assert "control_seconds" in service + for filename in ( + "manifest.json", + "request-trace.json", + "tensor-manifest.json", + "samples.jsonl", + "summary.json", + "stdout.log", + "stderr.log", + ): + assert filename in uploader From 4bcdcee5dfccdea04171fe688933f221b0b70cef Mon Sep 17 00:00:00 2001 From: Vivek Kalyan Date: Thu, 23 Jul 2026 19:46:03 -0700 Subject: [PATCH 8/9] fix: Make replay evidence receiver authoritative --- .../lora-update-replay-h200.sky.yaml | 72 --- .../benchmarks/prepare_lora_update_replay.py | 85 --- .../benchmarks/service_lora_update_replay.py | 539 ++++++++++++------ .../benchmarks/upload_lora_update_replay.py | 25 +- src/art/megatron/service.py | 69 ++- src/art/megatron/weights/update_replay.py | 137 ++++- .../test_runtime_project_isolation.py | 11 + .../test_service_runtime_boundary.py | 30 + tests/unit/test_weight_update_replay.py | 122 +++- .../src/art_vllm_runtime/dedicated_server.py | 42 +- .../src/art_vllm_runtime/policy_spans.py | 22 + 11 files changed, 801 insertions(+), 353 deletions(-) delete mode 100644 scripts/benchmarks/lora-update-replay-h200.sky.yaml delete mode 100644 scripts/benchmarks/prepare_lora_update_replay.py diff --git a/scripts/benchmarks/lora-update-replay-h200.sky.yaml b/scripts/benchmarks/lora-update-replay-h200.sky.yaml deleted file mode 100644 index c62d54cbe..000000000 --- a/scripts/benchmarks/lora-update-replay-h200.sky.yaml +++ /dev/null @@ -1,72 +0,0 @@ -name: art-lora-update-replay-h200 - -workdir: . - -resources: - cloud: kubernetes - accelerators: H200:4 - cpus: "32+" - memory: "256+" - -envs: - CUDA_HOME: /usr/local/cuda-12.8 - TORCH_CUDA_ARCH_LIST: "9.0" - UV_CONCURRENT_BUILDS: "1" - UV_LINK_MODE: copy - WANDB_ENTITY: wb-training - WANDB_PROJECT: bench - -setup: | - set -euo pipefail - curl -LsSf https://astral.sh/uv/install.sh | sh - export PATH="$HOME/.local/bin:$PATH" - uv sync --extra megatron - uv run --no-sync bash src/art/megatron/setup.sh - cd vllm_runtime - uv sync - -run: | - set -euo pipefail - export PATH="$HOME/.local/bin:$PATH" - mkdir -p scratch/lora-update-replay - uv run --no-sync python scripts/benchmarks/prepare_lora_update_replay.py \ - --output scratch/lora-update-replay/inputs - - request_jsonl="$PWD/scratch/lora-update-replay/inputs/requests.jsonl" - mkdir -p scratch/lora-update-replay/fixed-load - uv run --no-sync python scripts/benchmarks/service_lora_update_replay.py \ - --request-jsonl "$request_jsonl" \ - --mode fixed-load \ - --warmups 5 \ - --updates 30 \ - --output scratch/lora-update-replay/fixed-load \ - > scratch/lora-update-replay/fixed-load/stdout.log \ - 2> scratch/lora-update-replay/fixed-load/stderr.log - uv run --no-sync python scripts/benchmarks/upload_lora_update_replay.py \ - --input scratch/lora-update-replay/fixed-load \ - --mode fixed-load - - mkdir -p scratch/lora-update-replay/idle - uv run --no-sync python scripts/benchmarks/service_lora_update_replay.py \ - --request-jsonl "$request_jsonl" \ - --mode idle \ - --warmups 5 \ - --updates 30 \ - --output scratch/lora-update-replay/idle \ - > scratch/lora-update-replay/idle/stdout.log \ - 2> scratch/lora-update-replay/idle/stderr.log - uv run --no-sync python scripts/benchmarks/upload_lora_update_replay.py \ - --input scratch/lora-update-replay/idle \ - --mode idle - -config: - kubernetes: - pod_config: - spec: - schedulerName: binpack-scheduler - activeDeadlineSeconds: 21600 - containers: - - name: ray-node - envFrom: - - secretRef: - name: vivek-dev-api-keys diff --git a/scripts/benchmarks/prepare_lora_update_replay.py b/scripts/benchmarks/prepare_lora_update_replay.py deleted file mode 100644 index 7626c592c..000000000 --- a/scripts/benchmarks/prepare_lora_update_replay.py +++ /dev/null @@ -1,85 +0,0 @@ -#!/usr/bin/env python3 - -import argparse -import hashlib -import json -from pathlib import Path - -BONNIE_ARTIFACT = "wb-training/bench/bonnie:v1" -BONNIE_TRAIN_SHA256 = "fd26c97423de94a158d5a54e38485ec55b3fd3504f5f0346850fcfa79ad3ce66" - - -def _sha256(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as handle: - for chunk in iter(lambda: handle.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--output", type=Path, required=True) - args = parser.parse_args() - artifact_dir = args.output / "bonnie-v1" - artifact_dir.mkdir(parents=True, exist_ok=True) - - import wandb - - wandb.Api().artifact(BONNIE_ARTIFACT).download(root=str(artifact_dir)) - train_path = artifact_dir / "train.jsonl" - digest = _sha256(train_path) - if digest != BONNIE_TRAIN_SHA256: - raise RuntimeError( - f"Bonnie train artifact checksum is {digest}; " - f"expected {BONNIE_TRAIN_SHA256}" - ) - request_path = args.output / "requests.jsonl" - requests: list[dict[str, object]] = [] - with train_path.open(encoding="utf-8") as handle: - for source_index, line in enumerate(handle): - if not line.strip(): - continue - row = json.loads(line) - messages = [dict(message) for message in row["messages"]] - for index, message in enumerate(messages): - if index > 0 and message.get("role") == "system": - message["role"] = "user" - message["content"] = "[System note]\n" + str( - message.get("content") or "" - ) - request: dict[str, object] = { - "messages": messages, - "n": 8, - "temperature": 0.8, - "max_tokens": 256, - "logprobs": True, - "stream": True, - "stream_options": {"include_usage": True}, - "chat_template_kwargs": {"enable_thinking": False}, - } - if row.get("tools"): - request["tools"] = row["tools"] - requests.append( - { - "id": str( - row.get("id") - or ( - f"mar31_v2:train:{row.get('trace_id', source_index)}:" - f"assistant-{row.get('assistant_turn_index', 0)}" - ) - ), - "request": request, - } - ) - if len(requests) == 64: - break - request_path.write_text( - "".join(json.dumps(row, sort_keys=True) + "\n" for row in requests) - ) - print(f"BONNIE_TRAIN_JSONL={train_path}") - print(f"REQUEST_JSONL={request_path}") - - -if __name__ == "__main__": - main() diff --git a/scripts/benchmarks/service_lora_update_replay.py b/scripts/benchmarks/service_lora_update_replay.py index dc09091b2..ba2e6a8fe 100644 --- a/scripts/benchmarks/service_lora_update_replay.py +++ b/scripts/benchmarks/service_lora_update_replay.py @@ -9,6 +9,7 @@ import json from pathlib import Path import random +import resource import shutil import statistics import subprocess @@ -19,45 +20,61 @@ import art from art.megatron.service import MegatronService -from art.megatron.weights.update_replay import inspect_adapter +from art.megatron.weights.update_replay import ( + evaluate_replay_acceptance, + inspect_adapter, + validate_bench_request_trace, +) BASE_MODEL = "Qwen/Qwen3.6-35B-A3B" -def _hash_request(request: dict[str, Any]) -> str: - encoded = json.dumps(request, sort_keys=True, separators=(",", ":")).encode() - return hashlib.sha256(encoded).hexdigest() +def _canonical_sha256(value: Any) -> str: + payload = json.dumps( + value, ensure_ascii=False, separators=(",", ":"), sort_keys=True + ).encode() + return hashlib.sha256(payload).hexdigest() -def _load_requests(path: Path) -> list[dict[str, Any]]: - rows = [json.loads(line) for line in path.read_text().splitlines() if line] - if len(rows) != 64: - raise ValueError(f"Expected 64 pinned request rows, got {len(rows)}") - ids: set[str] = set() - hashes: set[str] = set() +def _effective_requests( + request_jsonl: Path, trace_manifest: Path, *, bench_commit: str +) -> tuple[list[dict[str, Any]], dict[str, Any]]: + rows, provenance = validate_bench_request_trace( + request_jsonl, trace_manifest, bench_commit=bench_commit + ) + generation = provenance["manifest"].get("generation", {}) + expected = { + "n": 8, + "temperature": 0.8, + "max_tokens": 256, + "logprobs": True, + "enable_thinking": False, + "load_concurrency": 8, + } + if {key: generation.get(key) for key in expected} != expected: + raise ValueError(f"Unexpected Bench generation contract: {generation!r}") for row in rows: - request = row["request"] - request["model"] = "bonnie-replay:active" - expected = { - "n": 8, - "temperature": 0.8, - "max_tokens": 256, - "logprobs": True, - "stream": True, - } - for key, value in expected.items(): - if request.get(key) != value: - raise ValueError( - f"Request {row.get('id')} has {key}={request.get(key)}" - ) - if request.get("chat_template_kwargs") != {"enable_thinking": False}: - raise ValueError(f"Request {row.get('id')} enables thinking") - row["request_hash"] = _hash_request(request) - ids.add(str(row["id"])) - hashes.add(row["request_hash"]) - if len(ids) != 64 or len(hashes) != 64: - raise ValueError("Pinned request trace contains duplicate IDs or requests") - return rows + effective = dict(row["request"]) + effective.update( + { + "model": "bonnie-replay:active", + "stream": True, + "stream_options": {"include_usage": True}, + } + ) + row["effective_request"] = effective + row["effective_request_sha256"] = _canonical_sha256(effective) + return rows, provenance + + +async def _runtime_metrics(client: httpx.AsyncClient) -> dict[str, float]: + response = await client.get("/art/metrics") + response.raise_for_status() + return { + key: float(value) + for key, value in response.json().get("metrics", {}).items() + if isinstance(value, (int, float)) + } async def _load_worker( @@ -67,21 +84,30 @@ async def _load_worker( worker: int, stop: asyncio.Event, samples: list[dict[str, Any]], + token_events: list[dict[str, Any]], phase: list[str], ) -> None: + sequence = 0 index = worker while not stop.is_set(): row = requests[index % len(requests)] index += 8 + request_id = f"load-{worker}-{sequence:06d}" + sequence += 1 started = time.perf_counter() + phase_at_start = phase[0] error = None prompt_tokens = completion_tokens = 0 ttft_s: float | None = None last_token_time_by_choice: dict[int, float] = {} inter_token_intervals_s: list[float] = [] + spans_by_choice: dict[int, list[dict[str, Any]]] = {} try: async with client.stream( - "POST", "/v1/chat/completions", json=row["request"] + "POST", + "/v1/chat/completions", + json=row["effective_request"], + headers={"X-Request-Id": request_id}, ) as response: response.raise_for_status() async for line in response.aiter_lines(): @@ -95,6 +121,10 @@ async def _load_worker( ) for choice in chunk.get("choices") or []: choice_index = int(choice.get("index", 0)) + spans = choice.get("policy_token_spans") or [] + spans_by_choice.setdefault(choice_index, []).extend( + dict(span) for span in spans + ) logprobs = choice.get("logprobs") or {} emitted = len(logprobs.get("content") or []) if emitted == 0: @@ -112,19 +142,32 @@ async def _load_worker( if emitted > 1: inter_token_intervals_s.extend([0.0] * (emitted - 1)) last_token_time_by_choice[choice_index] = now + token_events.append( + { + "request_id": request_id, + "monotonic_s": now, + "tokens": emitted, + } + ) except Exception as exc: error = f"{type(exc).__name__}: {exc}" + ended = time.perf_counter() samples.append( { + "request_id": request_id, "worker": worker, - "phase": phase[0], - "id": row["id"], - "request_hash": row["request_hash"], - "latency_s": time.perf_counter() - started, + "phase_at_start": phase_at_start, + "datapoint_id": row["id"], + "canonical_request_sha256": row["canonical_request_sha256"], + "effective_request_sha256": row["effective_request_sha256"], + "started_monotonic_s": started, + "ended_monotonic_s": ended, + "latency_s": ended - started, "ttft_s": ttft_s, "inter_token_intervals_s": inter_token_intervals_s, "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, + "policy_token_spans_by_choice": spans_by_choice, "error": error, } ) @@ -132,37 +175,72 @@ async def _load_worker( async def _record_control( *, + client: httpx.AsyncClient, name: str, duration_s: float, phase: list[str], - samples: list[dict[str, Any]], ) -> dict[str, Any]: phase[0] = name - sample_start = len(samples) + before = await _runtime_metrics(client) started = time.perf_counter() await asyncio.sleep(duration_s) - elapsed_s = time.perf_counter() - started - sample_end = len(samples) - lane = samples[sample_start:sample_end] - completion_tokens = sum(int(row["completion_tokens"]) for row in lane) + ended = time.perf_counter() + after = await _runtime_metrics(client) + elapsed = ended - started + prompt_tokens = after["prompt_tokens_total"] - before["prompt_tokens_total"] + completion_tokens = ( + after["generation_tokens_total"] - before["generation_tokens_total"] + ) return { "name": name, - "duration_s": elapsed_s, - "sample_start": sample_start, - "sample_end": sample_end, - "request_count": len(lane), - "request_errors": sum(row["error"] is not None for row in lane), - "prompt_tokens": sum(int(row["prompt_tokens"]) for row in lane), + "started_monotonic_s": started, + "ended_monotonic_s": ended, + "duration_s": elapsed, + "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, - "completion_tokens_per_s": completion_tokens / elapsed_s, + "prompt_tokens_per_s": prompt_tokens / elapsed, + "completion_tokens_per_s": completion_tokens / elapsed, } -def _bootstrap_ci(values: list[float]) -> list[float]: - rng = random.Random(20260723) - means = [statistics.fmean(rng.choices(values, k=len(values))) for _ in range(2000)] - means.sort() - return [means[50], means[1949]] +async def _fixed_input_fingerprint( + client: httpx.AsyncClient, *, policy_version: int +) -> dict[str, Any]: + request = { + "model": "bonnie-replay:active", + "messages": [{"role": "user", "content": "Reply with exactly: OK"}], + "temperature": 0, + "max_tokens": 1, + "n": 1, + "logprobs": True, + "top_logprobs": 5, + "chat_template_kwargs": {"enable_thinking": False}, + } + fingerprints: list[str] = [] + for _ in range(2): + response = await client.post("/v1/chat/completions", json=request) + response.raise_for_status() + body = response.json() + choice = body["choices"][0] + spans = choice.get("policy_token_spans") or [] + versions = {int(span["policy_version"]) for span in spans} + if versions != {policy_version}: + raise RuntimeError( + f"Fixed-input request used policies {versions}; expected {policy_version}" + ) + receipt = { + "message": choice["message"], + "logprobs": choice.get("logprobs"), + } + fingerprints.append(_canonical_sha256(receipt)) + if len(set(fingerprints)) != 1: + raise RuntimeError("Fixed-input receiver logit parity failed") + return { + "method": "two deterministic receiver requests after commit", + "sha256": fingerprints[0], + "policy_version": policy_version, + "policy_token_spans": spans, + } def _gpu_memory_snapshot() -> list[dict[str, int | str]]: @@ -190,10 +268,60 @@ def _gpu_memory_snapshot() -> list[dict[str, int | str]]: return rows +def _cpu_memory_snapshot() -> dict[str, int]: + meminfo: dict[str, int] = {} + for line in Path("/proc/meminfo").read_text().splitlines(): + key, raw = line.split(":", maxsplit=1) + meminfo[key] = int(raw.strip().split()[0]) + return { + "system_used_kib": meminfo["MemTotal"] - meminfo["MemAvailable"], + "runner_peak_rss_kib": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss, + } + + +def _bootstrap_ci(values: list[float]) -> list[float]: + rng = random.Random(20260723) + means = [statistics.fmean(rng.choices(values, k=len(values))) for _ in range(2000)] + means.sort() + return [means[50], means[1949]] + + +def _attach_policy_ages( + samples: list[dict[str, Any]], updates: list[dict[str, Any]] +) -> None: + commits = [ + ( + float(update["committed_monotonic_s"]), + int(update["receiver"]["committed_state"]["policy_version"]), + ) + for update in updates + ] + for sample in samples: + current = max( + ( + version + for committed_at, version in commits + if committed_at <= float(sample["ended_monotonic_s"]) + ), + default=0, + ) + versions = [ + int(span["policy_version"]) + for spans in sample["policy_token_spans_by_choice"].values() + for span in spans + ] + sample["receiver_policy_version_at_end"] = current + sample["policy_age_min"] = current - max(versions) if versions else None + sample["policy_age_max"] = current - min(versions) if versions else None + + async def run(args: argparse.Namespace) -> None: - requests = _load_requests(args.request_jsonl) + requests, trace_provenance = _effective_requests( + args.request_jsonl, args.trace_manifest, bench_commit=args.bench_commit + ) output = args.output.resolve() output.mkdir(parents=True, exist_ok=True) + shutil.copyfile(args.request_jsonl, output / "request-trace.jsonl") art.init_megatron_runtime_config( topology=art.MegatronTopologyConfig(tp=1, pp=1, cp=2, ep=2, etp=1), packed_sequence_length=122880, @@ -223,19 +351,19 @@ async def run(args: argparse.Namespace) -> None: ) seed = Path(service._resolve_active_lora_path()) await service.start_openai_server(None) - base_url = service._vllm_base_url - headers = service._runtime_headers() - if service._in_flight_lora_slot != "bonnie-replay:active": - raise RuntimeError("Replay slot name changed unexpectedly") samples: list[dict[str, Any]] = [] + token_events: list[dict[str, Any]] = [] controls: list[dict[str, Any]] = [] + updates: list[dict[str, Any]] = [] phase = ["idle"] stop = asyncio.Event() async with httpx.AsyncClient( - base_url=base_url, headers=headers, timeout=300.0 + base_url=service._vllm_base_url, + headers=service._runtime_headers(), + timeout=300.0, ) as client: - runtime_before = (await client.get("/art/metrics")).json() - if int(runtime_before.get("metrics", {}).get("world_size", 0)) != 2: + runtime_before = await _runtime_metrics(client) + if int(runtime_before.get("world_size", 0)) != 2: raise RuntimeError(f"Expected vLLM TP2 runtime: {runtime_before!r}") load_tasks = ( [ @@ -246,6 +374,7 @@ async def run(args: argparse.Namespace) -> None: worker=worker, stop=stop, samples=samples, + token_events=token_events, phase=phase, ) ) @@ -254,33 +383,46 @@ async def run(args: argparse.Namespace) -> None: if args.mode == "fixed-load" else [] ) - updates: list[dict[str, Any]] = [] hashes: set[str] = set() manifest_sha256: str | None = None try: if load_tasks: controls.append( await _record_control( + client=client, name="control-before", duration_s=args.control_seconds, phase=phase, - samples=samples, ) ) phase[0] = "updates" for index in range(args.warmups + args.updates): policy_version = index + 1 checkpoint = output / "checkpoints" / f"{policy_version:04d}" - sample_start = len(samples) - memory_before = _gpu_memory_snapshot() + memory_before = { + "gpu": _gpu_memory_snapshot(), + "cpu": _cpu_memory_snapshot(), + } + update_runtime_before = ( + await _runtime_metrics(client) + if index >= args.warmups and load_tasks + else None + ) update_started = time.perf_counter() - metrics = await service.replay_lora_update( + replay = await service.replay_lora_update( source_lora_path=str(seed), output_lora_path=str(checkpoint), policy_version=policy_version, content_version=index, ) - update_wall_s = time.perf_counter() - update_started + committed_at = time.perf_counter() + update_runtime_after = ( + await _runtime_metrics(client) + if update_runtime_before is not None + else None + ) + metrics = replay["metrics"] + receiver = replay["receiver"] expected_topology = { "topology/trainer_tp": 1.0, "topology/trainer_pp": 1.0, @@ -292,12 +434,10 @@ async def run(args: argparse.Namespace) -> None: key: metrics.get(key) for key in expected_topology } != expected_topology: raise RuntimeError(f"Unexpected trainer topology: {metrics!r}") - if len(memory_before) != 4: - raise RuntimeError( - f"Expected four visible H200s, got {memory_before!r}" - ) - if any("H200" not in str(row["name"]) for row in memory_before): - raise RuntimeError(f"Replay requires H200 GPUs: {memory_before!r}") + if len(memory_before["gpu"]) != 4 or any( + "H200" not in str(row["name"]) for row in memory_before["gpu"] + ): + raise RuntimeError(f"Replay requires four H200s: {memory_before!r}") snapshot = inspect_adapter(checkpoint) if snapshot.file_sha256 in hashes: raise RuntimeError("Replay produced duplicate update contents") @@ -306,20 +446,43 @@ async def run(args: argparse.Namespace) -> None: elif snapshot.manifest_sha256 != manifest_sha256: raise RuntimeError("Replay changed the serving tensor manifest") hashes.add(snapshot.file_sha256) + parity = await _fixed_input_fingerprint( + client, policy_version=policy_version + ) updates.append( { "index": index, "measured": index >= args.warmups, - "policy_version": policy_version, + "requested_policy_version": policy_version, "checkpoint": str(checkpoint), "sha256": snapshot.file_sha256, "manifest_sha256": snapshot.manifest_sha256, "metrics": metrics, - "update_wall_s": update_wall_s, - "gpu_memory_before": memory_before, - "gpu_memory_after": _gpu_memory_snapshot(), - "sample_start": sample_start, - "sample_end": len(samples), + "receiver": receiver, + "fixed_input_logit_parity": parity, + "update_started_monotonic_s": update_started, + "committed_monotonic_s": committed_at, + "update_wall_s": committed_at - update_started, + "runtime_counter_delta": ( + { + "prompt_tokens": ( + update_runtime_after["prompt_tokens_total"] + - update_runtime_before["prompt_tokens_total"] + ), + "completion_tokens": ( + update_runtime_after["generation_tokens_total"] + - update_runtime_before["generation_tokens_total"] + ), + } + if update_runtime_before is not None + and update_runtime_after is not None + else None + ), + "memory_before": memory_before, + "memory_after": { + "gpu": _gpu_memory_snapshot(), + "cpu": _cpu_memory_snapshot(), + }, } ) if len(updates) > 2: @@ -327,10 +490,10 @@ async def run(args: argparse.Namespace) -> None: if load_tasks: controls.append( await _record_control( + client=client, name="control-after", duration_s=args.control_seconds, phase=phase, - samples=samples, ) ) before_rate = float(controls[0]["completion_tokens_per_s"]) @@ -343,80 +506,84 @@ async def run(args: argparse.Namespace) -> None: if drift > 0.05: controls.append( await _record_control( + client=client, name="control-repeat", duration_s=args.control_seconds, phase=phase, - samples=samples, ) ) finally: stop.set() await asyncio.gather(*load_tasks) - runtime_after = (await client.get("/art/metrics")).json() + runtime_after = await _runtime_metrics(client) await service.aclose() measured = [row for row in updates if row["measured"]] + acceptance = evaluate_replay_acceptance([row["metrics"] for row in measured]) totals = [ float(row["metrics"]["time/weight_update_trainer_publish_s"]) + float(row["metrics"]["time/weight_update_service_rpc_s"]) for row in measured ] - measured_sample_rows = [ - sample - for update in measured - for sample in samples[update["sample_start"] : update["sample_end"]] - ] - (output / "updates.jsonl").write_text( - "".join(json.dumps(row, sort_keys=True) + "\n" for row in updates) + _attach_policy_ages(samples, updates) + prompt_tokens = sum( + float((row["runtime_counter_delta"] or {})["prompt_tokens"]) + for row in measured + if row["runtime_counter_delta"] is not None ) - (output / "samples.jsonl").write_text( - "".join(json.dumps(row, sort_keys=True) + "\n" for row in samples) + completion_tokens = sum( + float((row["runtime_counter_delta"] or {})["completion_tokens"]) + for row in measured + if row["runtime_counter_delta"] is not None ) - (output / "runtime-before.json").write_text(json.dumps(runtime_before, indent=2)) - (output / "runtime-after.json").write_text(json.dumps(runtime_after, indent=2)) - topology = { - "trainer_gpu_ids": [0, 1], - "inference_gpu_ids": [2, 3], - "inference_tp": 2, - "trainer": {"tp": 1, "pp": 1, "cp": 2, "ep": 2, "etp": 1}, - } - (output / "topology.json").write_text(json.dumps(topology, indent=2)) - reference_snapshot = inspect_adapter(Path(updates[-1]["checkpoint"])) - tensor_manifest = { - "schema_version": 1, - "manifest_sha256": reference_snapshot.manifest_sha256, - "rank": reference_snapshot.rank, - "base_model": reference_snapshot.base_model, - "logical_bytes": reference_snapshot.logical_bytes, - "transported_bytes": reference_snapshot.transported_bytes, - "tensor_count": reference_snapshot.tensor_count, - "tensors": [asdict(tensor) for tensor in reference_snapshot.tensors], - } - (output / "tensor-manifest.json").write_text(json.dumps(tensor_manifest, indent=2)) - request_trace = { - "schema_version": 1, - "source_artifact": "wb-training/bench/bonnie:v1", - "source_train_jsonl_sha256": ( - "fd26c97423de94a158d5a54e38485ec55b3fd3504f5f0346850fcfa79ad3ce66" - ), - "selection": "first 64 non-empty train.jsonl rows in artifact order", - "requests": requests, - } - (output / "request-trace.json").write_text( - json.dumps(request_trace, indent=2, sort_keys=True) + measured_duration = sum(float(row["update_wall_s"]) for row in measured) + control_completion_rate = ( + statistics.fmean(row["completion_tokens_per_s"] for row in controls) + if controls + else None ) - summary = { + ages = [ + int(sample["policy_age_max"]) + for sample in samples + if sample["policy_age_max"] is not None + ] + gpu_snapshots = [ + gpu + for update in updates + for boundary in ("memory_before", "memory_after") + for gpu in update[boundary]["gpu"] + ] + cpu_snapshots = [ + update[boundary]["cpu"] + for update in updates + for boundary in ("memory_before", "memory_after") + ] + summary: dict[str, Any] = { + "acceptance": acceptance, "update_total_mean_s": statistics.fmean(totals), "update_total_p50_s": statistics.median(totals), "update_total_p95_s": sorted(totals)[round(0.95 * (len(totals) - 1))], - "update_total_mean_95ci_s": _bootstrap_ci(totals), "update_total_stddev_s": statistics.stdev(totals), - "prompt_tokens": sum(row["prompt_tokens"] for row in samples), - "completion_tokens": sum(row["completion_tokens"] for row in samples), - "request_latency_mean_s": ( - statistics.fmean(row["latency_s"] for row in samples) if samples else None + "update_total_mean_95ci_s": _bootstrap_ci(totals), + "measured_lane_duration_s": measured_duration, + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "prompt_tokens_per_s": prompt_tokens / measured_duration, + "completion_tokens_per_s": completion_tokens / measured_duration, + "lost_tokens_per_update": ( + (control_completion_rate * measured_duration - completion_tokens) + / len(measured) + if control_completion_rate is not None + else None ), "request_errors": sum(row["error"] is not None for row in samples), + "policy_age_p50": statistics.median(ages) if ages else None, + "policy_age_p95": ( + sorted(ages)[round(0.95 * (len(ages) - 1))] if ages else None + ), + "peak_gpu_used_mib": max(int(row["used_mib"]) for row in gpu_snapshots), + "peak_system_used_kib": max(row["system_used_kib"] for row in cpu_snapshots), + "peak_runner_rss_kib": max(row["runner_peak_rss_kib"] for row in cpu_snapshots), "controls": controls, } ttfts = [ @@ -430,69 +597,95 @@ async def run(args: argparse.Namespace) -> None: if row["error"] is None for interval in row["inter_token_intervals_s"] ] - if samples: - summary["ttft_p50_s"] = statistics.median(ttfts) if ttfts else None - summary["ttft_p95_s"] = ( - sorted(ttfts)[round(0.95 * (len(ttfts) - 1))] if ttfts else None - ) - summary["inter_token_latency_p50_s"] = statistics.median(itls) if itls else None - summary["inter_token_latency_p95_s"] = ( - sorted(itls)[round(0.95 * (len(itls) - 1))] if itls else None - ) - if controls: - invalid_controls = [ - row for row in controls if not row["request_count"] or row["request_errors"] - ] - if invalid_controls: - raise RuntimeError(f"Invalid no-update control lanes: {invalid_controls}") - control_rate = statistics.fmean( - float(row["completion_tokens_per_s"]) for row in controls - ) - measured_lane_s = sum(float(row["update_wall_s"]) for row in measured) - measured_completion_tokens = sum( - int(row["completion_tokens"]) for row in measured_sample_rows - ) - summary["lost_tokens_per_update"] = ( - control_rate * measured_lane_s - measured_completion_tokens - ) / len(measured) + summary.update( + { + "ttft_p50_s": statistics.median(ttfts) if ttfts else None, + "ttft_p95_s": ( + sorted(ttfts)[round(0.95 * (len(ttfts) - 1))] if ttfts else None + ), + "inter_token_latency_p50_s": statistics.median(itls) if itls else None, + "inter_token_latency_p95_s": ( + sorted(itls)[round(0.95 * (len(itls) - 1))] if itls else None + ), + } + ) + (output / "updates.jsonl").write_text( + "".join(json.dumps(row, sort_keys=True) + "\n" for row in updates) + ) + (output / "samples.jsonl").write_text( + "".join(json.dumps(row, sort_keys=True) + "\n" for row in samples) + ) + (output / "token-events.jsonl").write_text( + "".join(json.dumps(row, sort_keys=True) + "\n" for row in token_events) + ) (output / "summary.json").write_text(json.dumps(summary, indent=2)) - manifest = { + (output / "runtime-before.json").write_text(json.dumps(runtime_before, indent=2)) + (output / "runtime-after.json").write_text(json.dumps(runtime_after, indent=2)) + reference_snapshot = inspect_adapter(Path(updates[-1]["checkpoint"])) + (output / "tensor-manifest.json").write_text( + json.dumps( + { + "schema_version": 1, + "manifest_sha256": reference_snapshot.manifest_sha256, + "rank": reference_snapshot.rank, + "base_model": reference_snapshot.base_model, + "logical_bytes": reference_snapshot.logical_bytes, + "transported_bytes": reference_snapshot.transported_bytes, + "tensor_count": reference_snapshot.tensor_count, + "tensors": [asdict(tensor) for tensor in reference_snapshot.tensors], + }, + indent=2, + ) + ) + required_result_files = [ + "manifest.json", + "request-trace.jsonl", + "tensor-manifest.json", + "samples.jsonl", + "summary.json", + "stdout.log", + "stderr.log", + ] + result_manifest = { "schema_version": 1, "mode": args.mode, "art_commit": subprocess.check_output( ["git", "rev-parse", "HEAD"], text=True ).strip(), - "request_trace": str(args.request_jsonl), - "request_ids_and_hashes": [ - {"id": row["id"], "request_hash": row["request_hash"]} for row in requests - ], - "routed_expert_collection": { - "supported": False, - "reason": "The replay uses vLLM's OpenAI JSON endpoint.", + "bench_trace": trace_provenance, + "effective_request_transform": { + "model": "bonnie-replay:active", + "stream": True, + "stream_options": {"include_usage": True}, + }, + "topology": { + "trainer_gpu_ids": [0, 1], + "inference_gpu_ids": [2, 3], + "inference_tp": 2, + "trainer": {"tp": 1, "pp": 1, "cp": 2, "ep": 2, "etp": 1}, }, - "artifacts": [ + "required_result_files": required_result_files, + "additional_result_files": [ "updates.jsonl", - "samples.jsonl", + "token-events.jsonl", "runtime-before.json", "runtime-after.json", - "topology.json", - "tensor-manifest.json", - "request-trace.json", - "summary.json", - "stdout.log", - "stderr.log", ], } - (output / "manifest.json").write_text(json.dumps(manifest, indent=2)) + (output / "manifest.json").write_text(json.dumps(result_manifest, indent=2)) if summary["request_errors"]: raise RuntimeError( f"Evidence run observed {summary['request_errors']} request errors" ) + if not acceptance["passed"]: + raise RuntimeError(f"Replay acceptance gates failed: {acceptance['failures']}") def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--request-jsonl", type=Path, required=True) + parser.add_argument("--trace-manifest", type=Path, required=True) + parser.add_argument("--bench-commit", required=True) parser.add_argument("--mode", choices=("idle", "fixed-load"), required=True) parser.add_argument("--warmups", type=int, default=5) parser.add_argument("--updates", type=int, default=30) diff --git a/scripts/benchmarks/upload_lora_update_replay.py b/scripts/benchmarks/upload_lora_update_replay.py index 969730c8b..33bd2db0f 100644 --- a/scripts/benchmarks/upload_lora_update_replay.py +++ b/scripts/benchmarks/upload_lora_update_replay.py @@ -9,7 +9,7 @@ REQUIRED_FILES = { "manifest.json", - "request-trace.json", + "request-trace.jsonl", "tensor-manifest.json", "samples.jsonl", "summary.json", @@ -29,6 +29,24 @@ def main() -> None: ) if missing: raise RuntimeError(f"Refusing to upload incomplete evidence: {missing}") + manifest = json.loads((args.input / "manifest.json").read_text()) + claimed = set(manifest.get("required_result_files", ())) + if claimed != REQUIRED_FILES: + raise RuntimeError( + f"Result manifest claims {sorted(claimed)}; expected {sorted(REQUIRED_FILES)}" + ) + summary = json.loads((args.input / "summary.json").read_text()) + acceptance = summary.get("acceptance") + if not isinstance(acceptance, dict) or acceptance.get("passed") is not True: + raise RuntimeError("Refusing to upload a result that failed acceptance gates") + additional = set(manifest.get("additional_result_files", ())) + missing_additional = sorted( + name for name in additional if not (args.input / name).is_file() + ) + if missing_additional: + raise RuntimeError( + f"Result manifest claims missing additional files: {missing_additional}" + ) import wandb @@ -42,9 +60,10 @@ def main() -> None: artifact = wandb.Artifact( name=f"art-lora-update-replay-{args.mode}", type="lora-update-replay-result", - metadata=json.loads((args.input / "summary.json").read_text()), + metadata=summary, ) - artifact.add_dir(str(args.input)) + for name in sorted(REQUIRED_FILES | additional): + artifact.add_file(str(args.input / name), name=name) logged = run.log_artifact(artifact) logged.wait() artifact_ref = f"{logged.entity}/{logged.project}/{logged.name}" diff --git a/src/art/megatron/service.py b/src/art/megatron/service.py index ed8dd0ae7..86e55ddb2 100644 --- a/src/art/megatron/service.py +++ b/src/art/megatron/service.py @@ -772,8 +772,24 @@ async def _reload_adapter(self, checkpoint_path: str, step: int) -> None: async def _update_in_flight_adapter( self, checkpoint_path: str, step: int ) -> dict[str, float]: + metrics, _receipt = await self._install_in_flight_adapter(checkpoint_path, step) + return metrics + + async def _install_in_flight_adapter( + self, + checkpoint_path: str, + step: int, + *, + verify_receiver_content: bool = False, + ) -> tuple[dict[str, float], dict[str, Any]]: import httpx + from art.megatron.weights.update_replay import ( + inspect_adapter, + validate_committed_policy_version, + validate_receiver_update_receipt, + ) + self._raise_if_child_failed() self.serving_capabilities.require( "in_flight_lora_updates", operation="In-flight LoRA updates" @@ -798,6 +814,38 @@ async def _update_in_flight_adapter( rpc_s = time.perf_counter() - rpc_started json_response = getattr(response, "json", None) response_body = json_response() if callable(json_response) else {} + verification_s = 0.0 + receipt: dict[str, Any] + if verify_receiver_content: + verification_started = time.perf_counter() + async with httpx.AsyncClient() as client: + verification_response = await client.get( + f"{self._vllm_base_url}/art/in_flight_lora_state", + params={"lora_slot": self._in_flight_lora_slot}, + **self._runtime_request_kwargs(), + timeout=60.0, + ) + verification_response.raise_for_status() + verification_body = verification_response.json() + response_body = { + **response_body, + "committed_state": verification_body.get("committed_state"), + "installed_content": verification_body.get("installed_content"), + } + verification_s = time.perf_counter() - verification_started + if verify_receiver_content: + snapshot = inspect_adapter(checkpoint_path) + receipt = validate_receiver_update_receipt( + response_body, + expected_policy_version=step, + expected_safetensors_sha256=snapshot.file_sha256, + ) + else: + validate_committed_policy_version(response_body, expected=step) + receipt = { + "committed_state": dict(response_body["committed_state"]), + "installed_content": None, + } runtime_timing = response_body.get("timing_s", {}) if not isinstance(runtime_timing, dict): runtime_timing = {} @@ -805,8 +853,12 @@ async def _update_in_flight_adapter( self._loaded_adapter_steps.add(step) metrics = { "time/weight_update_service_rpc_s": rpc_s, - "weight_update/policy_version": float(step), + "weight_update/policy_version": float( + receipt["committed_state"]["policy_version"] + ), } + if verify_receiver_content: + metrics["time/weight_update_receiver_verification_s"] = verification_s for phase in ( "begin_update", "load_adapter", @@ -817,7 +869,7 @@ async def _update_in_flight_adapter( value = runtime_timing.get(phase) if isinstance(value, (int, float)): metrics[f"time/weight_update_vllm_{phase}_s"] = float(value) - return metrics + return metrics, receipt async def _load_rollout_lora_for_step( self, checkpoint_path: str, step: int @@ -943,7 +995,7 @@ async def replay_lora_update( output_lora_path: str, policy_version: int, content_version: int, - ) -> dict[str, float]: + ) -> dict[str, Any]: """Publish real resident Megatron tensors, then install that exact output.""" if not self.is_dedicated: raise RuntimeError("LoRA update replay requires dedicated trainer GPUs") @@ -979,10 +1031,15 @@ async def replay_lora_update( publish_metrics = self._weight_update_metrics_from_event(result) if publish_metrics is None: raise RuntimeError("Replay worker produced no LoRA-ready metrics") - install_metrics = await self._update_in_flight_adapter( - output_lora_path, policy_version + install_metrics, receiver_receipt = await self._install_in_flight_adapter( + output_lora_path, + policy_version, + verify_receiver_content=True, ) - return {**publish_metrics, **install_metrics} + return { + "metrics": {**publish_metrics, **install_metrics}, + "receiver": receiver_receipt, + } async def _sleep_runtime(self) -> None: import httpx diff --git a/src/art/megatron/weights/update_replay.py b/src/art/megatron/weights/update_replay.py index cbae0bed2..f901d0ed4 100644 --- a/src/art/megatron/weights/update_replay.py +++ b/src/art/megatron/weights/update_replay.py @@ -156,9 +156,144 @@ def validate_replay_snapshots( def validate_committed_policy_version( response_body: dict[str, Any], *, expected: int ) -> None: - committed = int(response_body.get("policy_version", -1)) + committed_state = response_body.get("committed_state") + if not isinstance(committed_state, dict): + raise RuntimeError("Runtime response has no committed receiver state") + committed = int(committed_state.get("policy_version", -1)) if committed != expected: raise RuntimeError(f"Runtime committed policy {committed}; expected {expected}") + if committed_state.get("update_active") or committed_state.get("admission_blocked"): + raise RuntimeError(f"Runtime did not finish policy commit: {committed_state!r}") + + +def validate_receiver_update_receipt( + response_body: dict[str, Any], + *, + expected_policy_version: int, + expected_safetensors_sha256: str, +) -> dict[str, Any]: + validate_committed_policy_version(response_body, expected=expected_policy_version) + installed_content = response_body.get("installed_content") + if not isinstance(installed_content, dict): + raise RuntimeError("Runtime response has no installed-content receipt") + source = str(installed_content.get("source_safetensors_sha256", "")) + installed = str(installed_content.get("installed_safetensors_sha256", "")) + if ( + source != expected_safetensors_sha256 + or installed != expected_safetensors_sha256 + ): + raise RuntimeError( + "Runtime installed checksum does not match the published adapter: " + f"published={expected_safetensors_sha256}, source={source}, " + f"installed={installed}" + ) + return { + "committed_state": dict(response_body["committed_state"]), + "installed_content": dict(installed_content), + } + + +def validate_bench_request_trace( + request_jsonl: str | Path, + manifest_path: str | Path, + *, + bench_commit: str, +) -> tuple[list[dict[str, Any]], dict[str, Any]]: + if len(bench_commit) != 40 or any( + character not in "0123456789abcdef" for character in bench_commit.lower() + ): + raise ValueError("--bench-commit must be an exact 40-character Git commit") + trace_path = Path(request_jsonl) + trace_payload = trace_path.read_bytes() + manifest = json.loads(Path(manifest_path).read_text()) + requests = manifest.get("requests") + if not isinstance(requests, dict): + raise ValueError("Bench trace manifest has no requests object") + digest = hashlib.sha256(trace_payload).hexdigest() + if digest != requests.get("sha256"): + raise ValueError( + f"Bench request trace digest is {digest}; expected {requests.get('sha256')}" + ) + rows = [ + json.loads(line) for line in trace_payload.decode("utf-8").splitlines() if line + ] + expected_hashes = requests.get("hashes") + selected_ids = manifest.get("selection", {}).get("selected_ids") + if ( + not isinstance(expected_hashes, list) + or not isinstance(selected_ids, list) + or len(rows) != len(expected_hashes) + or len(rows) != len(selected_ids) + ): + raise ValueError("Bench trace row IDs and hashes do not align") + validated: list[dict[str, Any]] = [] + for request, expected, datapoint_id in zip( + rows, expected_hashes, selected_ids, strict=True + ): + if not isinstance(expected, dict): + raise ValueError("Bench request hash entry is not an object") + if str(expected.get("datapoint_id")) != str(datapoint_id): + raise ValueError("Bench request hash IDs do not align with selected IDs") + canonical = json.dumps( + request, ensure_ascii=False, separators=(",", ":"), sort_keys=True + ).encode() + request_sha256 = hashlib.sha256(canonical).hexdigest() + if request_sha256 != expected.get("sha256"): + raise ValueError( + f"Bench request {datapoint_id} hash is {request_sha256}; " + f"expected {expected.get('sha256')}" + ) + validated.append( + { + "id": str(datapoint_id), + "canonical_request_sha256": request_sha256, + "request": request, + } + ) + return validated, { + "bench_commit": bench_commit, + "trace_sha256": digest, + "manifest": manifest, + } + + +def evaluate_replay_acceptance( + measured_metrics: list[dict[str, float]], +) -> dict[str, Any]: + if len(measured_metrics) != 30: + raise RuntimeError( + f"Acceptance requires exactly 30 measured updates; got {len(measured_metrics)}" + ) + combined = [ + row["time/weight_update_trainer_publish_s"] + + row["time/weight_update_service_rpc_s"] + for row in measured_metrics + ] + trainer = [row["time/weight_update_trainer_publish_s"] for row in measured_metrics] + receiver = [row["time/weight_update_vllm_total_s"] for row in measured_metrics] + gates = { + "combined_mean": { + "value_s": statistics.fmean(combined), + "lower_s": 5.001 * 0.85, + "upper_s": 5.001 * 1.15, + }, + "trainer_publish_mean": { + "value_s": statistics.fmean(trainer), + "lower_s": 2.672 * 0.75, + "upper_s": 2.672 * 1.25, + }, + "receiver_install_mean": { + "value_s": statistics.fmean(receiver), + "lower_s": 2.320 * 0.75, + "upper_s": 2.320 * 1.25, + }, + } + failures = [ + name + for name, gate in gates.items() + if not gate["lower_s"] <= gate["value_s"] <= gate["upper_s"] + ] + return {"passed": not failures, "failures": failures, "gates": gates} class FixedLoad: diff --git a/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py b/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py index 6ffc9a7c5..1f3a61f0e 100644 --- a/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py +++ b/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py @@ -173,6 +173,7 @@ async def admit(): blocked = not admission.done() await coordinator.commit_update(slot, 5, new) version, admitted_lora = await admission + committed_state = await coordinator.committed_state(slot) stale_rejected = False try: await coordinator.begin_update(slot, 5) @@ -197,6 +198,7 @@ async def admit(): "lora_path": admitted_lora.lora_path, "recovered_after_cancel": recovered, "stale_rejected": stale_rejected, + "committed_state": committed_state, }, sort_keys=True)) asyncio.run(main()) @@ -211,6 +213,15 @@ async def admit(): "policy_version": 5, "recovered_after_cancel": True, "stale_rejected": True, + "committed_state": { + "active_admissions": 0, + "admission_blocked": False, + "installed_lora_int_id": None, + "installed_lora_path": "new", + "lora_slot": "model:active", + "policy_version": 5, + "update_active": False, + }, } diff --git a/tests/integration/megatron/runtime_isolation/test_service_runtime_boundary.py b/tests/integration/megatron/runtime_isolation/test_service_runtime_boundary.py index c6c32685a..0b2932c90 100644 --- a/tests/integration/megatron/runtime_isolation/test_service_runtime_boundary.py +++ b/tests/integration/megatron/runtime_isolation/test_service_runtime_boundary.py @@ -40,6 +40,19 @@ class _AsyncOkResponse: def raise_for_status(self) -> None: return None + def json(self) -> dict[str, object]: + return { + "committed_state": { + "policy_version": 4, + "update_active": False, + "admission_blocked": False, + }, + "installed_content": { + "source_safetensors_sha256": "abc", + "installed_safetensors_sha256": "abc", + }, + } + class _RecordingAsyncClient: def __init__( @@ -64,6 +77,16 @@ async def post( self._posts.append((url, json if json is not None else params, timeout)) return _AsyncOkResponse() + async def get( + self, + url: str, + *, + params: dict[str, object] | None = None, + timeout: float, + ) -> _AsyncOkResponse: + self._posts.append((url, params, timeout)) + return _AsyncOkResponse() + def test_megatron_default_lora_adapter_config_uses_model_lora_config( tmp_path: Path, @@ -103,6 +126,10 @@ async def test_megatron_in_flight_eval_uses_immutable_adapter_slot( service._vllm_runtime.port = 8123 posts: list[tuple[str, dict[str, object] | None, float]] = [] monkeypatch.setattr(httpx, "AsyncClient", lambda: _RecordingAsyncClient(posts)) + monkeypatch.setattr( + "art.megatron.weights.update_replay.inspect_adapter", + lambda _path: SimpleNamespace(file_sha256="abc"), + ) checkpoint_path = str(tmp_path / "checkpoints" / "4") assert ( @@ -152,6 +179,9 @@ async def test_external_in_flight_update_maps_checkpoint_path( monkeypatch: pytest.MonkeyPatch, ) -> None: local_root = str(tmp_path / "local") + monkeypatch.setattr( + MegatronService, "_validate_megatron_dependencies", lambda _self: None + ) service = MegatronService( model_name="test-model", base_model="Qwen/Qwen3-0.6B", diff --git a/tests/unit/test_weight_update_replay.py b/tests/unit/test_weight_update_replay.py index 4b9fb763c..c9063e6a3 100644 --- a/tests/unit/test_weight_update_replay.py +++ b/tests/unit/test_weight_update_replay.py @@ -1,3 +1,4 @@ +import hashlib import json from pathlib import Path @@ -13,8 +14,11 @@ from art.megatron.service import MegatronService from art.megatron.weights.lora_publish import LoraPublishMetrics from art.megatron.weights.update_replay import ( + evaluate_replay_acceptance, inspect_adapter, + validate_bench_request_trace, validate_committed_policy_version, + validate_receiver_update_receipt, validate_replay_snapshots, ) @@ -150,10 +154,60 @@ def test_lora_publish_job_roundtrips_through_worker_protocol() -> None: def test_policy_version_must_match_committed_runtime_version() -> None: - validate_committed_policy_version({"policy_version": 12}, expected=12) + validate_committed_policy_version( + { + "policy_version": 999, + "committed_state": { + "policy_version": 12, + "update_active": False, + "admission_blocked": False, + }, + }, + expected=12, + ) with pytest.raises(RuntimeError, match="committed policy 11; expected 12"): - validate_committed_policy_version({"policy_version": 11}, expected=12) + validate_committed_policy_version( + { + "committed_state": { + "policy_version": 11, + "update_active": False, + "admission_blocked": False, + } + }, + expected=12, + ) + + +def test_receiver_receipt_requires_committed_state_and_installed_checksum() -> None: + receipt = { + "committed_state": { + "policy_version": 7, + "update_active": False, + "admission_blocked": False, + }, + "installed_content": { + "source_safetensors_sha256": "abc", + "installed_safetensors_sha256": "abc", + }, + } + + assert ( + validate_receiver_update_receipt( + receipt, + expected_policy_version=7, + expected_safetensors_sha256="abc", + )["committed_state"]["policy_version"] + == 7 + ) + + receipt["installed_content"]["installed_safetensors_sha256"] = "stale" + with pytest.raises(RuntimeError, match="installed checksum"): + validate_receiver_update_receipt( + receipt, + expected_policy_version=7, + expected_safetensors_sha256="abc", + ) def test_snapshot_config_is_machine_readable(tmp_path: Path) -> None: @@ -164,23 +218,67 @@ def test_snapshot_config_is_machine_readable(tmp_path: Path) -> None: assert config["art_lora_format"] == "vllm" -def test_h200_replay_launcher_preserves_evidence_contract() -> None: +def test_bench_trace_requires_exact_external_provenance(tmp_path: Path) -> None: + request = { + "model": "Qwen/Qwen3.6-35B-A3B", + "messages": [{"role": "user", "content": "hello"}], + } + canonical = json.dumps( + request, ensure_ascii=False, separators=(",", ":"), sort_keys=True + ).encode() + trace = tmp_path / "request-trace.jsonl" + trace.write_bytes(canonical + b"\n") + digest = hashlib.sha256(trace.read_bytes()).hexdigest() + request_digest = hashlib.sha256(canonical).hexdigest() + manifest = tmp_path / "manifest.json" + manifest.write_text( + json.dumps( + { + "selection": {"selected_ids": ["row-1"]}, + "requests": { + "sha256": digest, + "hashes": [{"datapoint_id": "row-1", "sha256": request_digest}], + }, + } + ) + ) + + rows, provenance = validate_bench_request_trace( + trace, manifest, bench_commit="a" * 40 + ) + + assert rows[0]["id"] == "row-1" + assert provenance["trace_sha256"] == digest + assert provenance["bench_commit"] == "a" * 40 + + +def test_replay_acceptance_enforces_combined_and_component_gates() -> None: + passing = [ + { + "time/weight_update_trainer_publish_s": 2.672, + "time/weight_update_service_rpc_s": 2.329, + "time/weight_update_vllm_total_s": 2.320, + } + for _ in range(30) + ] + assert evaluate_replay_acceptance(passing)["passed"] is True + + failing = [dict(row) for row in passing] + failing[0]["time/weight_update_vllm_total_s"] = 100.0 + assert evaluate_replay_acceptance(failing)["passed"] is False + + +def test_replay_uploader_matches_bench_required_files() -> None: root = Path(__file__).parents[2] - launcher = ( - root / "scripts/benchmarks/lora-update-replay-h200.sky.yaml" - ).read_text() service = (root / "scripts/benchmarks/service_lora_update_replay.py").read_text() uploader = (root / "scripts/benchmarks/upload_lora_update_replay.py").read_text() - assert "accelerators: H200:4" in launcher - assert "name: vivek-dev-api-keys" in launcher - assert "--mode fixed-load" in launcher - assert "--mode idle" in launcher assert 'max_model_len": 32769' in service - assert "control_seconds" in service + assert "validate_bench_request_trace" in service + assert "--bench-commit" in service for filename in ( "manifest.json", - "request-trace.json", + "request-trace.jsonl", "tensor-manifest.json", "samples.jsonl", "summary.json", diff --git a/vllm_runtime/src/art_vllm_runtime/dedicated_server.py b/vllm_runtime/src/art_vllm_runtime/dedicated_server.py index 1cf16b943..299d0c72e 100644 --- a/vllm_runtime/src/art_vllm_runtime/dedicated_server.py +++ b/vllm_runtime/src/art_vllm_runtime/dedicated_server.py @@ -2,6 +2,7 @@ import argparse import asyncio +import hashlib from http import HTTPStatus import json import os @@ -48,6 +49,15 @@ class _InFlightLoraUpdateRequest(BaseModel): is_3d_lora_weight: bool = False +def _adapter_safetensors_sha256(lora_path: str) -> str: + path = os.path.join(lora_path, "adapter_model.safetensors") + digest = hashlib.sha256() + with open(path, "rb") as handle: + for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser(description="ART dedicated vLLM server") parser.add_argument("--model", required=True, help="Base model name or path") @@ -264,6 +274,7 @@ async def in_flight_lora_update( models.lora_requests[lora_slot], ) commit_s = time.perf_counter() - commit_started + committed_state = await coordinator.committed_state(lora_slot) timing_s = { "begin_update": begin_s, "load_adapter": load_s, @@ -291,12 +302,41 @@ async def in_flight_lora_update( "status": "updated", "model_name": public_model_name, "lora_slot": lora_slot, - "policy_version": policy_version, + "policy_version": committed_state["policy_version"], + "committed_state": committed_state, "waiting_cache_salt": waiting_cache_salt, "timing_s": timing_s, } ) + @router.get("/art/in_flight_lora_state") + async def in_flight_lora_state( + raw_request: Request, + lora_slot: str = Query(min_length=1), + ) -> JSONResponse: + models = raw_request.app.state.openai_serving_models + engine_client = engine(raw_request) + coordinator = lora_update_coordinator(models, engine_client) + committed_state = await coordinator.committed_state(lora_slot) + installed_path = committed_state.get("installed_lora_path") + installed_sha256 = ( + _adapter_safetensors_sha256(str(installed_path)) + if installed_path is not None + else None + ) + return JSONResponse( + content={ + "committed_state": committed_state, + "installed_content": { + "source_safetensors_sha256": installed_sha256, + "installed_safetensors_sha256": installed_sha256, + "verification_scope": ( + "runtime-selected safetensors after successful GPU load" + ), + }, + } + ) + app.include_router(router) return app diff --git a/vllm_runtime/src/art_vllm_runtime/policy_spans.py b/vllm_runtime/src/art_vllm_runtime/policy_spans.py index 4a1459018..637e9cea9 100644 --- a/vllm_runtime/src/art_vllm_runtime/policy_spans.py +++ b/vllm_runtime/src/art_vllm_runtime/policy_spans.py @@ -162,6 +162,28 @@ async def fail_update(self, lora_slot: str) -> None: state.blocked = True state.condition.notify_all() + async def committed_state(self, lora_slot: str) -> dict[str, Any]: + state = self._state(lora_slot) + async with state.condition: + lora_request = state.lora_request + return { + "lora_slot": lora_slot, + "policy_version": state.policy_version, + "update_active": state.update_active, + "admission_blocked": state.blocked, + "active_admissions": state.active_admissions, + "installed_lora_path": ( + str(lora_request.lora_path) if lora_request is not None else None + ), + "installed_lora_int_id": ( + int(lora_int_id) + if lora_request is not None + and (lora_int_id := getattr(lora_request, "lora_int_id", None)) + is not None + else None + ), + } + def lora_update_coordinator(models: Any, engine_client: Any) -> LoraUpdateCoordinator: coordinator = getattr(models, _LORA_UPDATE_COORDINATOR_FIELD, None) From 967e5c684312a70f853d8112f808dd41ab4c8fe4 Mon Sep 17 00:00:00 2001 From: Vivek Kalyan Date: Thu, 23 Jul 2026 19:48:47 -0700 Subject: [PATCH 9/9] chore: Remove unpublished replay experiments --- scripts/benchmarks/replay_lora_updates.py | 6 - .../benchmarks/service_lora_update_replay.py | 698 ------------------ .../benchmarks/upload_lora_update_replay.py | 78 -- src/art/megatron/runtime/jobs.py | 11 - src/art/megatron/service.py | 198 +---- src/art/megatron/train.py | 58 +- src/art/megatron/weights/lora_publish.py | 55 +- src/art/megatron/weights/update_replay.py | 540 -------------- .../test_runtime_project_isolation.py | 18 - .../test_service_runtime_boundary.py | 30 - tests/unit/test_weight_update_replay.py | 288 -------- .../src/art_vllm_runtime/dedicated_server.py | 68 +- vllm_runtime/src/art_vllm_runtime/metrics.py | 22 - .../src/art_vllm_runtime/policy_spans.py | 43 +- 14 files changed, 21 insertions(+), 2092 deletions(-) delete mode 100644 scripts/benchmarks/replay_lora_updates.py delete mode 100644 scripts/benchmarks/service_lora_update_replay.py delete mode 100644 scripts/benchmarks/upload_lora_update_replay.py delete mode 100644 src/art/megatron/weights/update_replay.py delete mode 100644 tests/unit/test_weight_update_replay.py diff --git a/scripts/benchmarks/replay_lora_updates.py b/scripts/benchmarks/replay_lora_updates.py deleted file mode 100644 index f3f0b3abc..000000000 --- a/scripts/benchmarks/replay_lora_updates.py +++ /dev/null @@ -1,6 +0,0 @@ -#!/usr/bin/env python3 - -from art.megatron.weights.update_replay import main - -if __name__ == "__main__": - main() diff --git a/scripts/benchmarks/service_lora_update_replay.py b/scripts/benchmarks/service_lora_update_replay.py deleted file mode 100644 index ba2e6a8fe..000000000 --- a/scripts/benchmarks/service_lora_update_replay.py +++ /dev/null @@ -1,698 +0,0 @@ -#!/usr/bin/env python3 - -from __future__ import annotations - -import argparse -import asyncio -from dataclasses import asdict -import hashlib -import json -from pathlib import Path -import random -import resource -import shutil -import statistics -import subprocess -import time -from typing import Any - -import httpx - -import art -from art.megatron.service import MegatronService -from art.megatron.weights.update_replay import ( - evaluate_replay_acceptance, - inspect_adapter, - validate_bench_request_trace, -) - -BASE_MODEL = "Qwen/Qwen3.6-35B-A3B" - - -def _canonical_sha256(value: Any) -> str: - payload = json.dumps( - value, ensure_ascii=False, separators=(",", ":"), sort_keys=True - ).encode() - return hashlib.sha256(payload).hexdigest() - - -def _effective_requests( - request_jsonl: Path, trace_manifest: Path, *, bench_commit: str -) -> tuple[list[dict[str, Any]], dict[str, Any]]: - rows, provenance = validate_bench_request_trace( - request_jsonl, trace_manifest, bench_commit=bench_commit - ) - generation = provenance["manifest"].get("generation", {}) - expected = { - "n": 8, - "temperature": 0.8, - "max_tokens": 256, - "logprobs": True, - "enable_thinking": False, - "load_concurrency": 8, - } - if {key: generation.get(key) for key in expected} != expected: - raise ValueError(f"Unexpected Bench generation contract: {generation!r}") - for row in rows: - effective = dict(row["request"]) - effective.update( - { - "model": "bonnie-replay:active", - "stream": True, - "stream_options": {"include_usage": True}, - } - ) - row["effective_request"] = effective - row["effective_request_sha256"] = _canonical_sha256(effective) - return rows, provenance - - -async def _runtime_metrics(client: httpx.AsyncClient) -> dict[str, float]: - response = await client.get("/art/metrics") - response.raise_for_status() - return { - key: float(value) - for key, value in response.json().get("metrics", {}).items() - if isinstance(value, (int, float)) - } - - -async def _load_worker( - *, - client: httpx.AsyncClient, - requests: list[dict[str, Any]], - worker: int, - stop: asyncio.Event, - samples: list[dict[str, Any]], - token_events: list[dict[str, Any]], - phase: list[str], -) -> None: - sequence = 0 - index = worker - while not stop.is_set(): - row = requests[index % len(requests)] - index += 8 - request_id = f"load-{worker}-{sequence:06d}" - sequence += 1 - started = time.perf_counter() - phase_at_start = phase[0] - error = None - prompt_tokens = completion_tokens = 0 - ttft_s: float | None = None - last_token_time_by_choice: dict[int, float] = {} - inter_token_intervals_s: list[float] = [] - spans_by_choice: dict[int, list[dict[str, Any]]] = {} - try: - async with client.stream( - "POST", - "/v1/chat/completions", - json=row["effective_request"], - headers={"X-Request-Id": request_id}, - ) as response: - response.raise_for_status() - async for line in response.aiter_lines(): - if not line.startswith("data: ") or line == "data: [DONE]": - continue - chunk = json.loads(line.removeprefix("data: ")) - usage = chunk.get("usage") or {} - prompt_tokens = int(usage.get("prompt_tokens", prompt_tokens)) - completion_tokens = int( - usage.get("completion_tokens", completion_tokens) - ) - for choice in chunk.get("choices") or []: - choice_index = int(choice.get("index", 0)) - spans = choice.get("policy_token_spans") or [] - spans_by_choice.setdefault(choice_index, []).extend( - dict(span) for span in spans - ) - logprobs = choice.get("logprobs") or {} - emitted = len(logprobs.get("content") or []) - if emitted == 0: - delta = choice.get("delta") or {} - emitted = int( - bool(delta.get("content") or delta.get("tool_calls")) - ) - if emitted: - now = time.perf_counter() - if ttft_s is None: - ttft_s = now - started - previous = last_token_time_by_choice.get(choice_index) - if previous is not None: - inter_token_intervals_s.append(now - previous) - if emitted > 1: - inter_token_intervals_s.extend([0.0] * (emitted - 1)) - last_token_time_by_choice[choice_index] = now - token_events.append( - { - "request_id": request_id, - "monotonic_s": now, - "tokens": emitted, - } - ) - except Exception as exc: - error = f"{type(exc).__name__}: {exc}" - ended = time.perf_counter() - samples.append( - { - "request_id": request_id, - "worker": worker, - "phase_at_start": phase_at_start, - "datapoint_id": row["id"], - "canonical_request_sha256": row["canonical_request_sha256"], - "effective_request_sha256": row["effective_request_sha256"], - "started_monotonic_s": started, - "ended_monotonic_s": ended, - "latency_s": ended - started, - "ttft_s": ttft_s, - "inter_token_intervals_s": inter_token_intervals_s, - "prompt_tokens": prompt_tokens, - "completion_tokens": completion_tokens, - "policy_token_spans_by_choice": spans_by_choice, - "error": error, - } - ) - - -async def _record_control( - *, - client: httpx.AsyncClient, - name: str, - duration_s: float, - phase: list[str], -) -> dict[str, Any]: - phase[0] = name - before = await _runtime_metrics(client) - started = time.perf_counter() - await asyncio.sleep(duration_s) - ended = time.perf_counter() - after = await _runtime_metrics(client) - elapsed = ended - started - prompt_tokens = after["prompt_tokens_total"] - before["prompt_tokens_total"] - completion_tokens = ( - after["generation_tokens_total"] - before["generation_tokens_total"] - ) - return { - "name": name, - "started_monotonic_s": started, - "ended_monotonic_s": ended, - "duration_s": elapsed, - "prompt_tokens": prompt_tokens, - "completion_tokens": completion_tokens, - "prompt_tokens_per_s": prompt_tokens / elapsed, - "completion_tokens_per_s": completion_tokens / elapsed, - } - - -async def _fixed_input_fingerprint( - client: httpx.AsyncClient, *, policy_version: int -) -> dict[str, Any]: - request = { - "model": "bonnie-replay:active", - "messages": [{"role": "user", "content": "Reply with exactly: OK"}], - "temperature": 0, - "max_tokens": 1, - "n": 1, - "logprobs": True, - "top_logprobs": 5, - "chat_template_kwargs": {"enable_thinking": False}, - } - fingerprints: list[str] = [] - for _ in range(2): - response = await client.post("/v1/chat/completions", json=request) - response.raise_for_status() - body = response.json() - choice = body["choices"][0] - spans = choice.get("policy_token_spans") or [] - versions = {int(span["policy_version"]) for span in spans} - if versions != {policy_version}: - raise RuntimeError( - f"Fixed-input request used policies {versions}; expected {policy_version}" - ) - receipt = { - "message": choice["message"], - "logprobs": choice.get("logprobs"), - } - fingerprints.append(_canonical_sha256(receipt)) - if len(set(fingerprints)) != 1: - raise RuntimeError("Fixed-input receiver logit parity failed") - return { - "method": "two deterministic receiver requests after commit", - "sha256": fingerprints[0], - "policy_version": policy_version, - "policy_token_spans": spans, - } - - -def _gpu_memory_snapshot() -> list[dict[str, int | str]]: - output = subprocess.check_output( - [ - "nvidia-smi", - "--query-gpu=index,name,memory.used,memory.total", - "--format=csv,noheader,nounits", - ], - text=True, - ) - rows: list[dict[str, int | str]] = [] - for line in output.splitlines(): - raw_index, name, raw_used, raw_total = ( - part.strip() for part in line.split(",", maxsplit=3) - ) - rows.append( - { - "gpu_index": int(raw_index), - "name": name, - "used_mib": int(raw_used), - "total_mib": int(raw_total), - } - ) - return rows - - -def _cpu_memory_snapshot() -> dict[str, int]: - meminfo: dict[str, int] = {} - for line in Path("/proc/meminfo").read_text().splitlines(): - key, raw = line.split(":", maxsplit=1) - meminfo[key] = int(raw.strip().split()[0]) - return { - "system_used_kib": meminfo["MemTotal"] - meminfo["MemAvailable"], - "runner_peak_rss_kib": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss, - } - - -def _bootstrap_ci(values: list[float]) -> list[float]: - rng = random.Random(20260723) - means = [statistics.fmean(rng.choices(values, k=len(values))) for _ in range(2000)] - means.sort() - return [means[50], means[1949]] - - -def _attach_policy_ages( - samples: list[dict[str, Any]], updates: list[dict[str, Any]] -) -> None: - commits = [ - ( - float(update["committed_monotonic_s"]), - int(update["receiver"]["committed_state"]["policy_version"]), - ) - for update in updates - ] - for sample in samples: - current = max( - ( - version - for committed_at, version in commits - if committed_at <= float(sample["ended_monotonic_s"]) - ), - default=0, - ) - versions = [ - int(span["policy_version"]) - for spans in sample["policy_token_spans_by_choice"].values() - for span in spans - ] - sample["receiver_policy_version_at_end"] = current - sample["policy_age_min"] = current - max(versions) if versions else None - sample["policy_age_max"] = current - min(versions) if versions else None - - -async def run(args: argparse.Namespace) -> None: - requests, trace_provenance = _effective_requests( - args.request_jsonl, args.trace_manifest, bench_commit=args.bench_commit - ) - output = args.output.resolve() - output.mkdir(parents=True, exist_ok=True) - shutil.copyfile(args.request_jsonl, output / "request-trace.jsonl") - art.init_megatron_runtime_config( - topology=art.MegatronTopologyConfig(tp=1, pp=1, cp=2, ep=2, etp=1), - packed_sequence_length=122880, - streaming_weight_offload=True, - ) - service = MegatronService( - model_name="bonnie-replay", - base_model=BASE_MODEL, - config={ - "trainer_gpu_ids": [0, 1], - "inference_gpu_ids": [2, 3], - "rollout_weights_mode": "lora", - "rollout_weight_update_mode": "in_flight_lora", - "allow_unvalidated_arch": True, - "chat_template_kwargs": {"enable_thinking": False}, - "lora_config": {"rank": 8}, - "engine_args": { - "tensor_parallel_size": 2, - "max_model_len": 32769, - "max_num_batched_tokens": 131072, - "max_num_seqs": 256, - "enable_prefix_caching": True, - "gpu_memory_utilization": 0.9, - }, - }, - output_dir=str(output / "service"), - ) - seed = Path(service._resolve_active_lora_path()) - await service.start_openai_server(None) - samples: list[dict[str, Any]] = [] - token_events: list[dict[str, Any]] = [] - controls: list[dict[str, Any]] = [] - updates: list[dict[str, Any]] = [] - phase = ["idle"] - stop = asyncio.Event() - async with httpx.AsyncClient( - base_url=service._vllm_base_url, - headers=service._runtime_headers(), - timeout=300.0, - ) as client: - runtime_before = await _runtime_metrics(client) - if int(runtime_before.get("world_size", 0)) != 2: - raise RuntimeError(f"Expected vLLM TP2 runtime: {runtime_before!r}") - load_tasks = ( - [ - asyncio.create_task( - _load_worker( - client=client, - requests=requests, - worker=worker, - stop=stop, - samples=samples, - token_events=token_events, - phase=phase, - ) - ) - for worker in range(8) - ] - if args.mode == "fixed-load" - else [] - ) - hashes: set[str] = set() - manifest_sha256: str | None = None - try: - if load_tasks: - controls.append( - await _record_control( - client=client, - name="control-before", - duration_s=args.control_seconds, - phase=phase, - ) - ) - phase[0] = "updates" - for index in range(args.warmups + args.updates): - policy_version = index + 1 - checkpoint = output / "checkpoints" / f"{policy_version:04d}" - memory_before = { - "gpu": _gpu_memory_snapshot(), - "cpu": _cpu_memory_snapshot(), - } - update_runtime_before = ( - await _runtime_metrics(client) - if index >= args.warmups and load_tasks - else None - ) - update_started = time.perf_counter() - replay = await service.replay_lora_update( - source_lora_path=str(seed), - output_lora_path=str(checkpoint), - policy_version=policy_version, - content_version=index, - ) - committed_at = time.perf_counter() - update_runtime_after = ( - await _runtime_metrics(client) - if update_runtime_before is not None - else None - ) - metrics = replay["metrics"] - receiver = replay["receiver"] - expected_topology = { - "topology/trainer_tp": 1.0, - "topology/trainer_pp": 1.0, - "topology/trainer_cp": 2.0, - "topology/trainer_ep": 2.0, - "topology/trainer_etp": 1.0, - } - if { - key: metrics.get(key) for key in expected_topology - } != expected_topology: - raise RuntimeError(f"Unexpected trainer topology: {metrics!r}") - if len(memory_before["gpu"]) != 4 or any( - "H200" not in str(row["name"]) for row in memory_before["gpu"] - ): - raise RuntimeError(f"Replay requires four H200s: {memory_before!r}") - snapshot = inspect_adapter(checkpoint) - if snapshot.file_sha256 in hashes: - raise RuntimeError("Replay produced duplicate update contents") - if manifest_sha256 is None: - manifest_sha256 = snapshot.manifest_sha256 - elif snapshot.manifest_sha256 != manifest_sha256: - raise RuntimeError("Replay changed the serving tensor manifest") - hashes.add(snapshot.file_sha256) - parity = await _fixed_input_fingerprint( - client, policy_version=policy_version - ) - updates.append( - { - "index": index, - "measured": index >= args.warmups, - "requested_policy_version": policy_version, - "checkpoint": str(checkpoint), - "sha256": snapshot.file_sha256, - "manifest_sha256": snapshot.manifest_sha256, - "metrics": metrics, - "receiver": receiver, - "fixed_input_logit_parity": parity, - "update_started_monotonic_s": update_started, - "committed_monotonic_s": committed_at, - "update_wall_s": committed_at - update_started, - "runtime_counter_delta": ( - { - "prompt_tokens": ( - update_runtime_after["prompt_tokens_total"] - - update_runtime_before["prompt_tokens_total"] - ), - "completion_tokens": ( - update_runtime_after["generation_tokens_total"] - - update_runtime_before["generation_tokens_total"] - ), - } - if update_runtime_before is not None - and update_runtime_after is not None - else None - ), - "memory_before": memory_before, - "memory_after": { - "gpu": _gpu_memory_snapshot(), - "cpu": _cpu_memory_snapshot(), - }, - } - ) - if len(updates) > 2: - shutil.rmtree(Path(updates[-3]["checkpoint"])) - if load_tasks: - controls.append( - await _record_control( - client=client, - name="control-after", - duration_s=args.control_seconds, - phase=phase, - ) - ) - before_rate = float(controls[0]["completion_tokens_per_s"]) - after_rate = float(controls[-1]["completion_tokens_per_s"]) - drift = ( - abs(after_rate - before_rate) / before_rate - if before_rate > 0 - else float("inf") - ) - if drift > 0.05: - controls.append( - await _record_control( - client=client, - name="control-repeat", - duration_s=args.control_seconds, - phase=phase, - ) - ) - finally: - stop.set() - await asyncio.gather(*load_tasks) - runtime_after = await _runtime_metrics(client) - await service.aclose() - - measured = [row for row in updates if row["measured"]] - acceptance = evaluate_replay_acceptance([row["metrics"] for row in measured]) - totals = [ - float(row["metrics"]["time/weight_update_trainer_publish_s"]) - + float(row["metrics"]["time/weight_update_service_rpc_s"]) - for row in measured - ] - _attach_policy_ages(samples, updates) - prompt_tokens = sum( - float((row["runtime_counter_delta"] or {})["prompt_tokens"]) - for row in measured - if row["runtime_counter_delta"] is not None - ) - completion_tokens = sum( - float((row["runtime_counter_delta"] or {})["completion_tokens"]) - for row in measured - if row["runtime_counter_delta"] is not None - ) - measured_duration = sum(float(row["update_wall_s"]) for row in measured) - control_completion_rate = ( - statistics.fmean(row["completion_tokens_per_s"] for row in controls) - if controls - else None - ) - ages = [ - int(sample["policy_age_max"]) - for sample in samples - if sample["policy_age_max"] is not None - ] - gpu_snapshots = [ - gpu - for update in updates - for boundary in ("memory_before", "memory_after") - for gpu in update[boundary]["gpu"] - ] - cpu_snapshots = [ - update[boundary]["cpu"] - for update in updates - for boundary in ("memory_before", "memory_after") - ] - summary: dict[str, Any] = { - "acceptance": acceptance, - "update_total_mean_s": statistics.fmean(totals), - "update_total_p50_s": statistics.median(totals), - "update_total_p95_s": sorted(totals)[round(0.95 * (len(totals) - 1))], - "update_total_stddev_s": statistics.stdev(totals), - "update_total_mean_95ci_s": _bootstrap_ci(totals), - "measured_lane_duration_s": measured_duration, - "prompt_tokens": prompt_tokens, - "completion_tokens": completion_tokens, - "prompt_tokens_per_s": prompt_tokens / measured_duration, - "completion_tokens_per_s": completion_tokens / measured_duration, - "lost_tokens_per_update": ( - (control_completion_rate * measured_duration - completion_tokens) - / len(measured) - if control_completion_rate is not None - else None - ), - "request_errors": sum(row["error"] is not None for row in samples), - "policy_age_p50": statistics.median(ages) if ages else None, - "policy_age_p95": ( - sorted(ages)[round(0.95 * (len(ages) - 1))] if ages else None - ), - "peak_gpu_used_mib": max(int(row["used_mib"]) for row in gpu_snapshots), - "peak_system_used_kib": max(row["system_used_kib"] for row in cpu_snapshots), - "peak_runner_rss_kib": max(row["runner_peak_rss_kib"] for row in cpu_snapshots), - "controls": controls, - } - ttfts = [ - float(row["ttft_s"]) - for row in samples - if row["ttft_s"] is not None and row["error"] is None - ] - itls = [ - float(interval) - for row in samples - if row["error"] is None - for interval in row["inter_token_intervals_s"] - ] - summary.update( - { - "ttft_p50_s": statistics.median(ttfts) if ttfts else None, - "ttft_p95_s": ( - sorted(ttfts)[round(0.95 * (len(ttfts) - 1))] if ttfts else None - ), - "inter_token_latency_p50_s": statistics.median(itls) if itls else None, - "inter_token_latency_p95_s": ( - sorted(itls)[round(0.95 * (len(itls) - 1))] if itls else None - ), - } - ) - (output / "updates.jsonl").write_text( - "".join(json.dumps(row, sort_keys=True) + "\n" for row in updates) - ) - (output / "samples.jsonl").write_text( - "".join(json.dumps(row, sort_keys=True) + "\n" for row in samples) - ) - (output / "token-events.jsonl").write_text( - "".join(json.dumps(row, sort_keys=True) + "\n" for row in token_events) - ) - (output / "summary.json").write_text(json.dumps(summary, indent=2)) - (output / "runtime-before.json").write_text(json.dumps(runtime_before, indent=2)) - (output / "runtime-after.json").write_text(json.dumps(runtime_after, indent=2)) - reference_snapshot = inspect_adapter(Path(updates[-1]["checkpoint"])) - (output / "tensor-manifest.json").write_text( - json.dumps( - { - "schema_version": 1, - "manifest_sha256": reference_snapshot.manifest_sha256, - "rank": reference_snapshot.rank, - "base_model": reference_snapshot.base_model, - "logical_bytes": reference_snapshot.logical_bytes, - "transported_bytes": reference_snapshot.transported_bytes, - "tensor_count": reference_snapshot.tensor_count, - "tensors": [asdict(tensor) for tensor in reference_snapshot.tensors], - }, - indent=2, - ) - ) - required_result_files = [ - "manifest.json", - "request-trace.jsonl", - "tensor-manifest.json", - "samples.jsonl", - "summary.json", - "stdout.log", - "stderr.log", - ] - result_manifest = { - "schema_version": 1, - "mode": args.mode, - "art_commit": subprocess.check_output( - ["git", "rev-parse", "HEAD"], text=True - ).strip(), - "bench_trace": trace_provenance, - "effective_request_transform": { - "model": "bonnie-replay:active", - "stream": True, - "stream_options": {"include_usage": True}, - }, - "topology": { - "trainer_gpu_ids": [0, 1], - "inference_gpu_ids": [2, 3], - "inference_tp": 2, - "trainer": {"tp": 1, "pp": 1, "cp": 2, "ep": 2, "etp": 1}, - }, - "required_result_files": required_result_files, - "additional_result_files": [ - "updates.jsonl", - "token-events.jsonl", - "runtime-before.json", - "runtime-after.json", - ], - } - (output / "manifest.json").write_text(json.dumps(result_manifest, indent=2)) - if summary["request_errors"]: - raise RuntimeError( - f"Evidence run observed {summary['request_errors']} request errors" - ) - if not acceptance["passed"]: - raise RuntimeError(f"Replay acceptance gates failed: {acceptance['failures']}") - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--request-jsonl", type=Path, required=True) - parser.add_argument("--trace-manifest", type=Path, required=True) - parser.add_argument("--bench-commit", required=True) - parser.add_argument("--mode", choices=("idle", "fixed-load"), required=True) - parser.add_argument("--warmups", type=int, default=5) - parser.add_argument("--updates", type=int, default=30) - parser.add_argument("--control-seconds", type=float, default=60.0) - parser.add_argument("--output", type=Path, required=True) - asyncio.run(run(parser.parse_args())) - - -if __name__ == "__main__": - main() diff --git a/scripts/benchmarks/upload_lora_update_replay.py b/scripts/benchmarks/upload_lora_update_replay.py deleted file mode 100644 index 33bd2db0f..000000000 --- a/scripts/benchmarks/upload_lora_update_replay.py +++ /dev/null @@ -1,78 +0,0 @@ -#!/usr/bin/env python3 - -from __future__ import annotations - -import argparse -import json -import os -from pathlib import Path - -REQUIRED_FILES = { - "manifest.json", - "request-trace.jsonl", - "tensor-manifest.json", - "samples.jsonl", - "summary.json", - "stdout.log", - "stderr.log", -} - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--input", type=Path, required=True) - parser.add_argument("--mode", choices=("idle", "fixed-load"), required=True) - args = parser.parse_args() - - missing = sorted( - name for name in REQUIRED_FILES if not (args.input / name).is_file() - ) - if missing: - raise RuntimeError(f"Refusing to upload incomplete evidence: {missing}") - manifest = json.loads((args.input / "manifest.json").read_text()) - claimed = set(manifest.get("required_result_files", ())) - if claimed != REQUIRED_FILES: - raise RuntimeError( - f"Result manifest claims {sorted(claimed)}; expected {sorted(REQUIRED_FILES)}" - ) - summary = json.loads((args.input / "summary.json").read_text()) - acceptance = summary.get("acceptance") - if not isinstance(acceptance, dict) or acceptance.get("passed") is not True: - raise RuntimeError("Refusing to upload a result that failed acceptance gates") - additional = set(manifest.get("additional_result_files", ())) - missing_additional = sorted( - name for name in additional if not (args.input / name).is_file() - ) - if missing_additional: - raise RuntimeError( - f"Result manifest claims missing additional files: {missing_additional}" - ) - - import wandb - - run = wandb.init( - entity=os.environ.get("WANDB_ENTITY", "wb-training"), - project=os.environ.get("WANDB_PROJECT", "bench"), - job_type="lora-update-replay", - config={"mode": args.mode, "artifact_schema_version": 1}, - ) - assert run is not None - artifact = wandb.Artifact( - name=f"art-lora-update-replay-{args.mode}", - type="lora-update-replay-result", - metadata=summary, - ) - for name in sorted(REQUIRED_FILES | additional): - artifact.add_file(str(args.input / name), name=name) - logged = run.log_artifact(artifact) - logged.wait() - artifact_ref = f"{logged.entity}/{logged.project}/{logged.name}" - run.summary["result_artifact"] = artifact_ref - run.summary["result_artifact_digest"] = logged.digest - run.finish() - print(f"WANDB_ARTIFACT={artifact_ref}") - print(f"WANDB_ARTIFACT_DIGEST={logged.digest}") - - -if __name__ == "__main__": - main() diff --git a/src/art/megatron/runtime/jobs.py b/src/art/megatron/runtime/jobs.py index 022be1a72..4285c88bc 100644 --- a/src/art/megatron/runtime/jobs.py +++ b/src/art/megatron/runtime/jobs.py @@ -59,16 +59,6 @@ class MegatronSyncJob(BaseModel): log_path: str = DEFAULT_TRAINING_LOG_PATH -class MegatronLoraPublishJob(BaseModel): - kind: Literal["publish_lora"] = "publish_lora" - step: int = Field(ge=0) - source_lora_path: str - output_lora_path: str - content_version: int = Field(ge=0) - allow_unvalidated_arch: bool = False - log_path: str = DEFAULT_TRAINING_LOG_PATH - - class MegatronOptimizerSaveJob(BaseModel): kind: Literal["save_optimizer"] = "save_optimizer" step: int = Field(ge=0) @@ -99,7 +89,6 @@ class MegatronSFTTrainingJob(BaseModel): MegatronTrainingJob | MegatronMergedTrainingJob | MegatronSyncJob - | MegatronLoraPublishJob | MegatronOptimizerSaveJob | MegatronSFTTrainingJob, Field(discriminator="kind"), diff --git a/src/art/megatron/service.py b/src/art/megatron/service.py index 86e55ddb2..fe04050f9 100644 --- a/src/art/megatron/service.py +++ b/src/art/megatron/service.py @@ -8,7 +8,6 @@ import socket import subprocess import sys -import time from typing import Any, AsyncIterator, Literal, TypedDict, cast from urllib.parse import urlparse import uuid @@ -73,7 +72,6 @@ from .runtime.jobs import ( LORA_READY_EVENT, OPTIMIZER_READY_EVENT, - MegatronLoraPublishJob, MegatronMergedTrainingJob, MegatronOptimizerSaveJob, MegatronSFTTrainingJob, @@ -769,27 +767,9 @@ async def _reload_adapter(self, checkpoint_path: str, step: int) -> None: self._latest_step = step self._loaded_adapter_steps.add(step) - async def _update_in_flight_adapter( - self, checkpoint_path: str, step: int - ) -> dict[str, float]: - metrics, _receipt = await self._install_in_flight_adapter(checkpoint_path, step) - return metrics - - async def _install_in_flight_adapter( - self, - checkpoint_path: str, - step: int, - *, - verify_receiver_content: bool = False, - ) -> tuple[dict[str, float], dict[str, Any]]: + async def _update_in_flight_adapter(self, checkpoint_path: str, step: int) -> None: import httpx - from art.megatron.weights.update_replay import ( - inspect_adapter, - validate_committed_policy_version, - validate_receiver_update_receipt, - ) - self._raise_if_child_failed() self.serving_capabilities.require( "in_flight_lora_updates", operation="In-flight LoRA updates" @@ -797,7 +777,6 @@ async def _install_in_flight_adapter( self.serving_capabilities.require( "policy_token_spans", operation="In-flight LoRA updates" ) - rpc_started = time.perf_counter() async with httpx.AsyncClient() as client: response = await client.post( f"{self._vllm_base_url}/art/in_flight_lora_update", @@ -811,73 +790,16 @@ async def _install_in_flight_adapter( timeout=60.0, ) response.raise_for_status() - rpc_s = time.perf_counter() - rpc_started - json_response = getattr(response, "json", None) - response_body = json_response() if callable(json_response) else {} - verification_s = 0.0 - receipt: dict[str, Any] - if verify_receiver_content: - verification_started = time.perf_counter() - async with httpx.AsyncClient() as client: - verification_response = await client.get( - f"{self._vllm_base_url}/art/in_flight_lora_state", - params={"lora_slot": self._in_flight_lora_slot}, - **self._runtime_request_kwargs(), - timeout=60.0, - ) - verification_response.raise_for_status() - verification_body = verification_response.json() - response_body = { - **response_body, - "committed_state": verification_body.get("committed_state"), - "installed_content": verification_body.get("installed_content"), - } - verification_s = time.perf_counter() - verification_started - if verify_receiver_content: - snapshot = inspect_adapter(checkpoint_path) - receipt = validate_receiver_update_receipt( - response_body, - expected_policy_version=step, - expected_safetensors_sha256=snapshot.file_sha256, - ) - else: - validate_committed_policy_version(response_body, expected=step) - receipt = { - "committed_state": dict(response_body["committed_state"]), - "installed_content": None, - } - runtime_timing = response_body.get("timing_s", {}) - if not isinstance(runtime_timing, dict): - runtime_timing = {} self._latest_step = step self._loaded_adapter_steps.add(step) - metrics = { - "time/weight_update_service_rpc_s": rpc_s, - "weight_update/policy_version": float( - receipt["committed_state"]["policy_version"] - ), - } - if verify_receiver_content: - metrics["time/weight_update_receiver_verification_s"] = verification_s - for phase in ( - "begin_update", - "load_adapter", - "update_waiting_cache", - "commit_update", - "total", - ): - value = runtime_timing.get(phase) - if isinstance(value, (int, float)): - metrics[f"time/weight_update_vllm_{phase}_s"] = float(value) - return metrics, receipt async def _load_rollout_lora_for_step( self, checkpoint_path: str, step: int - ) -> dict[str, float]: + ) -> None: if self.rollout_weight_update_mode == "in_flight_lora": - return await self._update_in_flight_adapter(checkpoint_path, step) - await self._reload_adapter(checkpoint_path, step) - return {} + await self._update_in_flight_adapter(checkpoint_path, step) + else: + await self._reload_adapter(checkpoint_path, step) async def acquire_exact_adapter(self, step: int, checkpoint_path: str) -> str: if self.rollout_weights_mode != "lora": @@ -988,59 +910,6 @@ async def _sync_dedicated_merged_weights( pass self._latest_step = step - async def replay_lora_update( - self, - *, - source_lora_path: str, - output_lora_path: str, - policy_version: int, - content_version: int, - ) -> dict[str, Any]: - """Publish real resident Megatron tensors, then install that exact output.""" - if not self.is_dedicated: - raise RuntimeError("LoRA update replay requires dedicated trainer GPUs") - if self.rollout_weight_update_mode != "in_flight_lora": - raise RuntimeError("LoRA update replay requires in_flight_lora mode") - if os.path.exists(output_lora_path): - raise FileExistsError(f"Replay output already exists: {output_lora_path}") - await self._ensure_megatron_running() - self._clear_pending_jobs() - job_path, log_path = self._create_megatron_job_paths() - job = MegatronLoraPublishJob( - step=policy_version, - source_lora_path=source_lora_path, - output_lora_path=output_lora_path, - content_version=content_version, - allow_unvalidated_arch=self._allow_unvalidated_arch, - log_path=log_path, - ) - write_megatron_job(job, job_path=job_path) - publish_metrics: dict[str, float] | None = None - async for result in stream_megatron_job( - job, - job_path=job_path, - process=self._megatron_process, - process_log_path=self._megatron_log_path, - ): - if result.get("event") != LORA_READY_EVENT: - raise RuntimeError(f"Unexpected replay worker event: {result!r}") - if int(result.get("step", -1)) != policy_version: - raise RuntimeError(f"Replay worker published wrong step: {result!r}") - if int(result.get("content_version", -1)) != content_version: - raise RuntimeError(f"Replay worker published wrong content: {result!r}") - publish_metrics = self._weight_update_metrics_from_event(result) - if publish_metrics is None: - raise RuntimeError("Replay worker produced no LoRA-ready metrics") - install_metrics, receiver_receipt = await self._install_in_flight_adapter( - output_lora_path, - policy_version, - verify_receiver_content=True, - ) - return { - "metrics": {**publish_metrics, **install_metrics}, - "receiver": receiver_receipt, - } - async def _sleep_runtime(self) -> None: import httpx @@ -1287,27 +1156,16 @@ async def _handle_training_lora_ready( checkpoint_dir: str | None, staging_lora_path: str, step: int, - ) -> tuple[str, dict[str, float]]: - metrics: dict[str, float] = {} + ) -> str: if checkpoint_dir is None: checkpoint_dir = self._publish_staged_training_checkpoint( staging_lora_path=staging_lora_path, step=step, ) if self.is_dedicated and self.rollout_weights_mode == "lora": - metrics = await self._load_rollout_lora_for_step(checkpoint_dir, step) + await self._load_rollout_lora_for_step(checkpoint_dir, step) self._status(f"Loaded checkpoint {step} into vLLM") - return checkpoint_dir, metrics - - @staticmethod - def _weight_update_metrics_from_event( - result: dict[str, Any], - ) -> dict[str, float]: - return { - key: float(value) - for key, value in result.items() - if key not in {"event", "step"} and isinstance(value, (int, float)) - } + return checkpoint_dir async def _finish_training_checkpoint( self, @@ -1468,7 +1326,6 @@ async def train( write_megatron_job(job, job_path=job_path) checkpoint_dir: str | None = None optimizer_world_size: int | None = None - pending_metrics: dict[str, float] | None = None async for result in stream_megatron_job( job, job_path=job_path, @@ -1476,20 +1333,11 @@ async def train( process_log_path=self._megatron_log_path, ): if result.get("event") == LORA_READY_EVENT: - ( - checkpoint_dir, - install_metrics, - ) = await self._handle_training_lora_ready( + checkpoint_dir = await self._handle_training_lora_ready( checkpoint_dir=checkpoint_dir, staging_lora_path=staging_lora_path, step=next_step, ) - if pending_metrics is None: - pending_metrics = {} - pending_metrics.update( - self._weight_update_metrics_from_event(result) - ) - pending_metrics.update(install_metrics) continue if ( world_size := self._optimizer_ready_world_size( @@ -1498,14 +1346,7 @@ async def train( ) is not None: optimizer_world_size = world_size continue - if pending_metrics is not None: - yield pending_metrics - pending_metrics = { - key: float(value) for key, value in result.items() - } - - if pending_metrics is not None: - yield pending_metrics + yield {key: float(value) for key, value in result.items()} await self._finish_training_checkpoint( checkpoint_dir=checkpoint_dir, @@ -1545,7 +1386,6 @@ async def train( checkpoint_dir = None optimizer_world_size = None - pending_metrics = None async for result in stream_megatron_job( job, job_path=job_path, @@ -1553,20 +1393,11 @@ async def train( process_log_path=self._megatron_log_path, ): if result.get("event") == LORA_READY_EVENT: - ( - checkpoint_dir, - install_metrics, - ) = await self._handle_training_lora_ready( + checkpoint_dir = await self._handle_training_lora_ready( checkpoint_dir=checkpoint_dir, staging_lora_path=staging_lora_path, step=next_step, ) - if pending_metrics is None: - pending_metrics = {} - pending_metrics.update( - self._weight_update_metrics_from_event(result) - ) - pending_metrics.update(install_metrics) continue if ( world_size := self._optimizer_ready_world_size( @@ -1575,12 +1406,7 @@ async def train( ) is not None: optimizer_world_size = world_size continue - if pending_metrics is not None: - yield pending_metrics - pending_metrics = {key: float(value) for key, value in result.items()} - - if pending_metrics is not None: - yield pending_metrics + yield {key: float(value) for key, value in result.items()} await self._finish_training_checkpoint( checkpoint_dir=checkpoint_dir, diff --git a/src/art/megatron/train.py b/src/art/megatron/train.py index 7bbfdf57b..284620115 100644 --- a/src/art/megatron/train.py +++ b/src/art/megatron/train.py @@ -70,7 +70,6 @@ LORA_READY_EVENT, OPTIMIZER_READY_EVENT, MegatronJob, - MegatronLoraPublishJob, MegatronMergedTrainingJob, MegatronOptimizerSaveJob, MegatronSFTTrainingJob, @@ -890,49 +889,6 @@ def _run_megatron_job(runtime: TrainingRuntime, job: MegatronJob) -> None: pause_generation=False, ) return - if isinstance(job, MegatronLoraPublishJob): - adapter_model = _load_adapter_into_model( - runtime.model, - job.source_lora_path, - runtime.rank, - handler=runtime.model_support_handler, - ) - adapter_dtypes = {key: tensor.dtype for key, tensor in adapter_model.items()} - del adapter_model - from art.megatron.weights.lora_publish import iter_lora_modules - - content_value = float(job.content_version + 1) / 1024.0 - with torch.no_grad(): - for module in iter_lora_modules(runtime.model): - module.B_T.view(-1)[0] = content_value - publish_metrics = save_vllm_lora_from_model( - model=runtime.model, - adapter_dtypes=adapter_dtypes, - handler=runtime.model_support_handler, - adapter_config=load_adapter_config(job.source_lora_path), - output_dir=job.output_lora_path, - rank=runtime.rank, - world_size=runtime.world_size, - ) - if runtime.rank == 0: - assert publish_metrics is not None - _write_job_event( - job.log_path, - LORA_READY_EVENT, - step=job.step, - content_version=job.content_version, - **{ - "topology/trainer_tp": ps.get_tensor_model_parallel_world_size(), - "topology/trainer_pp": ps.get_pipeline_model_parallel_world_size(), - "topology/trainer_cp": ps.get_context_parallel_world_size(), - "topology/trainer_ep": ps.get_expert_model_parallel_world_size(), - "topology/trainer_etp": ( - ps.get_expert_tensor_parallel_world_size() - ), - }, - **publish_metrics.as_training_metrics(), - ) - return if isinstance(job, MegatronSFTTrainingJob): run_megatron_sft_job(runtime, job) return @@ -947,9 +903,7 @@ def _run_megatron_job(runtime: TrainingRuntime, job: MegatronJob) -> None: def _job_cleanup_path(job: MegatronJob) -> str | None: - if isinstance( - job, (MegatronOptimizerSaveJob, MegatronSyncJob, MegatronLoraPublishJob) - ): + if isinstance(job, (MegatronOptimizerSaveJob, MegatronSyncJob)): return None if isinstance(job, MegatronSFTTrainingJob): return job.sft_data_dir @@ -1072,7 +1026,7 @@ def _save_lora_and_optimizer( optimizer_ready_log_path: str | None = None, ) -> None: assert runtime.optimizer is not None - publish_metrics = save_vllm_lora_from_model( + save_vllm_lora_from_model( model=runtime.model, adapter_dtypes=adapter_dtypes, handler=runtime.model_support_handler, @@ -1082,13 +1036,7 @@ def _save_lora_and_optimizer( world_size=runtime.world_size, ) if lora_ready_log_path is not None and runtime.rank == 0: - assert publish_metrics is not None - _write_job_event( - lora_ready_log_path, - LORA_READY_EVENT, - step=step, - **publish_metrics.as_training_metrics(), - ) + _write_job_event(lora_ready_log_path, LORA_READY_EVENT, step=step) if _should_save_optimizer( runtime, step=step, diff --git a/src/art/megatron/weights/lora_publish.py b/src/art/megatron/weights/lora_publish.py index 715521e7d..e9d8b4b08 100644 --- a/src/art/megatron/weights/lora_publish.py +++ b/src/art/megatron/weights/lora_publish.py @@ -1,7 +1,4 @@ from collections.abc import Iterable, Sequence -from dataclasses import dataclass -from pathlib import Path -import time from typing import Any, NamedTuple import torch @@ -22,34 +19,6 @@ from art.megatron.training.model_chunks import ModelChunks -@dataclass(frozen=True) -class LoraPublishMetrics: - """Phase-level measurements for one trainer-side LoRA publication.""" - - gather_pack_s: float - stage_to_cpu_s: float - write_s: float - total_s: float - logical_bytes: int - transported_bytes: int - tensor_count: int - - def as_training_metrics(self) -> dict[str, int | float]: - return { - "time/weight_update_trainer_gather_pack_s": self.gather_pack_s, - "time/weight_update_trainer_stage_to_cpu_s": self.stage_to_cpu_s, - "time/weight_update_trainer_write_s": self.write_s, - "time/weight_update_trainer_publish_s": self.total_s, - "weight_update/logical_bytes": self.logical_bytes, - "weight_update/transported_bytes": self.transported_bytes, - "weight_update/tensor_count": self.tensor_count, - } - - -def _tensor_bytes(tensor: torch.Tensor) -> int: - return tensor.numel() * tensor.element_size() - - class PackedExpertShardMeta(NamedTuple): key: str owner_rank: int @@ -788,9 +757,7 @@ def save_vllm_lora_from_model( rank: int, world_size: int, slot_ref: LoRASlotRef | None = None, -) -> LoraPublishMetrics | None: - total_started = time.perf_counter() - gather_pack_started = total_started +) -> None: result = build_vllm_lora_tensors_from_model( model=model, adapter_dtypes=adapter_dtypes, @@ -800,28 +767,10 @@ def save_vllm_lora_from_model( world_size=world_size, slot_ref=slot_ref, ) - gather_pack_s = time.perf_counter() - gather_pack_started if result is None: - return None + return vllm_tensors, published_config = result - logical_bytes = sum(_tensor_bytes(tensor) for tensor in vllm_tensors.values()) - - stage_started = time.perf_counter() stager = _PinnedCpuStager() published_tensors = _stage_published_tensors(vllm_tensors, stager) stager.finish() - stage_to_cpu_s = time.perf_counter() - stage_started - - write_started = time.perf_counter() save_vllm_lora_tensors(output_dir, published_tensors, published_config) - write_s = time.perf_counter() - write_started - adapter_path = Path(output_dir) / "adapter_model.safetensors" - return LoraPublishMetrics( - gather_pack_s=gather_pack_s, - stage_to_cpu_s=stage_to_cpu_s, - write_s=write_s, - total_s=time.perf_counter() - total_started, - logical_bytes=logical_bytes, - transported_bytes=adapter_path.stat().st_size, - tensor_count=len(vllm_tensors), - ) diff --git a/src/art/megatron/weights/update_replay.py b/src/art/megatron/weights/update_replay.py deleted file mode 100644 index f901d0ed4..000000000 --- a/src/art/megatron/weights/update_replay.py +++ /dev/null @@ -1,540 +0,0 @@ -"""Receiver-only replay of pre-published LoRA updates. - -This diagnostic deliberately excludes trainer publication. Use -``MegatronService.replay_lora_update`` for end-to-end measurements. -""" - -from __future__ import annotations - -import argparse -import asyncio -from dataclasses import asdict, dataclass -import hashlib -import json -from pathlib import Path -import statistics -import time -from typing import Any - -import httpx -from safetensors import safe_open - -from art.megatron.model_support.lora_disk import load_adapter_config - - -@dataclass(frozen=True) -class TensorSpec: - name: str - shape: tuple[int, ...] - dtype: str - bytes: int - - -@dataclass(frozen=True) -class AdapterSnapshot: - path: str - file_sha256: str - manifest_sha256: str - logical_bytes: int - transported_bytes: int - tensor_count: int - rank: int - base_model: str - tensors: tuple[TensorSpec, ...] - - -@dataclass -class LoadCounters: - prompt_tokens: int = 0 - completion_tokens: int = 0 - requests: int = 0 - errors: int = 0 - - def __sub__(self, other: LoadCounters) -> LoadCounters: - return LoadCounters( - prompt_tokens=self.prompt_tokens - other.prompt_tokens, - completion_tokens=self.completion_tokens - other.completion_tokens, - requests=self.requests - other.requests, - errors=self.errors - other.errors, - ) - - -def _file_sha256(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as handle: - for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def inspect_adapter(path: str | Path) -> AdapterSnapshot: - adapter_dir = Path(path).resolve() - tensor_path = adapter_dir / "adapter_model.safetensors" - config = load_adapter_config(adapter_dir) - rank = int(config.get("r", config.get("rank", 0))) - base_model = str(config.get("base_model_name_or_path", "")) - tensors: list[TensorSpec] = [] - logical_bytes = 0 - with safe_open(tensor_path, framework="pt", device="cpu") as handle: - for name in sorted(handle.keys()): - tensor = handle.get_tensor(name) - tensor_bytes = tensor.numel() * tensor.element_size() - logical_bytes += tensor_bytes - tensors.append( - TensorSpec( - name=name, - shape=tuple(int(dim) for dim in tensor.shape), - dtype=str(tensor.dtype), - bytes=tensor_bytes, - ) - ) - manifest_payload = json.dumps( - [asdict(tensor) for tensor in tensors], - sort_keys=True, - separators=(",", ":"), - ).encode() - return AdapterSnapshot( - path=str(adapter_dir), - file_sha256=_file_sha256(tensor_path), - manifest_sha256=hashlib.sha256(manifest_payload).hexdigest(), - logical_bytes=logical_bytes, - transported_bytes=tensor_path.stat().st_size, - tensor_count=len(tensors), - rank=rank, - base_model=base_model, - tensors=tuple(tensors), - ) - - -def validate_replay_snapshots( - snapshots: list[AdapterSnapshot], - *, - expected_rank: int, - expected_experts: int, - expected_model_substring: str, -) -> None: - if not snapshots: - raise ValueError("At least one adapter snapshot is required") - paths = [snapshot.path for snapshot in snapshots] - hashes = [snapshot.file_sha256 for snapshot in snapshots] - if len(set(paths)) != len(paths): - raise ValueError("Receiver-only replay requires unique snapshot paths") - if len(set(hashes)) != len(hashes): - raise ValueError("Receiver-only replay requires unique snapshot contents") - reference = snapshots[0] - for snapshot in snapshots: - if snapshot.rank != expected_rank: - raise ValueError( - f"{snapshot.path} has rank {snapshot.rank}; expected {expected_rank}" - ) - if expected_model_substring not in snapshot.base_model: - raise ValueError( - f"{snapshot.path} targets {snapshot.base_model!r}; expected a model " - f"containing {expected_model_substring!r}" - ) - if snapshot.manifest_sha256 != reference.manifest_sha256: - raise ValueError( - f"{snapshot.path} does not have the reference serving tensor layout" - ) - expert_tensors = [ - tensor - for tensor in reference.tensors - if ".mlp.experts." in tensor.name or ".mlp.experts.base_layer." in tensor.name - ] - if not expert_tensors: - raise ValueError("Adapter has no packed Qwen MoE expert tensors") - mismatched = [ - tensor.name for tensor in expert_tensors if expected_experts not in tensor.shape - ] - if mismatched: - raise ValueError( - "Packed expert tensors do not preserve the expected expert dimension: " - + ", ".join(mismatched[:3]) - ) - - -def validate_committed_policy_version( - response_body: dict[str, Any], *, expected: int -) -> None: - committed_state = response_body.get("committed_state") - if not isinstance(committed_state, dict): - raise RuntimeError("Runtime response has no committed receiver state") - committed = int(committed_state.get("policy_version", -1)) - if committed != expected: - raise RuntimeError(f"Runtime committed policy {committed}; expected {expected}") - if committed_state.get("update_active") or committed_state.get("admission_blocked"): - raise RuntimeError(f"Runtime did not finish policy commit: {committed_state!r}") - - -def validate_receiver_update_receipt( - response_body: dict[str, Any], - *, - expected_policy_version: int, - expected_safetensors_sha256: str, -) -> dict[str, Any]: - validate_committed_policy_version(response_body, expected=expected_policy_version) - installed_content = response_body.get("installed_content") - if not isinstance(installed_content, dict): - raise RuntimeError("Runtime response has no installed-content receipt") - source = str(installed_content.get("source_safetensors_sha256", "")) - installed = str(installed_content.get("installed_safetensors_sha256", "")) - if ( - source != expected_safetensors_sha256 - or installed != expected_safetensors_sha256 - ): - raise RuntimeError( - "Runtime installed checksum does not match the published adapter: " - f"published={expected_safetensors_sha256}, source={source}, " - f"installed={installed}" - ) - return { - "committed_state": dict(response_body["committed_state"]), - "installed_content": dict(installed_content), - } - - -def validate_bench_request_trace( - request_jsonl: str | Path, - manifest_path: str | Path, - *, - bench_commit: str, -) -> tuple[list[dict[str, Any]], dict[str, Any]]: - if len(bench_commit) != 40 or any( - character not in "0123456789abcdef" for character in bench_commit.lower() - ): - raise ValueError("--bench-commit must be an exact 40-character Git commit") - trace_path = Path(request_jsonl) - trace_payload = trace_path.read_bytes() - manifest = json.loads(Path(manifest_path).read_text()) - requests = manifest.get("requests") - if not isinstance(requests, dict): - raise ValueError("Bench trace manifest has no requests object") - digest = hashlib.sha256(trace_payload).hexdigest() - if digest != requests.get("sha256"): - raise ValueError( - f"Bench request trace digest is {digest}; expected {requests.get('sha256')}" - ) - rows = [ - json.loads(line) for line in trace_payload.decode("utf-8").splitlines() if line - ] - expected_hashes = requests.get("hashes") - selected_ids = manifest.get("selection", {}).get("selected_ids") - if ( - not isinstance(expected_hashes, list) - or not isinstance(selected_ids, list) - or len(rows) != len(expected_hashes) - or len(rows) != len(selected_ids) - ): - raise ValueError("Bench trace row IDs and hashes do not align") - validated: list[dict[str, Any]] = [] - for request, expected, datapoint_id in zip( - rows, expected_hashes, selected_ids, strict=True - ): - if not isinstance(expected, dict): - raise ValueError("Bench request hash entry is not an object") - if str(expected.get("datapoint_id")) != str(datapoint_id): - raise ValueError("Bench request hash IDs do not align with selected IDs") - canonical = json.dumps( - request, ensure_ascii=False, separators=(",", ":"), sort_keys=True - ).encode() - request_sha256 = hashlib.sha256(canonical).hexdigest() - if request_sha256 != expected.get("sha256"): - raise ValueError( - f"Bench request {datapoint_id} hash is {request_sha256}; " - f"expected {expected.get('sha256')}" - ) - validated.append( - { - "id": str(datapoint_id), - "canonical_request_sha256": request_sha256, - "request": request, - } - ) - return validated, { - "bench_commit": bench_commit, - "trace_sha256": digest, - "manifest": manifest, - } - - -def evaluate_replay_acceptance( - measured_metrics: list[dict[str, float]], -) -> dict[str, Any]: - if len(measured_metrics) != 30: - raise RuntimeError( - f"Acceptance requires exactly 30 measured updates; got {len(measured_metrics)}" - ) - combined = [ - row["time/weight_update_trainer_publish_s"] - + row["time/weight_update_service_rpc_s"] - for row in measured_metrics - ] - trainer = [row["time/weight_update_trainer_publish_s"] for row in measured_metrics] - receiver = [row["time/weight_update_vllm_total_s"] for row in measured_metrics] - gates = { - "combined_mean": { - "value_s": statistics.fmean(combined), - "lower_s": 5.001 * 0.85, - "upper_s": 5.001 * 1.15, - }, - "trainer_publish_mean": { - "value_s": statistics.fmean(trainer), - "lower_s": 2.672 * 0.75, - "upper_s": 2.672 * 1.25, - }, - "receiver_install_mean": { - "value_s": statistics.fmean(receiver), - "lower_s": 2.320 * 0.75, - "upper_s": 2.320 * 1.25, - }, - } - failures = [ - name - for name, gate in gates.items() - if not gate["lower_s"] <= gate["value_s"] <= gate["upper_s"] - ] - return {"passed": not failures, "failures": failures, "gates": gates} - - -class FixedLoad: - def __init__( - self, - client: httpx.AsyncClient, - *, - request: dict[str, Any], - concurrency: int, - ) -> None: - self._client = client - self._request = request - self._concurrency = concurrency - self._running = False - self._tasks: list[asyncio.Task[None]] = [] - self.counters = LoadCounters() - - async def _worker(self) -> None: - while self._running: - try: - response = await self._client.post( - "/v1/chat/completions", json=self._request - ) - response.raise_for_status() - usage = response.json().get("usage", {}) - self.counters.prompt_tokens += int(usage.get("prompt_tokens", 0)) - self.counters.completion_tokens += int( - usage.get("completion_tokens", 0) - ) - self.counters.requests += 1 - except BaseException: - self.counters.errors += 1 - if not self._running: - return - - def snapshot(self) -> LoadCounters: - return LoadCounters(**asdict(self.counters)) - - async def start(self) -> None: - self._running = True - self._tasks = [ - asyncio.create_task(self._worker()) for _ in range(self._concurrency) - ] - - async def stop(self) -> None: - self._running = False - await asyncio.gather(*self._tasks, return_exceptions=True) - - -def _percentile(values: list[float], percentile: float) -> float: - if not values: - return 0.0 - ordered = sorted(values) - index = min(round((len(ordered) - 1) * percentile), len(ordered) - 1) - return ordered[index] - - -async def replay_updates(args: argparse.Namespace) -> dict[str, Any]: - snapshots = [inspect_adapter(path) for path in args.adapter_path] - validate_replay_snapshots( - snapshots, - expected_rank=args.expected_rank, - expected_experts=args.expected_experts, - expected_model_substring=args.expected_model_substring, - ) - server_paths = args.server_adapter_path or [item.path for item in snapshots] - if len(server_paths) != len(snapshots): - raise ValueError("--server-adapter-path must match --adapter-path count") - total_updates = args.warmups + args.updates - if len(snapshots) != total_updates: - raise ValueError( - "Receiver-only replay requires exactly one unique snapshot per update: " - f"{len(snapshots)} != {total_updates}" - ) - - headers = {"Authorization": f"Bearer {args.api_key}"} if args.api_key else {} - timeout = httpx.Timeout(args.timeout_s) - async with httpx.AsyncClient( - base_url=args.server_url.rstrip("/"), - headers=headers, - timeout=timeout, - ) as client: - metrics_response = await client.get("/art/metrics") - metrics_response.raise_for_status() - runtime_metrics = metrics_response.json() - runtime_world_size = int( - runtime_metrics.get("metrics", {}).get("world_size", 0) - ) - if runtime_world_size != args.expected_inference_world_size: - raise ValueError( - f"Runtime world size is {runtime_world_size}; expected " - f"{args.expected_inference_world_size}" - ) - - request: dict[str, Any] | None = None - if args.mode == "fixed-load": - if args.request_json is None: - raise ValueError("--request-json is required in fixed-load mode") - request = json.loads(Path(args.request_json).read_text()) - fixed_load = ( - FixedLoad(client, request=request, concurrency=args.load_concurrency) - if request is not None - else None - ) - if fixed_load is not None: - await fixed_load.start() - - update_records: list[dict[str, Any]] = [] - measured_update_s = 0.0 - try: - for update_index in range(total_updates): - snapshot_index = update_index - policy_version = args.first_policy_version + update_index - before_load = fixed_load.snapshot() if fixed_load else LoadCounters() - started = time.perf_counter() - response = await client.post( - "/art/in_flight_lora_update", - json={ - "model_name": args.model_name, - "base_model_name": args.base_model_name, - "lora_slot": args.lora_slot, - "lora_path": server_paths[snapshot_index], - "policy_version": policy_version, - }, - ) - response.raise_for_status() - wall_s = time.perf_counter() - started - body = response.json() - validate_committed_policy_version(body, expected=policy_version) - after_load = fixed_load.snapshot() if fixed_load else LoadCounters() - if update_index >= args.warmups: - measured_update_s += wall_s - update_records.append( - { - "update_index": update_index - args.warmups, - "policy_version": policy_version, - "snapshot_index": snapshot_index, - "snapshot_sha256": snapshots[snapshot_index].file_sha256, - "wall_s": wall_s, - "runtime_timing_s": body.get("timing_s", {}), - "load_during_update": asdict(after_load - before_load), - } - ) - - control_load = LoadCounters() - control_s = 0.0 - if fixed_load is not None and measured_update_s > 0: - before_control = fixed_load.snapshot() - control_started = time.perf_counter() - await asyncio.sleep(measured_update_s) - control_s = time.perf_counter() - control_started - control_load = fixed_load.snapshot() - before_control - finally: - if fixed_load is not None: - await fixed_load.stop() - - post_snapshots = [inspect_adapter(path) for path in args.adapter_path] - if [item.file_sha256 for item in snapshots] != [ - item.file_sha256 for item in post_snapshots - ]: - raise RuntimeError("Replay mutated a captured adapter snapshot") - - update_wall = [float(record["wall_s"]) for record in update_records] - update_completion_tokens = sum( - int(record["load_during_update"]["completion_tokens"]) - for record in update_records - ) - control_completion_rate = ( - control_load.completion_tokens / control_s if control_s > 0 else 0.0 - ) - lost_completion_tokens = max( - control_completion_rate * measured_update_s - update_completion_tokens, - 0.0, - ) - return { - "schema_version": 1, - "measurement_scope": "receiver_only_pre_published_snapshots", - "mode": args.mode, - "warmups": args.warmups, - "measured_updates": args.updates, - "runtime_world_size": runtime_world_size, - "snapshots": [ - {key: value for key, value in asdict(snapshot).items() if key != "tensors"} - for snapshot in snapshots - ], - "updates": update_records, - "control": { - "wall_s": control_s, - "load": asdict(control_load), - "completion_tok_per_s": control_completion_rate, - }, - "summary": { - "update_wall_mean_s": statistics.fmean(update_wall), - "update_wall_p50_s": _percentile(update_wall, 0.50), - "update_wall_p95_s": _percentile(update_wall, 0.95), - "measured_update_wall_s": measured_update_s, - "completion_tokens_during_updates": update_completion_tokens, - "lost_completion_tokens": lost_completion_tokens, - "lost_completion_tokens_per_update": lost_completion_tokens - / max(args.updates, 1), - }, - } - - -def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser( - description="Replay captured Qwen3.6 LoRA updates against ART vLLM" - ) - parser.add_argument("--server-url", required=True) - parser.add_argument("--adapter-path", action="append", required=True) - parser.add_argument("--server-adapter-path", action="append") - parser.add_argument("--model-name", required=True) - parser.add_argument("--base-model-name", required=True) - parser.add_argument("--lora-slot", required=True) - parser.add_argument("--api-key") - parser.add_argument("--mode", choices=("idle", "fixed-load"), default="idle") - parser.add_argument("--request-json") - parser.add_argument("--load-concurrency", type=int, default=8) - parser.add_argument("--warmups", type=int, default=5) - parser.add_argument("--updates", type=int, default=30) - parser.add_argument("--first-policy-version", type=int, default=1) - parser.add_argument("--timeout-s", type=float, default=120.0) - parser.add_argument("--expected-rank", type=int, default=8) - parser.add_argument("--expected-experts", type=int, default=256) - parser.add_argument("--expected-model-substring", default="Qwen3.6-35B-A3B") - parser.add_argument("--expected-inference-world-size", type=int, default=2) - parser.add_argument("--output", type=Path, required=True) - return parser - - -def main(argv: list[str] | None = None) -> None: - args = build_parser().parse_args(argv) - if args.warmups < 0 or args.updates < 1: - raise SystemExit( - "--warmups must be non-negative and --updates must be positive" - ) - result = asyncio.run(replay_updates(args)) - args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n") - print(json.dumps(result["summary"], sort_keys=True)) - - -if __name__ == "__main__": - main() diff --git a/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py b/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py index 1f3a61f0e..eb7501217 100644 --- a/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py +++ b/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py @@ -173,12 +173,6 @@ async def admit(): blocked = not admission.done() await coordinator.commit_update(slot, 5, new) version, admitted_lora = await admission - committed_state = await coordinator.committed_state(slot) - stale_rejected = False - try: - await coordinator.begin_update(slot, 5) - except ValueError: - stale_rejected = True async with coordinator.admission(slot): cancelled_update = asyncio.create_task(coordinator.begin_update(slot)) @@ -197,8 +191,6 @@ async def admit(): "policy_version": version, "lora_path": admitted_lora.lora_path, "recovered_after_cancel": recovered, - "stale_rejected": stale_rejected, - "committed_state": committed_state, }, sort_keys=True)) asyncio.run(main()) @@ -212,16 +204,6 @@ async def admit(): "lora_path": "new", "policy_version": 5, "recovered_after_cancel": True, - "stale_rejected": True, - "committed_state": { - "active_admissions": 0, - "admission_blocked": False, - "installed_lora_int_id": None, - "installed_lora_path": "new", - "lora_slot": "model:active", - "policy_version": 5, - "update_active": False, - }, } diff --git a/tests/integration/megatron/runtime_isolation/test_service_runtime_boundary.py b/tests/integration/megatron/runtime_isolation/test_service_runtime_boundary.py index 0b2932c90..c6c32685a 100644 --- a/tests/integration/megatron/runtime_isolation/test_service_runtime_boundary.py +++ b/tests/integration/megatron/runtime_isolation/test_service_runtime_boundary.py @@ -40,19 +40,6 @@ class _AsyncOkResponse: def raise_for_status(self) -> None: return None - def json(self) -> dict[str, object]: - return { - "committed_state": { - "policy_version": 4, - "update_active": False, - "admission_blocked": False, - }, - "installed_content": { - "source_safetensors_sha256": "abc", - "installed_safetensors_sha256": "abc", - }, - } - class _RecordingAsyncClient: def __init__( @@ -77,16 +64,6 @@ async def post( self._posts.append((url, json if json is not None else params, timeout)) return _AsyncOkResponse() - async def get( - self, - url: str, - *, - params: dict[str, object] | None = None, - timeout: float, - ) -> _AsyncOkResponse: - self._posts.append((url, params, timeout)) - return _AsyncOkResponse() - def test_megatron_default_lora_adapter_config_uses_model_lora_config( tmp_path: Path, @@ -126,10 +103,6 @@ async def test_megatron_in_flight_eval_uses_immutable_adapter_slot( service._vllm_runtime.port = 8123 posts: list[tuple[str, dict[str, object] | None, float]] = [] monkeypatch.setattr(httpx, "AsyncClient", lambda: _RecordingAsyncClient(posts)) - monkeypatch.setattr( - "art.megatron.weights.update_replay.inspect_adapter", - lambda _path: SimpleNamespace(file_sha256="abc"), - ) checkpoint_path = str(tmp_path / "checkpoints" / "4") assert ( @@ -179,9 +152,6 @@ async def test_external_in_flight_update_maps_checkpoint_path( monkeypatch: pytest.MonkeyPatch, ) -> None: local_root = str(tmp_path / "local") - monkeypatch.setattr( - MegatronService, "_validate_megatron_dependencies", lambda _self: None - ) service = MegatronService( model_name="test-model", base_model="Qwen/Qwen3-0.6B", diff --git a/tests/unit/test_weight_update_replay.py b/tests/unit/test_weight_update_replay.py deleted file mode 100644 index c9063e6a3..000000000 --- a/tests/unit/test_weight_update_replay.py +++ /dev/null @@ -1,288 +0,0 @@ -import hashlib -import json -from pathlib import Path - -import pytest -import torch - -from art.megatron.model_support.lora_disk import save_vllm_lora_tensors -from art.megatron.runtime.jobs import ( - MegatronLoraPublishJob, - dump_megatron_job, - load_megatron_job, -) -from art.megatron.service import MegatronService -from art.megatron.weights.lora_publish import LoraPublishMetrics -from art.megatron.weights.update_replay import ( - evaluate_replay_acceptance, - inspect_adapter, - validate_bench_request_trace, - validate_committed_policy_version, - validate_receiver_update_receipt, - validate_replay_snapshots, -) - - -def _write_snapshot( - path: Path, - *, - rank: int = 8, - experts: int = 256, - value: float = 1.0, -) -> None: - save_vllm_lora_tensors( - path, - { - ( - "base_model.model.model.language_model.layers.0." - "mlp.experts.lora_A.weight" - ): torch.full((experts, rank, 2), value, dtype=torch.bfloat16), - ( - "base_model.model.model.language_model.layers.0." - "mlp.experts.lora_B.weight" - ): torch.full((experts, 4, rank), value, dtype=torch.bfloat16), - }, - { - "base_model_name_or_path": "Qwen/Qwen3.6-35B-A3B", - "r": rank, - "lora_alpha": rank, - }, - ) - - -def test_publish_metrics_use_stable_training_metric_names() -> None: - metrics = LoraPublishMetrics( - gather_pack_s=1.0, - stage_to_cpu_s=2.0, - write_s=3.0, - total_s=6.0, - logical_bytes=10, - transported_bytes=12, - tensor_count=4, - ) - - assert metrics.as_training_metrics() == { - "time/weight_update_trainer_gather_pack_s": 1.0, - "time/weight_update_trainer_stage_to_cpu_s": 2.0, - "time/weight_update_trainer_write_s": 3.0, - "time/weight_update_trainer_publish_s": 6.0, - "weight_update/logical_bytes": 10, - "weight_update/transported_bytes": 12, - "weight_update/tensor_count": 4, - } - - -def test_weight_update_event_metrics_exclude_control_fields() -> None: - assert MegatronService._weight_update_metrics_from_event( - { - "event": "lora_ready", - "step": 7, - "time/weight_update_trainer_publish_s": 5.0, - "weight_update/logical_bytes": 1024, - "diagnostic": "ignored", - } - ) == { - "time/weight_update_trainer_publish_s": 5.0, - "weight_update/logical_bytes": 1024.0, - } - - -def test_replay_snapshot_preserves_layout_bytes_and_integrity(tmp_path: Path) -> None: - first_path = tmp_path / "first" - second_path = tmp_path / "second" - _write_snapshot(first_path, value=1.0) - _write_snapshot(second_path, value=2.0) - - first = inspect_adapter(first_path) - second = inspect_adapter(second_path) - validate_replay_snapshots( - [first, second], - expected_rank=8, - expected_experts=256, - expected_model_substring="Qwen3.6-35B-A3B", - ) - - assert first.manifest_sha256 == second.manifest_sha256 - assert first.file_sha256 != second.file_sha256 - assert first.logical_bytes == 256 * 8 * (2 + 4) * 2 - assert ( - first.transported_bytes - == (first_path / "adapter_model.safetensors").stat().st_size - ) - assert first.tensor_count == 2 - - -def test_replay_snapshot_rejects_nonrepresentative_rank(tmp_path: Path) -> None: - path = tmp_path / "rank-one" - _write_snapshot(path, rank=1) - - with pytest.raises(ValueError, match="has rank 1; expected 8"): - validate_replay_snapshots( - [inspect_adapter(path)], - expected_rank=8, - expected_experts=256, - expected_model_substring="Qwen3.6-35B-A3B", - ) - - -def test_receiver_replay_rejects_duplicate_snapshot_contents(tmp_path: Path) -> None: - first_path = tmp_path / "first" - second_path = tmp_path / "second" - _write_snapshot(first_path, value=1.0) - _write_snapshot(second_path, value=1.0) - - with pytest.raises(ValueError, match="unique snapshot contents"): - validate_replay_snapshots( - [inspect_adapter(first_path), inspect_adapter(second_path)], - expected_rank=8, - expected_experts=256, - expected_model_substring="Qwen3.6-35B-A3B", - ) - - -def test_lora_publish_job_roundtrips_through_worker_protocol() -> None: - job = MegatronLoraPublishJob( - step=9, - source_lora_path="/checkpoints/seed", - output_lora_path="/scratch/replay/0009", - content_version=8, - allow_unvalidated_arch=True, - log_path="/scratch/replay/0009.jsonl", - ) - - assert load_megatron_job(dump_megatron_job(job)) == job - - -def test_policy_version_must_match_committed_runtime_version() -> None: - validate_committed_policy_version( - { - "policy_version": 999, - "committed_state": { - "policy_version": 12, - "update_active": False, - "admission_blocked": False, - }, - }, - expected=12, - ) - - with pytest.raises(RuntimeError, match="committed policy 11; expected 12"): - validate_committed_policy_version( - { - "committed_state": { - "policy_version": 11, - "update_active": False, - "admission_blocked": False, - } - }, - expected=12, - ) - - -def test_receiver_receipt_requires_committed_state_and_installed_checksum() -> None: - receipt = { - "committed_state": { - "policy_version": 7, - "update_active": False, - "admission_blocked": False, - }, - "installed_content": { - "source_safetensors_sha256": "abc", - "installed_safetensors_sha256": "abc", - }, - } - - assert ( - validate_receiver_update_receipt( - receipt, - expected_policy_version=7, - expected_safetensors_sha256="abc", - )["committed_state"]["policy_version"] - == 7 - ) - - receipt["installed_content"]["installed_safetensors_sha256"] = "stale" - with pytest.raises(RuntimeError, match="installed checksum"): - validate_receiver_update_receipt( - receipt, - expected_policy_version=7, - expected_safetensors_sha256="abc", - ) - - -def test_snapshot_config_is_machine_readable(tmp_path: Path) -> None: - path = tmp_path / "snapshot" - _write_snapshot(path) - - config = json.loads((path / "adapter_config.json").read_text()) - assert config["art_lora_format"] == "vllm" - - -def test_bench_trace_requires_exact_external_provenance(tmp_path: Path) -> None: - request = { - "model": "Qwen/Qwen3.6-35B-A3B", - "messages": [{"role": "user", "content": "hello"}], - } - canonical = json.dumps( - request, ensure_ascii=False, separators=(",", ":"), sort_keys=True - ).encode() - trace = tmp_path / "request-trace.jsonl" - trace.write_bytes(canonical + b"\n") - digest = hashlib.sha256(trace.read_bytes()).hexdigest() - request_digest = hashlib.sha256(canonical).hexdigest() - manifest = tmp_path / "manifest.json" - manifest.write_text( - json.dumps( - { - "selection": {"selected_ids": ["row-1"]}, - "requests": { - "sha256": digest, - "hashes": [{"datapoint_id": "row-1", "sha256": request_digest}], - }, - } - ) - ) - - rows, provenance = validate_bench_request_trace( - trace, manifest, bench_commit="a" * 40 - ) - - assert rows[0]["id"] == "row-1" - assert provenance["trace_sha256"] == digest - assert provenance["bench_commit"] == "a" * 40 - - -def test_replay_acceptance_enforces_combined_and_component_gates() -> None: - passing = [ - { - "time/weight_update_trainer_publish_s": 2.672, - "time/weight_update_service_rpc_s": 2.329, - "time/weight_update_vllm_total_s": 2.320, - } - for _ in range(30) - ] - assert evaluate_replay_acceptance(passing)["passed"] is True - - failing = [dict(row) for row in passing] - failing[0]["time/weight_update_vllm_total_s"] = 100.0 - assert evaluate_replay_acceptance(failing)["passed"] is False - - -def test_replay_uploader_matches_bench_required_files() -> None: - root = Path(__file__).parents[2] - service = (root / "scripts/benchmarks/service_lora_update_replay.py").read_text() - uploader = (root / "scripts/benchmarks/upload_lora_update_replay.py").read_text() - - assert 'max_model_len": 32769' in service - assert "validate_bench_request_trace" in service - assert "--bench-commit" in service - for filename in ( - "manifest.json", - "request-trace.jsonl", - "tensor-manifest.json", - "samples.jsonl", - "summary.json", - "stdout.log", - "stderr.log", - ): - assert filename in uploader diff --git a/vllm_runtime/src/art_vllm_runtime/dedicated_server.py b/vllm_runtime/src/art_vllm_runtime/dedicated_server.py index 299d0c72e..36b8a0ffd 100644 --- a/vllm_runtime/src/art_vllm_runtime/dedicated_server.py +++ b/vllm_runtime/src/art_vllm_runtime/dedicated_server.py @@ -2,11 +2,9 @@ import argparse import asyncio -import hashlib from http import HTTPStatus import json import os -import time from fastapi.responses import JSONResponse from pydantic import BaseModel, Field @@ -49,15 +47,6 @@ class _InFlightLoraUpdateRequest(BaseModel): is_3d_lora_weight: bool = False -def _adapter_safetensors_sha256(lora_path: str) -> str: - path = os.path.join(lora_path, "adapter_model.safetensors") - digest = hashlib.sha256() - with open(path, "rb") as handle: - for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser(description="ART dedicated vLLM server") parser.add_argument("--model", required=True, help="Base model name or path") @@ -238,12 +227,8 @@ async def in_flight_lora_update( models = raw_request.app.state.openai_serving_models engine_client = engine(raw_request) coordinator = lora_update_coordinator(models, engine_client) - update_started = time.perf_counter() - begin_started = update_started - await coordinator.begin_update(lora_slot, policy_version) - begin_s = time.perf_counter() - begin_started + await coordinator.begin_update(lora_slot) try: - load_started = time.perf_counter() load_result = await models.load_lora_adapter( LoadLoRAAdapterRequest( lora_name=lora_slot, @@ -253,39 +238,23 @@ async def in_flight_lora_update( ), base_model_name=body.base_model_name, ) - load_s = time.perf_counter() - load_started if isinstance(load_result, ErrorResponse): await coordinator.fail_update(lora_slot) return JSONResponse( content=load_result.model_dump(mode="python"), status_code=load_result.error.code, ) - cache_started = time.perf_counter() waiting_cache_salt = await engine_client.engine_core.call_utility_async( "art_update_waiting_lora_cache_salt", lora_slot, policy_version, ) - update_waiting_cache_s = time.perf_counter() - cache_started - commit_started = time.perf_counter() await coordinator.commit_update( lora_slot, policy_version, models.lora_requests[lora_slot], ) - commit_s = time.perf_counter() - commit_started - committed_state = await coordinator.committed_state(lora_slot) - timing_s = { - "begin_update": begin_s, - "load_adapter": load_s, - "update_waiting_cache": update_waiting_cache_s, - "commit_update": commit_s, - "total": time.perf_counter() - update_started, - } - from art_vllm_runtime.metrics import ( - record_policy_cache_waiting_update, - record_weight_update, - ) + from art_vllm_runtime.metrics import record_policy_cache_waiting_update record_policy_cache_waiting_update( updated=int(waiting_cache_salt["updated_waiting_requests"]), @@ -293,7 +262,6 @@ async def in_flight_lora_update( waiting_cache_salt["skipped_started_waiting_requests"] ), ) - record_weight_update(timing_s) except BaseException: await coordinator.fail_update(lora_slot) raise @@ -302,38 +270,8 @@ async def in_flight_lora_update( "status": "updated", "model_name": public_model_name, "lora_slot": lora_slot, - "policy_version": committed_state["policy_version"], - "committed_state": committed_state, + "policy_version": policy_version, "waiting_cache_salt": waiting_cache_salt, - "timing_s": timing_s, - } - ) - - @router.get("/art/in_flight_lora_state") - async def in_flight_lora_state( - raw_request: Request, - lora_slot: str = Query(min_length=1), - ) -> JSONResponse: - models = raw_request.app.state.openai_serving_models - engine_client = engine(raw_request) - coordinator = lora_update_coordinator(models, engine_client) - committed_state = await coordinator.committed_state(lora_slot) - installed_path = committed_state.get("installed_lora_path") - installed_sha256 = ( - _adapter_safetensors_sha256(str(installed_path)) - if installed_path is not None - else None - ) - return JSONResponse( - content={ - "committed_state": committed_state, - "installed_content": { - "source_safetensors_sha256": installed_sha256, - "installed_safetensors_sha256": installed_sha256, - "verification_scope": ( - "runtime-selected safetensors after successful GPU load" - ), - }, } ) diff --git a/vllm_runtime/src/art_vllm_runtime/metrics.py b/vllm_runtime/src/art_vllm_runtime/metrics.py index 0f1db59c0..0c8be3d8f 100644 --- a/vllm_runtime/src/art_vllm_runtime/metrics.py +++ b/vllm_runtime/src/art_vllm_runtime/metrics.py @@ -32,12 +32,6 @@ def __init__(self) -> None: "policy_cache_unsalted_lora_requests_total": 0.0, "policy_cache_waiting_requests_updated_total": 0.0, "policy_cache_started_waiting_requests_skipped_total": 0.0, - "weight_update_count_total": 0.0, - "weight_update_begin_s_total": 0.0, - "weight_update_load_adapter_s_total": 0.0, - "weight_update_waiting_cache_s_total": 0.0, - "weight_update_commit_s_total": 0.0, - "weight_update_s_total": 0.0, } def configure(self, vllm_config: Any, *, engine_idx: int) -> None: @@ -212,18 +206,6 @@ def record_policy_cache_waiting_update( float(skipped_started) ) - def record_weight_update(self, timing_s: dict[str, float]) -> None: - with self._lock: - self._counters["weight_update_count_total"] += 1.0 - for phase, metric in ( - ("begin_update", "weight_update_begin_s_total"), - ("load_adapter", "weight_update_load_adapter_s_total"), - ("update_waiting_cache", "weight_update_waiting_cache_s_total"), - ("commit_update", "weight_update_commit_s_total"), - ("total", "weight_update_s_total"), - ): - self._counters[metric] += float(timing_s.get(phase, 0.0)) - _STATE = _ArtRuntimeMetricsState() @@ -264,7 +246,3 @@ def record_policy_cache_waiting_update(*, updated: int, skipped_started: int) -> updated=updated, skipped_started=skipped_started, ) - - -def record_weight_update(timing_s: dict[str, float]) -> None: - _STATE.record_weight_update(timing_s) diff --git a/vllm_runtime/src/art_vllm_runtime/policy_spans.py b/vllm_runtime/src/art_vllm_runtime/policy_spans.py index 637e9cea9..da90c7192 100644 --- a/vllm_runtime/src/art_vllm_runtime/policy_spans.py +++ b/vllm_runtime/src/art_vllm_runtime/policy_spans.py @@ -101,23 +101,12 @@ async def admission( state.active_admissions -= 1 state.condition.notify_all() - async def begin_update( - self, lora_slot: str, policy_version: int | None = None - ) -> None: + async def begin_update(self, lora_slot: str) -> None: state = self._state(lora_slot) async with state.condition: acquired = False try: await state.condition.wait_for(lambda: not state.update_active) - if ( - policy_version is not None - and state.policy_version is not None - and int(policy_version) <= state.policy_version - ): - raise ValueError( - f"LoRA slot {lora_slot!r} policy version must increase: " - f"{policy_version} <= {state.policy_version}" - ) state.update_active = True state.blocked = True acquired = True @@ -139,14 +128,6 @@ async def commit_update( async with state.condition: if not state.update_active: raise RuntimeError(f"No active LoRA update for slot {lora_slot!r}") - if ( - state.policy_version is not None - and int(policy_version) <= state.policy_version - ): - raise ValueError( - f"LoRA slot {lora_slot!r} policy version must increase: " - f"{policy_version} <= {state.policy_version}" - ) state.policy_version = int(policy_version) state.lora_request = lora_request state.update_active = False @@ -162,28 +143,6 @@ async def fail_update(self, lora_slot: str) -> None: state.blocked = True state.condition.notify_all() - async def committed_state(self, lora_slot: str) -> dict[str, Any]: - state = self._state(lora_slot) - async with state.condition: - lora_request = state.lora_request - return { - "lora_slot": lora_slot, - "policy_version": state.policy_version, - "update_active": state.update_active, - "admission_blocked": state.blocked, - "active_admissions": state.active_admissions, - "installed_lora_path": ( - str(lora_request.lora_path) if lora_request is not None else None - ), - "installed_lora_int_id": ( - int(lora_int_id) - if lora_request is not None - and (lora_int_id := getattr(lora_request, "lora_int_id", None)) - is not None - else None - ), - } - def lora_update_coordinator(models: Any, engine_client: Any) -> LoraUpdateCoordinator: coordinator = getattr(models, _LORA_UPDATE_COORDINATOR_FIELD, None)