From fe448971d39414e5b33cd663660bdebeb15794d4 Mon Sep 17 00:00:00 2001 From: Harry Lee Date: Mon, 5 Oct 2026 14:25:58 -0700 Subject: [PATCH 1/7] feat(replay): qualify AgentX G2 on existing mocker APIs Signed-off-by: Harry Lee --- Cargo.lock | 4 +- Cargo.toml | 2 +- benchmarks/pyproject.toml | 2 +- components/src/dynamo/replay/api.py | 13 + components/src/dynamo/replay/config.py | 50 ++ components/src/dynamo/replay/simulation.py | 41 +- .../src/dynamo/replay/tests/test_agentx_g2.py | 451 ++++++++++++++++++ .../dynamo/replay/tests/test_simulation.py | 77 +++ container/deps/requirements.aisimulate.txt | 2 +- .../dynosim-replay-cli-reference.mdx | 44 +- lib/bindings/kvbm/Cargo.lock | 4 +- lib/bindings/python/Cargo.lock | 4 +- lib/bindings/python/Cargo.toml | 2 +- lib/bindings/python/rust/llm/replay.rs | 74 ++- lib/bindings/python/src/dynamo/_core.pyi | 3 + lib/mocker/src/replay/agentic.rs | 287 +++++++++++ lib/mocker/src/replay/entrypoints.rs | 72 +++ lib/mocker/src/replay/mod.rs | 5 +- lib/mocker/src/replay/offline/entrypoints.rs | 287 ++++++++++- .../offline/extensions/kv_events/mod.rs | 95 ++++ pyproject.toml | 2 +- .../test_aisimulate_namespaces.py | 2 +- 22 files changed, 1486 insertions(+), 37 deletions(-) create mode 100644 components/src/dynamo/replay/tests/test_agentx_g2.py create mode 100644 lib/mocker/src/replay/agentic.rs diff --git a/Cargo.lock b/Cargo.lock index c30e4ec123ba..e427240b5cc7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -34,9 +34,9 @@ dependencies = [ [[package]] name = "aisimulate-core" -version = "0.13.0-dev.202610040000000065" +version = "0.13.0-dev.202610050000000066" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3b1ee30a0c0a8f270064ae7ac276391527da224ee870a3680d7426c9a749ab33" +checksum = "c9a5aa5cb79fb04ad0c8361efcc707afeef4a08680ef3d457b5a5d63eaf00ac8" dependencies = [ "anyhow", "bincode 1.3.3", diff --git a/Cargo.toml b/Cargo.toml index b47f9e7006d0..29c01f96c6a8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -59,7 +59,7 @@ keywords = ["llm", "genai", "inference", "nvidia", "distributed"] [workspace.dependencies] dynamo-custom-policy-builtin = { path = "lib/router-plugins/builtin", version = "1.6.0" } # Published separately by https://github.com/ai-dynamo/aisimulate. -aisimulate-core = { version = "=0.13.0-dev.202610040000000065" } +aisimulate-core = { version = "=0.13.0-dev.202610050000000066" } # Local crates dynamo-runtime = { path = "lib/runtime", version = "1.6.0" } diff --git a/benchmarks/pyproject.toml b/benchmarks/pyproject.toml index 88ec238391dc..6938ff370ad6 100644 --- a/benchmarks/pyproject.toml +++ b/benchmarks/pyproject.toml @@ -39,7 +39,7 @@ classifiers = [ ] dependencies = [ - "aisimulate==0.13.0.dev202610040000000065 ; python_version >= '3.11' and python_version < '3.14'", + "aisimulate==0.13.0.dev202610050000000066 ; python_version >= '3.11' and python_version < '3.14'", "aiperf==0.12.0", "matplotlib", "networkx", diff --git a/components/src/dynamo/replay/api.py b/components/src/dynamo/replay/api.py index ed08b76d0eb8..e2725b3e1fce 100644 --- a/components/src/dynamo/replay/api.py +++ b/components/src/dynamo/replay/api.py @@ -78,6 +78,9 @@ class _CommonReplayOptions(TypedDict, total=False): class _TraceReplayOptions(_CommonReplayOptions, total=False): agentic_lanes: int | None + agentic_snapshot: dict[str, Any] | None + agentic_warmup: bool + agentic_profile: dict[str, Any] | None execution_model: str | None weka_nested_timestamp_basis: Literal["auto", "absolute", "relative"] | None trace_block_size: int | None @@ -240,6 +243,9 @@ def run_trace_replay( num_decode_workers=1, replay_concurrency=None, agentic_lanes=None, + agentic_snapshot=None, + agentic_warmup=False, + agentic_profile=None, replay_mode="offline", router_mode="round_robin", arrival_speedup_ratio=1.0, @@ -276,6 +282,10 @@ def run_trace_replay( on the response path. ``None`` or ``0`` keeps synchronous updates. Offline KV-router replay only. + ``agentic_snapshot``, ``agentic_warmup`` and ``agentic_profile`` use AISimulate + workload controls on the existing offline trace path. Native validation owns + their schemas and lifecycle constraints. + Pass ``TelemetryOptions`` to enable policy-neutral sampling; omitting it leaves telemetry disabled. Callbacks and JSONL writes run synchronously on the replay loop, so their latency contributes to replay wall time. The @@ -311,6 +321,9 @@ def run_trace_replay( "num_decode_workers": num_decode_workers, "replay_concurrency": replay_concurrency, "agentic_lanes": agentic_lanes, + "agentic_snapshot": agentic_snapshot, + "agentic_warmup": agentic_warmup, + "agentic_profile": agentic_profile, "replay_mode": replay_mode, "router_mode": router_mode, "arrival_speedup_ratio": arrival_speedup_ratio, diff --git a/components/src/dynamo/replay/config.py b/components/src/dynamo/replay/config.py index 39ff38b86d11..d03ae50a4844 100644 --- a/components/src/dynamo/replay/config.py +++ b/components/src/dynamo/replay/config.py @@ -114,6 +114,26 @@ def resolve_ais_num_gpu_blocks(raw: dict[str, Any]) -> None: def lower_upstream_engine_args(payload: Mapping[str, Any]) -> dict[str, Any]: """Consume AISimulate's runner wire protocol at the Dynamo boundary.""" raw = dict(payload) + nested = raw.pop("rank", None) + if nested is not None: + if not isinstance(nested, Mapping): + raise TypeError("engine rank must be a mapping") + if set(raw) - {"dp_size", "tensor_parallel_size"}: + raise ValueError("engine arguments cannot mix nested rank with rank fields") + duplicates = set(raw) & set(nested) + if duplicates: + raise ValueError(f"duplicate nested engine arguments: {sorted(duplicates)}") + raw.update(nested) + for source, target in ( + ("backend", "engine_type"), + ("kv_transfer_bytes_per_token", "kv_bytes_per_token"), + ): + if source in raw: + if target in raw: + raise ValueError( + f"engine arguments cannot combine {source} and {target}" + ) + raw[target] = raw.pop(source) timing = raw.get("timing_model") has_custom_timing = isinstance(timing, dict) and timing.get("type") in { "fixed", @@ -144,6 +164,36 @@ def lower_upstream_engine_args(payload: Mapping[str, Any]) -> dict[str, Any]: raw.setdefault("dp_size", identity["attention_dp_size"]) if identity.get("tp_size") is not None: raw.setdefault("tensor_parallel_size", identity["tp_size"]) + host = raw.get("native_host_offload") + if ( + nested is None + and isinstance(host, Mapping) + and host.get("scope") == "cluster_shared" + and host.get("kv_layout_id") is None + ): + # Reuse AISimulate's resolved KV identity rather than inventing a + # connector-specific layout ID. Native nested ranks require an explicit ID. + from aisimulate.runner import _kv_layout_id + + backend = raw.get("engine_type", identity.get("backend", "vllm")) + if not isinstance(backend, str) or backend not in {"vllm", "sglang", "trtllm"}: + raise ValueError(f"unsupported engine_type for host offload: {backend!r}") + rank = { + "backend": backend, + "block_size": raw.get( + "block_size", {"vllm": 64, "sglang": 1, "trtllm": 32}[backend] + ), + "kv_cache_bytes_per_token": raw.get( + "kv_cache_bytes_per_token", raw.get("kv_bytes_per_token") + ), + "timing_model": {"config": raw.get("ais_perf_config", {})}, + } + raw["native_host_offload"] = { + **host, + "kv_layout_id": _kv_layout_id( + rank, identity.get("model_path"), raw.get("tensor_parallel_size", 1) + ), + } for name in tuple(raw): if name.startswith("aic_"): value = raw.pop(name) diff --git a/components/src/dynamo/replay/simulation.py b/components/src/dynamo/replay/simulation.py index 889080959f2f..024d97d1c127 100644 --- a/components/src/dynamo/replay/simulation.py +++ b/components/src/dynamo/replay/simulation.py @@ -85,6 +85,11 @@ def capabilities(self) -> RunnerCapabilities: "weka", ), supports_agentic_lanes=engine_capabilities.supports_agentic_lanes, + supports_agentic_host_offload=True, + supports_agentic_speculative_decoding=False, + supports_agentic_snapshots=True, + supports_agentic_warmup=True, + supports_agentic_profile=True, # Full AgentX runtime conformance remains a separate checkpoint. # Keep the public runner honest about the narrower integration here. supported_agentic_topologies=tuple( @@ -167,8 +172,10 @@ def run( metrics, metadata = self._normalize_report(report, output_requirements) trace_format = spec.workload.get("trace_format") agentic_lanes = spec.workload.get("agentic_lanes") - if trace_format in {"weka", "agentic_mooncake"} or ( - trace_format == "dynamo" and agentic_lanes is not None + if ( + "agentic_graph" in metadata + or trace_format in {"weka", "agentic_mooncake"} + or (trace_format == "dynamo" and agentic_lanes is not None) ): metadata.update( agentic_qualification=self.capabilities.agentic_qualification, @@ -309,9 +316,16 @@ def _execution_target_model(spec: ReplaySpec) -> str | None: candidates += [config.get("model_path"), config.get("model")] if isinstance(raw_engine_args, Mapping): candidates.append(raw_engine_args.get("aic_model_path")) - ais_config = raw_engine_args.get("ais_perf_config") - if isinstance(ais_config, Mapping): - candidates.append(ais_config.get("model")) + rank = raw_engine_args.get("rank", raw_engine_args) + if isinstance(rank, Mapping): + ais_config = rank.get("ais_perf_config") + if isinstance(ais_config, Mapping): + candidates.append(ais_config.get("model")) + timing = rank.get("timing_model") + if isinstance(timing, Mapping) and isinstance( + timing.get("config"), Mapping + ): + candidates.append(timing["config"].get("model")) for model in candidates: if isinstance(model, str) and model.strip(): return model.strip() @@ -352,6 +366,11 @@ def _run_trace( if not isinstance(trace_format, str): raise TypeError("trace workload requires a string trace_format") agentic_lanes = spec.workload.get("agentic_lanes") + agentic_options = { + name: spec.workload[name] + for name in ("agentic_snapshot", "agentic_warmup", "agentic_profile") + if name in spec.workload + } trace_block_size = spec.workload.get("trace_block_size") # Weka and Dynamo traces carry their source block size in the trace # metadata. Leave it unset so AISimulate can validate and use that @@ -368,6 +387,7 @@ def _run_trace( trace_block_size=trace_block_size, max_sim_time_ms=spec.workload.get("max_sim_time_ms"), agentic_lanes=agentic_lanes, + **agentic_options, execution_model=execution_model, extra_engine_args=self._engine_args(deployment.agg_engine_args), num_workers=deployment.num_workers, @@ -382,6 +402,7 @@ def _run_trace( trace_block_size=trace_block_size, max_sim_time_ms=spec.workload.get("max_sim_time_ms"), agentic_lanes=agentic_lanes, + **agentic_options, execution_model=execution_model, prefill_engine_args=self._engine_args(deployment.prefill_engine_args), decode_engine_args=self._engine_args(deployment.decode_engine_args), @@ -502,9 +523,15 @@ def _normalize_report( else: trace_report = dict(report) - for name in ("agentic_graph", "agentic_model_projection"): + for name in ( + "agentic_graph", + "agentic_model_projection", + "agentic_snapshots", + "agentic_phases", + "agentic_profile", + ): value = trace_report.get(name) - if isinstance(value, dict): + if isinstance(value, (dict, list)): metadata[name] = value resolved_weka_basis = trace_report.get("weka_nested_timestamp_basis") if isinstance(resolved_weka_basis, str): diff --git a/components/src/dynamo/replay/tests/test_agentx_g2.py b/components/src/dynamo/replay/tests/test_agentx_g2.py new file mode 100644 index 000000000000..42bf4659d90b --- /dev/null +++ b/components/src/dynamo/replay/tests/test_agentx_g2.py @@ -0,0 +1,451 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# ruff: noqa: E402 -- optional AISimulate preflight precedes integration imports. + +"""Real AgentX G2 qualification through Dynamo's existing replay API.""" + +from __future__ import annotations + +import copy +import json +from dataclasses import replace + +import pytest + +pytest.importorskip( + "aisimulate.sweeper", reason="AISimulate is an optional simulation dependency" +) + +from aisimulate.sweeper.provider import AdapterReplaySpec, RuntimeHookSpec +from aisimulate.sweeper.replay import ( + BackendDeploymentSpec, + ReplayOutputRequirements, + ReplaySpec, +) + +from dynamo.replay.simulation import DynamoReplayRunnerFactory + +pytestmark = [ + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.gpu_0, + pytest.mark.planner, + pytest.mark.timeout(30), +] + + +def _trace(tmp_path, trace_format): + """Authored data: fill G1, evict its prefix, then request that prefix again.""" + turns = ( + ("seed", 0, [1, 2, 3]), + ("pressure", 100, [11, 12, 13]), + ("restore", 200, [1, 2, 3]), + ("finish", 300, [21, 22, 23]), + ) + if trace_format == "dynamo": + rows = [ + { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": timestamp + 2, + "request": { + "request_id": request_id, + "request_received_ms": timestamp, + "total_time_ms": 2, + "output_tokens": 1, + "replay": { + "trace_block_size": 4, + "input_length": 9, + "input_sequence_hashes": hashes, + }, + }, + } + for request_id, timestamp, hashes in turns + ] + elif trace_format == "weka": + rows = [ + { + "id": "host-restore", + "models": ["target-model"], + "block_size": 4, + "hash_id_scope": "local", + "requests": [ + { + "t": timestamp / 1000, + "type": "s", + "model": "target-model", + "in": 9, + "out": 1, + "hash_ids": hashes, + "api_time": 0.002, + } + for _, timestamp, hashes in turns + ], + } + ] + else: + # Distinct sessions keep Mooncake-delta prompts independent. Explicit + # hashes provide the same cross-request prefix identity in both formats. + rows = [ + { + "request_id": request_id, + "session_id": request_id, + "timestamp": timestamp, + "input_length": 9, + "output_length": 1, + "hash_ids": hashes, + } + for request_id, timestamp, hashes in turns + ] + path = tmp_path / f"{trace_format}.jsonl" + path.write_text("\n".join(json.dumps(row) for row in rows) + "\n") + return path + + +def _spec(tmp_path, trace_format="weka", scope="dp_rank_local", *, disagg=False): + # Fixed timing and explicit bytes avoid models, GPUs and external services. + args = { + "engine_type": "vllm", + "worker_type": "aggregated", + "dp_size": 1, + "block_size": 4, + "num_gpu_blocks": 3, + "max_num_seqs": 1, + "max_num_batched_tokens": 16, + "kv_cache_bytes_per_token": 250000, + "enable_prefix_caching": True, + "timing_model": {"type": "fixed", "prefill_ms": 2, "decode_ms": 1}, + } + if scope is not None: + args["native_host_offload"] = { + "scope": scope, + "num_host_blocks": 8, + "d2h_bandwidth_gbps": 1, + "h2d_bandwidth_gbps": 1, + "kv_layout_id": "dynamo-agentx-g2-fixture-v1", + } + roles = ("prefill", "decode") if disagg else ("aggregated",) + workers = {role: {**copy.deepcopy(args), "worker_type": role} for role in roles} + deployment = BackendDeploymentSpec( + deployment_mode="disagg" if disagg else "agg", + backend="vllm", + backend_version="test", + agg_engine_args=workers.get("aggregated"), + prefill_engine_args=workers.get("prefill"), + decode_engine_args=workers.get("decode"), + num_workers=0 if disagg else 1, + num_prefill_workers=1 if disagg else 0, + num_decode_workers=1 if disagg else 0, + performance_model_metadata={ + role: {"config": {"model_path": "target-model"}} for role in roles + }, + ) + workload = { + "source_type": "trace", + "load_type": "trace_timestamps", + "trace_path": str(_trace(tmp_path, trace_format)), + "trace_format": trace_format, + "trace_block_size": 4, + } + if trace_format == "weka": + workload.update(agentic_lanes=1, weka_nested_timestamp_basis="absolute") + return ReplaySpec( + backend_deployment=deployment, + workload=workload, + goal={"target": "throughput"}, + adapters={ + "dynamo.router": AdapterReplaySpec( + runtime_hooks=( + RuntimeHookSpec( + provider="dynamo.router", + kind="placement_policy", + api_version=1, + config={"router_mode": "kv_router", "router_config": {}}, + ), + ) + ) + }, + ) + + +@pytest.fixture +def run_replay(): + runner = DynamoReplayRunnerFactory(trace_block_size=4).create(0) + yield lambda spec: _run(runner, spec) + runner.close() + + +def _run(runner, spec): + report = runner.run( + spec, + output_requirements=ReplayOutputRequirements( + include_raw_report=True, capture_per_request=True + ), + ) + native = report.metadata["native_report"] + # Existing Dynamo report envelope: summary contains native aggregate and + # lifecycle fields; per_request is a separate capture field. + return {**native["summary"], **native} + + +def _host_hits(report): + # HBM-only reports omit host reuse fields rather than serializing zero. + return [ + row + for row in report["per_request"] + if row.get("first_admission_host_reused_input_tokens", 0) > 0 + ] + + +@pytest.mark.parametrize( + "trace_format", ["weka", "dynamo", "mooncake", "mooncake-delta"] +) +@pytest.mark.parametrize("scope", ["dp_rank_local", "cluster_shared"]) +def test_existing_api_restores_g2_and_measures_h2d( + tmp_path, run_replay, trace_format, scope +): + spec = _spec(tmp_path, trace_format, scope) + fast = run_replay(spec) + slow_spec = copy.deepcopy(spec) + slow_spec.backend_deployment.agg_engine_args["native_host_offload"][ + "h2d_bandwidth_gbps" + ] = 0.1 + slow = run_replay(slow_spec) + hbm_spec = copy.deepcopy(spec) + hbm_spec.backend_deployment.agg_engine_args.pop("native_host_offload") + hbm = run_replay(hbm_spec) + + for report, wait_ms in ((fast, 2), (slow, 20)): + assert report["completed_requests"] == 4 + assert report["committed_prefill_tokens"] == 28 + hits = _host_hits(report) + assert len(hits) == 1 + hit = hits[0] + assert hit["first_admission_host_reused_input_tokens"] == 8 + assert hit["first_admission_g1_reused_input_tokens"] == 0 + assert hit["admission_history"][0]["host_reused_input_tokens"] == 8 + # Parent G2 routing may score HostPinned residency separately from G1. + assert hit["routing_history"] + assert hit["first_admit_ms"] - hit["arrival_time_ms"] == pytest.approx(wait_ms) + assert hit["ttft_ms"] == pytest.approx(wait_ms + 2) + if scope == "cluster_shared": + assert len(report["g2_domains"]) == 1 + assert report["g2_domains"][0]["capacity_blocks"] == 8 + assert hbm["completed_requests"] == 4 + assert not _host_hits(hbm) + assert hbm["committed_prefill_tokens"] > fast["committed_prefill_tokens"] + + +def test_1p1d_shared_pool_counts_capacity_and_first_reuse_once(tmp_path, run_replay): + spec = _spec(tmp_path, scope="cluster_shared", disagg=True) + report = run_replay(spec) + assert report["completed_requests"] == 4 + assert len(report["g2_domains"]) == 1 + assert report["g2_domains"][0]["capacity_blocks"] == 8 + hits = _host_hits(report) + assert len(hits) == 1 + assert hits[0]["first_admission_host_reused_input_tokens"] == 8 + assert {row["pool"] for row in hits[0]["admission_history"]} == { + "prefill", + "decode", + } + assert ( + sum( + row["first_admission_host_reused_input_tokens"] + for row in report["per_request"] + ) + == 8 + ) + + small = copy.deepcopy(spec) + for args in ( + small.backend_deployment.prefill_engine_args, + small.backend_deployment.decode_engine_args, + ): + args["native_host_offload"]["num_host_blocks"] = 1 + evicted = run_replay(small) + assert not _host_hits(evicted) + assert evicted["committed_prefill_tokens"] > report["committed_prefill_tokens"] + + +@pytest.mark.parametrize("scope", ["dp_rank_local", "cluster_shared"]) +def test_parent_resumes_from_g2_after_child_and_join(tmp_path, run_replay, scope): + spec = _spec(tmp_path, scope=scope) + path = tmp_path / "weka.jsonl" + play = json.loads(path.read_text()) + child_request = play["requests"][1] + play["requests"][1] = { + "t": 0.1, + "type": "subagent", + "agent_id": "pressure", + "subagent_type": "Explore", + "duration_ms": 2, + "status": "completed", + "models": ["target-model"], + "requests": [child_request], + } + path.write_text(json.dumps(play) + "\n") + report = run_replay(spec) + assert report["completed_requests"] == 4 + seed = min(report["per_request"], key=lambda row: row["arrival_time_ms"]) + children = [row for row in report["per_request"] if row["agentic"].get("parent_id")] + assert len(children) == 1 + child = children[0] + assert child["agentic"]["parent_id"] == seed["request_id"] + assert child["agentic"]["root_id"] == seed["request_id"] + assert child["agentic"]["conversation_id"] != seed["agentic"]["conversation_id"] + hits = _host_hits(report) + assert len(hits) == 1 + # Weka may split an independent outer turn into another conversation. The + # restored turn must specifically resume the original parent's conversation. + assert hits[0]["agentic"]["conversation_id"] == seed["agentic"]["conversation_id"] + assert hits[0]["agentic"]["root_id"] == seed["request_id"] + assert child["terminal_time_ms"] <= hits[0]["arrival_time_ms"] + assert hits[0]["first_admission_host_reused_input_tokens"] == 8 + assert hits[0]["first_admission_g1_reused_input_tokens"] == 0 + + +@pytest.mark.parametrize( + "field,value", + [ + ("kv_layout_id", "incompatible-layout"), + ("num_host_blocks", 16), + ("shared_h2d_bandwidth_gbps", 2), + ], +) +def test_1p1d_shared_pool_rejects_incompatible_participants( + tmp_path, run_replay, field, value +): + spec = _spec(tmp_path, scope="cluster_shared", disagg=True) + spec.backend_deployment.decode_engine_args["native_host_offload"][field] = value + # The existing binding maps native configuration errors to PyException; + # pin its diagnostic rather than accepting an arbitrary execution failure. + with pytest.raises(Exception, match="cluster_shared.*incompatible"): + run_replay(spec) + + +@pytest.mark.parametrize("trace_format", ["dynamo", "mooncake", "mooncake-delta"]) +def test_non_agentic_g2_keeps_multiple_worker_dp_support( + tmp_path, run_replay, trace_format +): + spec = _spec(tmp_path, trace_format, "cluster_shared") + spec.backend_deployment.agg_engine_args["dp_size"] = 2 + spec = replace( + spec, backend_deployment=replace(spec.backend_deployment, num_workers=2) + ) + report = run_replay(spec) + assert report["completed_requests"] == 4 + assert len(report["g2_domains"]) == 1 + assert report["g2_domains"][0]["capacity_blocks"] == 8 + assert all(row["routing_history"] for row in report["per_request"]) + + +@pytest.mark.parametrize("scope", [None, "dp_rank_local", "cluster_shared"]) +def test_snapshot_warmup_waits_for_g2_and_retains_cache(tmp_path, run_replay, scope): + spec = _spec(tmp_path, scope=scope) + if scope: + # Primer D2H lasts 200 ms, longer than the eleven fixed 2 ms passes. + spec.backend_deployment.agg_engine_args["native_host_offload"][ + "d2h_bandwidth_gbps" + ] = 0.01 + spec = replace(spec, workload={**spec.workload, "agentic_snapshot": {"seed": 5}}) + cold = run_replay(spec) + warm = run_replay(replace(spec, workload={**spec.workload, "agentic_warmup": True})) + assert cold["agentic_snapshots"] == warm["agentic_snapshots"] + assert 2 < warm["agentic_snapshots"][0]["t_star_ms"] < 100 + assert cold["completed_requests"] == warm["completed_requests"] == 3 + assert cold["committed_prefill_tokens"] == 27 + assert warm["committed_prefill_tokens"] == (19 if scope else 27) + assert sum( + row.get("first_admission_host_reused_input_tokens", 0) + for row in warm["per_request"] + ) == (8 if scope else 0) + phases = warm["agentic_phases"] + assert phases["phase"] == "profile" + assert phases["failure_reason"] is None + assert phases["lanes"][0]["warmup_completed"] == 10 + quiescent = max(row["quiescent_at_ms"] for row in phases["requests"]) + assert phases["profile_start_ms"] >= quiescent + if scope: + assert phases["profile_start_ms"] > quiescent + assert phases["profile_start_ms"] == pytest.approx(202) + assert all(row["agentic_phase"] == "profile" for row in warm["per_request"]) + + +@pytest.mark.parametrize("scope", [None, "dp_rank_local", "cluster_shared"]) +@pytest.mark.parametrize("warmup", [False, True]) +def test_profile_recycles_cache_identity_and_stops_at_cutoff( + tmp_path, run_replay, scope, warmup +): + spec = _spec(tmp_path, scope=scope) + spec = replace( + spec, + workload={ + **spec.workload, + "agentic_snapshot": {"seed": 5}, + "agentic_warmup": warmup, + "agentic_profile": {"duration_seconds": 0.7, "response_grace_seconds": 0}, + }, + ) + report = run_replay(spec) + profile = report["agentic_profile"] + assert profile["admission_closed"] + assert profile["admission_cutoff_ms"] - profile[ + "profile_start_ms" + ] == pytest.approx(700) + assert profile["finished_at_ms"] == profile["admission_cutoff_ms"] + assert profile["plays_started"] > 1 + assert profile["client_in_flight_requests"] == 0 + assert profile["unsettled_server_requests"] == 0 + assert not profile["cancel_drain_timed_out"] + assert all(0 <= row["arrival_time_ms"] < 700 for row in report["per_request"]) + cache_by_play = {} + for row in report["per_request"]: + identity = row["agentic"] + cache_by_play.setdefault(identity["play_id"], set()).add(identity["cache_id"]) + assert len(cache_by_play) > 1 + assert all(len(ids) == 1 for ids in cache_by_play.values()) + assert len({next(iter(ids)) for ids in cache_by_play.values()}) == len( + cache_by_play + ) + if warmup: + assert ( + profile["profile_start_ms"] + == report["agentic_phases"]["profile_start_ms"] + > 0 + ) + else: + assert profile["profile_start_ms"] == 0 + + +@pytest.mark.parametrize("scope", ["dp_rank_local", "cluster_shared"]) +def test_profile_cancels_incomplete_host_restore_without_draining_it( + tmp_path, run_replay, scope +): + spec = _spec(tmp_path, scope=scope) + spec.backend_deployment.agg_engine_args["native_host_offload"][ + "h2d_bandwidth_gbps" + ] = 0.001 + spec = replace( + spec, + workload={ + **spec.workload, + "agentic_snapshot": {"seed": 5}, + "agentic_warmup": True, + "agentic_profile": {"duration_seconds": 0.25, "response_grace_seconds": 0}, + }, + ) + report = run_replay(spec) + profile = report["agentic_profile"] + assert profile["canceled_requests"] == 1 + assert profile["client_in_flight_requests"] == 0 + assert profile["unsettled_server_requests"] == 0 + assert not profile["cancel_drain_timed_out"] + assert profile["finished_at_ms"] == pytest.approx(profile["profile_start_ms"] + 250) + canceled = [ + row for row in report["per_request"] if row["terminal_status"] == "canceled" + ] + assert len(canceled) == 1 + assert canceled[0]["dispatched_at_ms"] is not None + assert canceled[0]["first_admit_ms"] is None + assert canceled[0]["output_length"] == 0 diff --git a/components/src/dynamo/replay/tests/test_simulation.py b/components/src/dynamo/replay/tests/test_simulation.py index 1799bed16f57..c134e2434ad6 100644 --- a/components/src/dynamo/replay/tests/test_simulation.py +++ b/components/src/dynamo/replay/tests/test_simulation.py @@ -273,6 +273,30 @@ def fake_native_replay(_trace_files, **kwargs): {"ais_perf_config": {"model": " target-model "}}, id="engine-ais-perf-config", ), + pytest.param( + {}, + { + "timing_model": { + "type": "external", + "provider": "aic", + "config": {"model": " target-model "}, + } + }, + id="engine-timing-model", + ), + pytest.param( + {}, + { + "rank": { + "timing_model": { + "type": "external", + "provider": "aic", + "config": {"model": " target-model "}, + } + } + }, + id="nested-rank-timing-model", + ), ], ) def test_weka_runner_resolves_each_execution_target_model_source( @@ -824,3 +848,56 @@ def test_compiled_custom_timing_consumes_capacity_only_fields(timing): assert args.ais_perf_config is None report = simulation.DynamoReplayRunnerFactory().create(0).run(spec) assert report.metrics["completed_requests"] == raw["traffic"]["stop"]["requests"] + + +@pytest.mark.parametrize("nested", [False, True]) +def test_shared_g2_lowering_preserves_layout_and_link_controls(nested): + from aisimulate.runner import _materialize_engine_role + + from dynamo.replay.config import lower_upstream_engine_args + + host = { + "scope": "cluster_shared", + "num_host_blocks": 12, + "d2h_bandwidth_gbps": 7.0, + "h2d_bandwidth_gbps": 5.0, + "latency_to_first_byte_ms": 0.3, + "shared_d2h_bandwidth_gbps": 9.0, + "shared_h2d_bandwidth_gbps": 11.0, + } + raw = { + "engine_type": "vllm", + "aic_model_path": "original-model", + "block_size": 4, + "num_gpu_blocks": 3, + "tensor_parallel_size": 2, + "kv_cache_bytes_per_token": 250_000, + "timing_model": {"type": "fixed", "prefill_ms": 1.0, "decode_ms": 1.0}, + "native_host_offload": host, + } + materialized = _materialize_engine_role("vllm", "", {}, raw, "aggregated") + if nested: + raw = { + name: materialized[name] + for name in ("dp_size", "tensor_parallel_size", "rank") + } + lowered = lower_upstream_engine_args(raw) + assert lowered["native_host_offload"] == materialized["rank"]["native_host_offload"] + assert lowered["native_host_offload"].items() >= host.items() + assert lowered["tensor_parallel_size"] == 2 + assert lowered["kv_cache_bytes_per_token"] == 250_000 + assert "rank" not in lowered + + +@pytest.mark.parametrize( + "rank", + [ + {"backend": "vllm", "engine_type": "sglang"}, + {"kv_transfer_bytes_per_token": 16, "kv_bytes_per_token": 32}, + ], +) +def test_nested_engine_arguments_reject_conflicting_aliases(rank): + from dynamo.replay.config import lower_upstream_engine_args + + with pytest.raises(ValueError, match="cannot combine"): + lower_upstream_engine_args({"rank": rank}) diff --git a/container/deps/requirements.aisimulate.txt b/container/deps/requirements.aisimulate.txt index 4c018e5aab11..152e62993105 100644 --- a/container/deps/requirements.aisimulate.txt +++ b/container/deps/requirements.aisimulate.txt @@ -2,4 +2,4 @@ # SPDX-License-Identifier: Apache-2.0 # Published AISimulate wheel staged into planner and framework runtime wheelhouses. -aisimulate==0.13.0.dev202610040000000065 +aisimulate==0.13.0.dev202610050000000066 diff --git a/docs/fern/pages/reference/components/dynosim-replay-cli-reference.mdx b/docs/fern/pages/reference/components/dynosim-replay-cli-reference.mdx index cb59fe9a7e88..2cee6e9d64fa 100644 --- a/docs/fern/pages/reference/components/dynosim-replay-cli-reference.mdx +++ b/docs/fern/pages/reference/components/dynosim-replay-cli-reference.mdx @@ -177,8 +177,8 @@ AISimulate's Weka importer also changes two behaviors from the former Dynamo imp previously rejected. The fallback can affect inferred dependency edges and replay timing; it does not assert that the measured API duration was zero. -`mooncake-delta`, `agentic_mooncake`, and `weka` require aggregated engine mode. -The two agentic formats also require `trace_timestamps` and reject the virtual-time cutoff. +`mooncake-delta` requires aggregated engine mode. `agentic_mooncake` and `weka` +require `trace_timestamps`. `applied_compute_agentic` requires concurrency load. With the Dynamo stack, omit `planner` or set `planner.policy: disabled` for `mooncake-delta`, `agentic_mooncake`, and Dynamo traces that carry `agent_context` records. @@ -243,6 +243,44 @@ Agentic Mooncake v2 remains an optional materialized interchange format. Dynamo not ship a Weka parser, lowering implementation, or Weka-to-v2 converter; those source semantics and any future materialization tooling belong to AISimulate. +### AgentX host offload and lifecycle controls + +Configure G2 through `engine.workers..kv_cache.host_offload` in CLI YAML, +or `native_host_offload` in `MockEngineArgs`. Both use the existing offline replay +path for every trace format, including ordinary Dynamo, Mooncake, and +Mooncake-delta traces. No connector configuration is required. + +For AgentX graphs, G2 requires vLLM, one aggregated worker or one prefill plus +one decode worker, attention DP=1 on every role, static worker pools, and no +speculative decoding. Tensor parallelism may exceed one. These AgentX limits +also apply when a Dynamo trace contains agent context but `agentic_lanes` is +omitted. Ordinary traces retain their existing topology support; a standard +Dynamo trace with two workers and attention DP=2 is covered by the G2 tests. +G3 and online AgentX G2 replay are outside this support boundary. + +Each role can omit offload or select `dp_rank_local` or `cluster_shared`. +Shared roles must agree on layout, total capacity, and shared directional +bandwidths. `num_host_blocks` is the capacity of one shared pool and is counted +once. CLI lowering derives the layout identity; direct `MockEngineArgs` callers +must supply a compatible `kv_layout_id` for a shared pool. Host-tier events +remain distinct from GPU residency, while the KV router retains its configured +host-tier weighting. P-to-D handoff uses GPU-resident hashes. + +The Python `run_trace_replay` API also accepts `agentic_snapshot` (a mapping +containing `seed`), `agentic_warmup` (default `False`), and `agentic_profile` +(AISimulate's duration, response-grace, and cancellation-drain options). +Snapshots require positive `agentic_lanes`; warmup and profiling require a +snapshot. Profiling cannot be combined with `max_sim_time_ms`. Warmup retains +G1/G2 state across the measurement barrier. Results retain `agentic_snapshots`, +`agentic_phases`, and `agentic_profile` when those phases are enabled. + +With per-request capture, inspect `first_admission_g1_reused_input_tokens`, +`first_admission_host_reused_input_tokens`, and `admission_history` to distinguish +G1 hits, completed G2 restores, and P/D admission. HBM-only reports may omit the +host field. `g2_domains` reports shared capacity. This support is +`functional_only`; it does not establish performance equivalence with a +hardware offload recipe. + ## Engine and Adapter Rules - `engine.mode: aggregated` requires `engine.workers.aggregated`. @@ -260,7 +298,7 @@ and any future materialization tooling belong to AISimulate. - `timing.type: default` uses the AIConfigurator forward-pass model shipped in the `aisimulate` wheel. `fixed` requires both `prefill_ms` and `decode_ms`; `polynomial` selects the built-in polynomial model. -- `router.policy: kv_router` requires more than one routable worker. Set +- `router.policy: kv_router` supports one or more routable workers. Set `router.prefill_load_model.type` to `none` or `aic`. - `planner.policy` is `disabled` or `enabled`. When enabled, `planner.max_num_gpus` limits the Planner runtime budget; it is distinct from recommendation candidate constraints. diff --git a/lib/bindings/kvbm/Cargo.lock b/lib/bindings/kvbm/Cargo.lock index 254281a70e66..f034dbb95db6 100644 --- a/lib/bindings/kvbm/Cargo.lock +++ b/lib/bindings/kvbm/Cargo.lock @@ -34,9 +34,9 @@ dependencies = [ [[package]] name = "aisimulate-core" -version = "0.13.0-dev.202610040000000065" +version = "0.13.0-dev.202610050000000066" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3b1ee30a0c0a8f270064ae7ac276391527da224ee870a3680d7426c9a749ab33" +checksum = "c9a5aa5cb79fb04ad0c8361efcc707afeef4a08680ef3d457b5a5d63eaf00ac8" dependencies = [ "anyhow", "bincode 1.3.3", diff --git a/lib/bindings/python/Cargo.lock b/lib/bindings/python/Cargo.lock index a4e080270194..7961cd1d52a4 100644 --- a/lib/bindings/python/Cargo.lock +++ b/lib/bindings/python/Cargo.lock @@ -34,9 +34,9 @@ dependencies = [ [[package]] name = "aisimulate-core" -version = "0.13.0-dev.202610040000000065" +version = "0.13.0-dev.202610050000000066" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3b1ee30a0c0a8f270064ae7ac276391527da224ee870a3680d7426c9a749ab33" +checksum = "c9a5aa5cb79fb04ad0c8361efcc707afeef4a08680ef3d457b5a5d63eaf00ac8" dependencies = [ "anyhow", "bincode 1.3.3", diff --git a/lib/bindings/python/Cargo.toml b/lib/bindings/python/Cargo.toml index 11003a3de109..d18a46664312 100644 --- a/lib/bindings/python/Cargo.toml +++ b/lib/bindings/python/Cargo.toml @@ -59,7 +59,7 @@ request-trace-s3 = ["dynamo-llm/request-trace-s3"] # former AIC crate surface. The `python` feature is required for # AicEngineBuilder; this extension supplies the interpreter, so the dependency # must not enable `extension-module`. -aisimulate-core = { version = "=0.13.0-dev.202610040000000065", optional = true, features = ["python"] } +aisimulate-core = { version = "=0.13.0-dev.202610050000000066", optional = true, features = ["python"] } dynamo-runtime = { path = "../../runtime" } dynamo-mocker = { path = "../../mocker" } dynamo-parsers = { version = "9.0.0" } diff --git a/lib/bindings/python/rust/llm/replay.rs b/lib/bindings/python/rust/llm/replay.rs index 626a32187950..cf8a5a55eea8 100644 --- a/lib/bindings/python/rust/llm/replay.rs +++ b/lib/bindings/python/rust/llm/replay.rs @@ -18,8 +18,8 @@ use dynamo_mocker::loadgen::{ TraceFileFormat, WekaImportOptions, WekaNestedTimestampBasis, WekaResolvedTimestampBasis, }; use dynamo_mocker::replay::{ - ReplayArgsMode, ReplayRuntimeObservers, ReplayScalingDecision, ReplayScalingPolicy, - ReplayScalingSnapshot, ReplayTelemetryObserver, ReplayTelemetryOptions, + AgenticReplayOptions, ReplayArgsMode, ReplayRuntimeObservers, ReplayScalingDecision, + ReplayScalingPolicy, ReplayScalingSnapshot, ReplayTelemetryObserver, ReplayTelemetryOptions, ReplayTelemetrySnapshot, }; use parking_lot::Mutex; @@ -862,7 +862,7 @@ fn replay_paths_equal(left: &Path, right: &Path) -> bool { } #[pyfunction] -#[pyo3(signature = (trace_files, extra_engine_args=None, prefill_engine_args=None, decode_engine_args=None, router_config=None, ais_perf_config=None, num_workers=1, num_prefill_workers=1, num_decode_workers=1, replay_concurrency=None, replay_mode="offline", router_mode="round_robin", arrival_speedup_ratio=1.0, trace_block_size=None, trace_format="mooncake", trace_shared_prefix_ratio=0.0, trace_num_prefix_groups=0, report_jsonl_path=None, max_sim_time_ms=None, model_name=None, sla_ttft_ms=None, sla_itl_ms=None, sla_e2e_ms=None, capture_per_request=false, capture_planner_details=true, scaling_policy=None, agentic_lanes=None, execution_model=None, weka_nested_timestamp_basis=None, capture_telemetry=false, telemetry_sample_interval_ms=1_000.0, telemetry_callback=None, telemetry_jsonl_path=None, kv_event_lag_ms=None))] +#[pyo3(signature = (trace_files, extra_engine_args=None, prefill_engine_args=None, decode_engine_args=None, router_config=None, ais_perf_config=None, num_workers=1, num_prefill_workers=1, num_decode_workers=1, replay_concurrency=None, replay_mode="offline", router_mode="round_robin", arrival_speedup_ratio=1.0, trace_block_size=None, trace_format="mooncake", trace_shared_prefix_ratio=0.0, trace_num_prefix_groups=0, report_jsonl_path=None, max_sim_time_ms=None, model_name=None, sla_ttft_ms=None, sla_itl_ms=None, sla_e2e_ms=None, capture_per_request=false, capture_planner_details=true, scaling_policy=None, agentic_lanes=None, execution_model=None, weka_nested_timestamp_basis=None, capture_telemetry=false, telemetry_sample_interval_ms=1_000.0, telemetry_callback=None, telemetry_jsonl_path=None, kv_event_lag_ms=None, agentic_snapshot=None, agentic_warmup=false, agentic_profile=None))] #[allow(clippy::too_many_arguments)] pub fn run_mocker_trace_replay( py: Python<'_>, @@ -900,6 +900,9 @@ pub fn run_mocker_trace_replay( telemetry_callback: Option>, telemetry_jsonl_path: Option, kv_event_lag_ms: Option, + agentic_snapshot: Option<&Bound<'_, PyAny>>, + agentic_warmup: bool, + agentic_profile: Option<&Bound<'_, PyAny>>, ) -> PyResult { let kv_event_lag_ms = validate_kv_event_lag_ms(kv_event_lag_ms, replay_mode, router_mode)?; if telemetry_jsonl_path.as_deref().is_some_and(|path| { @@ -934,6 +937,38 @@ pub fn run_mocker_trace_replay( )?; let router_mode = parse_replay_router_mode(router_mode)?; let trace_format = parse_trace_file_format(trace_format)?; + let agentic_options = AgenticReplayOptions { + snapshot: agentic_snapshot.map(pythonize::depythonize).transpose()?, + warmup: agentic_warmup, + profile: agentic_profile.map(pythonize::depythonize).transpose()?, + arrival_speedup_ratio, + }; + agentic_options + .validate( + parse_agentic_lanes(agentic_lanes).map_err(to_pyerr)?, + max_sim_time_ms, + ) + .map_err(to_pyerr)?; + if agentic_options.has_phases() { + if replay_mode != "offline" { + return Err(PyValueError::new_err( + "agentic phases require offline replay", + )); + } + if replay_concurrency.is_some() { + return Err(PyValueError::new_err( + "agentic_snapshot requires agentic trace input without replay_concurrency", + )); + } + if !matches!( + trace_format, + TraceFileFormat::Weka | TraceFileFormat::Dynamo | TraceFileFormat::AgenticMooncake + ) { + return Err(PyValueError::new_err( + "agentic phases require an agentic trace", + )); + } + } let weka_options = parse_weka_import_options(trace_format, weka_nested_timestamp_basis).map_err(to_pyerr)?; let execution_model = match execution_model { @@ -1037,7 +1072,7 @@ pub fn run_mocker_trace_replay( } if matches!( trace_format, - TraceFileFormat::Dynamo | TraceFileFormat::Weka + TraceFileFormat::Dynamo | TraceFileFormat::Weka | TraceFileFormat::AgenticMooncake ) { let (trace, resolved_basis) = if trace_format == TraceFileFormat::Weka { let (graph, basis) = dynamo_mocker::loadgen::load_weka_agentic_graph_with_options( @@ -1046,6 +1081,15 @@ pub fn run_mocker_trace_replay( weka_options, )?; (DynamoRequestTrace::Agentic(graph), Some(basis)) + } else if trace_format == TraceFileFormat::AgenticMooncake { + ( + DynamoRequestTrace::Agentic( + dynamo_mocker::loadgen::AgenticTrace::from_agentic_mooncake( + &trace_files[0], + )?, + ), + None, + ) } else { ( DynamoRequestTrace::from_request_trace_files(&trace_files, trace_block_size)?, @@ -1063,6 +1107,7 @@ pub fn run_mocker_trace_replay( num_workers, replay_concurrency, agentic_lanes, + agentic_options, &replay_mode, arrival_speedup_ratio, router_mode, @@ -1269,6 +1314,7 @@ fn run_loaded_dynamo_request_trace( num_workers: usize, replay_concurrency: Option, agentic_lanes: Option, + agentic_options: AgenticReplayOptions, replay_mode: &str, arrival_speedup_ratio: f64, router_mode: dynamo_mocker::replay::ReplayRouterMode, @@ -1280,6 +1326,10 @@ fn run_loaded_dynamo_request_trace( ) -> anyhow::Result { match trace { DynamoRequestTrace::Standard(trace) => { + anyhow::ensure!( + !agentic_options.has_phases(), + "agentic phases require an agentic Dynamo request trace" + ); anyhow::ensure!( agentic_lanes.is_none(), "agentic_lanes requires an agentic Dynamo request trace" @@ -1379,11 +1429,13 @@ fn run_loaded_dynamo_request_trace( "agentic Dynamo request traces are not supported with replay_concurrency" ); } - let trace = trace - .normalize_starts() - .speed_up_timing(arrival_speedup_ratio)?; + let trace = if agentic_options.snapshot.is_some() { + trace + } else { + trace.normalize_starts() + }; match (args_selection, replay_mode) { - (ReplayArgsSelection::Aggregated(args), "offline") => dynamo_mocker::replay::simulate_agentic_trace_workload_with_router_mode_and_runtime_observers( + (ReplayArgsSelection::Aggregated(args), "offline") => dynamo_mocker::replay::simulate_agentic_trace_workload_with_options( *args, router_config, prefill_load_estimator, @@ -1395,19 +1447,20 @@ fn run_loaded_dynamo_request_trace( agentic_lanes, sla, take_runtime_observers(&mut scaling_policy, &mut telemetry), + agentic_options, ), (ReplayArgsSelection::Aggregated(args), "online") => dynamo_mocker::replay::simulate_agentic_trace_live_workload_with_router_mode_and_options( *args, router_config, prefill_load_estimator, - trace, + trace.speed_up_timing(arrival_speedup_ratio)?, num_workers, router_mode, record_per_request, agentic_lanes, sla, ), - (ReplayArgsSelection::Disagg(config), "offline") => dynamo_mocker::replay::simulate_agentic_trace_workload_disagg_with_router_mode_and_runtime_observers( + (ReplayArgsSelection::Disagg(config), "offline") => dynamo_mocker::replay::simulate_agentic_trace_workload_disagg_with_options( *config, router_config, prefill_load_estimator, @@ -1418,6 +1471,7 @@ fn run_loaded_dynamo_request_trace( agentic_lanes, sla, take_runtime_observers(&mut scaling_policy, &mut telemetry), + agentic_options, ), (ReplayArgsSelection::Disagg(_), "online") => anyhow::bail!( "online P/D agentic replay is not supported" diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index 54d7ec3c19a2..5f54a9e5f708 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -2556,6 +2556,9 @@ def run_mocker_trace_replay( telemetry_callback: Optional[ReplayTelemetryCallback] = None, telemetry_jsonl_path: Optional[str | os.PathLike[str]] = None, kv_event_lag_ms: Optional[float] = None, + agentic_snapshot: Optional[Dict[str, Any]] = None, + agentic_warmup: bool = False, + agentic_profile: Optional[Dict[str, Any]] = None, ) -> _OfflineReplayResult | Dict[str, Any]: """Replay mocker trace files and return the simulation report. diff --git a/lib/mocker/src/replay/agentic.rs b/lib/mocker/src/replay/agentic.rs new file mode 100644 index 000000000000..89a893a19113 --- /dev/null +++ b/lib/mocker/src/replay/agentic.rs @@ -0,0 +1,287 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +// SPDX-License-Identifier: Apache-2.0 + +use anyhow::{Context, Result, ensure}; + +use crate::loadgen::{AgenticProfileOptions, AgenticSnapshotOptions, AgenticTrace, WorkloadDriver}; + +/// Optional AgentX phases for the existing typed workload replay API. +#[derive(Clone, Debug)] +pub struct AgenticReplayOptions { + pub snapshot: Option, + pub warmup: bool, + pub profile: Option, + pub arrival_speedup_ratio: f64, +} + +impl Default for AgenticReplayOptions { + fn default() -> Self { + Self { + snapshot: None, + warmup: false, + profile: None, + arrival_speedup_ratio: 1.0, + } + } +} + +impl AgenticReplayOptions { + pub fn has_phases(&self) -> bool { + self.snapshot.is_some() || self.warmup || self.profile.is_some() + } + + pub fn validate(&self, lanes: Option, max_sim_time_ms: Option) -> Result<()> { + ensure!( + self.arrival_speedup_ratio.is_finite() && self.arrival_speedup_ratio > 0.0, + "arrival_speedup_ratio must be finite and greater than zero" + ); + if let Some(profile) = &self.profile { + profile.validate()?; + ensure!( + self.snapshot.is_some(), + "agentic_profile requires agentic_snapshot" + ); + ensure!( + max_sim_time_ms.is_none(), + "agentic_profile cannot be combined with max_sim_time_ms" + ); + } + ensure!( + !self.warmup || self.snapshot.is_some(), + "agentic_warmup requires agentic_snapshot" + ); + ensure!( + self.snapshot.is_none() || lanes.is_some_and(|lanes| lanes > 0), + "agentic_snapshot requires positive agentic_lanes" + ); + Ok(()) + } + + pub(crate) fn into_driver( + self, + trace: AgenticTrace, + block_size: usize, + include_replay_hashes: bool, + lanes: Option, + max_sim_time_ms: Option, + ) -> Result { + self.validate(lanes, max_sim_time_ms)?; + let Some(snapshot) = self.snapshot else { + return trace + .speed_up_timing(self.arrival_speedup_ratio)? + .into_trace_driver_with_options(block_size, include_replay_hashes, lanes); + }; + // Sample recorded source time first; speedup applies only to the saved + // frontier's remaining timers, just as it does in AISimulate's runner. + let prepared = trace.prepare_snapshots( + lanes.context("agentic_snapshot requires positive agentic_lanes")?, + snapshot, + )?; + let mut driver = if self.warmup { + WorkloadDriver::new_agentic_warmup( + prepared, + block_size, + include_replay_hashes, + self.arrival_speedup_ratio, + )? + } else { + WorkloadDriver::new_agentic_snapshots( + prepared, + block_size, + include_replay_hashes, + self.arrival_speedup_ratio, + )? + }; + if let Some(profile) = self.profile { + driver.enable_agentic_profile(profile)?; + } + Ok(driver) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::loadgen::AgenticGraphBuilder; + use serde_json::json; + + fn recorded_graph() -> AgenticTrace { + let mut graph = AgenticGraphBuilder::new( + serde_json::from_value(json!({ + "schema": "dynamo.agentic_mooncake", + "version": 2, + "block_size": 4, + "hash_id_scope": "local", + "source": {"format": "self-authored-test", "digest": "typed-agentic-options"} + })) + .unwrap(), + ) + .unwrap(); + for (id, start) in [("first", 100.0), ("second", 220.0)] { + graph + .push( + serde_json::from_value(json!({ + "request_id": id, + "play_id": "play", + "session_id": id, + "model": "fixture-model", + "input_length": 8, + "output_length": 1, + "hash_ids": [1, 2], + "not_before_ms": start, + "recorded_api_time_ms": 10.0, + "dependencies": [] + })) + .unwrap(), + ) + .unwrap(); + } + graph.finish().unwrap() + } + + #[test] + fn snapshot_sampling_keeps_source_clock_and_requested_lane_count() { + let build = |speedup| { + AgenticReplayOptions { + snapshot: Some(AgenticSnapshotOptions { seed: 7 }), + arrival_speedup_ratio: speedup, + ..Default::default() + } + .into_driver(recorded_graph(), 4, true, Some(3), None) + .unwrap() + }; + let regular = build(1.0); + let fast = build(4.0); + let snapshots = regular.agentic_snapshot_evidence().unwrap(); + assert_eq!( + snapshots.len(), + 3, + "snapshot lanes may exceed source play count" + ); + assert_eq!( + snapshots, + fast.agentic_snapshot_evidence().unwrap(), + "speedup must not change recorded-time sampling or cache identity" + ); + assert!(snapshots.iter().all(|snapshot| snapshot.t_star_ms >= 100.0)); + } + + #[test] + fn finite_workload_preserves_timestamps_and_scales_timers_once() { + for (lanes, first_ms, second_ms) in [(None, 100.0, 220.0), (Some(1), 0.0, 120.0)] { + let mut regular = AgenticReplayOptions::default() + .into_driver(recorded_graph(), 4, true, lanes, None) + .unwrap(); + let mut fast = AgenticReplayOptions { + arrival_speedup_ratio: 4.0, + ..Default::default() + } + .into_driver(recorded_graph(), 4, true, lanes, None) + .unwrap(); + // Finite lanes start a play at zero; unbounded replay keeps its + // timestamps. Both paths scale the next independent root once. + assert_eq!(regular.next_ready_time_ms(), Some(first_ms)); + assert_eq!(fast.next_ready_time_ms(), Some(first_ms / 4.0)); + assert_eq!(regular.pop_ready(first_ms, 1).len(), 1); + assert_eq!(fast.pop_ready(first_ms / 4.0, 1).len(), 1); + assert_eq!(regular.next_ready_time_ms(), Some(second_ms)); + assert_eq!(fast.next_ready_time_ms(), Some(second_ms / 4.0)); + assert!(regular.agentic_snapshot_evidence().is_none()); + assert!(regular.agentic_profile_report().is_none()); + } + } + + #[test] + fn warmup_defers_profile_clock_until_runtime_barrier() { + let options = AgenticReplayOptions { + snapshot: Some(AgenticSnapshotOptions { seed: 7 }), + warmup: true, + profile: Some(AgenticProfileOptions { + duration_seconds: 0.5, + response_grace_seconds: 0.2, + cancel_drain_seconds: 0.1, + ..Default::default() + }), + ..Default::default() + }; + let warm = options + .clone() + .into_driver(recorded_graph(), 4, true, Some(1), None) + .unwrap(); + assert!(warm.is_agentic_preparing()); + assert_eq!( + warm.agentic_profile_deadlines(), + None, + "creating a driver must not bypass the native preparation barrier" + ); + let cold = AgenticReplayOptions { + warmup: false, + ..options + } + .into_driver(recorded_graph(), 4, true, Some(1), None) + .unwrap(); + assert!(!cold.is_agentic_preparing()); + assert_eq!( + cold.agentic_profile_deadlines(), + Some((500.0, 700.0, 800.0)) + ); + } + + #[test] + fn invalid_phase_combinations_fail_before_execution() { + let cases = [ + ( + AgenticReplayOptions { + warmup: true, + ..Default::default() + }, + Some(1), + None, + "agentic_warmup requires agentic_snapshot", + ), + ( + AgenticReplayOptions { + profile: Some(AgenticProfileOptions::default()), + ..Default::default() + }, + Some(1), + None, + "agentic_profile requires agentic_snapshot", + ), + ( + AgenticReplayOptions { + snapshot: Some(AgenticSnapshotOptions { seed: 7 }), + ..Default::default() + }, + None, + None, + "agentic_snapshot requires positive agentic_lanes", + ), + ( + AgenticReplayOptions { + snapshot: Some(AgenticSnapshotOptions { seed: 7 }), + profile: Some(AgenticProfileOptions::default()), + ..Default::default() + }, + Some(1), + Some(100.0), + "agentic_profile cannot be combined with max_sim_time_ms", + ), + ( + AgenticReplayOptions { + arrival_speedup_ratio: f64::NAN, + ..Default::default() + }, + None, + None, + "arrival_speedup_ratio must be finite and greater than zero", + ), + ]; + for (options, lanes, time_limit, expected) in cases { + assert_eq!( + options.validate(lanes, time_limit).unwrap_err().to_string(), + expected + ); + } + } +} diff --git a/lib/mocker/src/replay/entrypoints.rs b/lib/mocker/src/replay/entrypoints.rs index ec7e1c0fa0d7..bc7f98d5a141 100644 --- a/lib/mocker/src/replay/entrypoints.rs +++ b/lib/mocker/src/replay/entrypoints.rs @@ -568,6 +568,7 @@ pub fn simulate_trace_file_with_router_mode_and_format_and_runtime_observers( record_per_request, max_sim_time_ms, agentic_lanes, + super::AgenticReplayOptions::default(), sla, scaling_policy, telemetry, @@ -724,6 +725,7 @@ pub fn simulate_trace_file_disagg_with_router_mode_and_format_and_runtime_observ record_per_request, max_sim_time_ms, agentic_lanes, + super::AgenticReplayOptions::default(), sla, scaling_policy, telemetry, @@ -873,6 +875,10 @@ pub fn simulate_trace_live_file_with_router_mode_and_format_and_options( let args = args.normalized()?; validate_online_replay_args(&args, num_workers)?; if is_agentic_trace_format(trace_format) { + anyhow::ensure!( + args.native_host_offload.is_none(), + "agentic host offload requires offline backend=vllm" + ); let trace = load_agentic_trace_from_file( trace_path, trace_block_size, @@ -2242,6 +2248,37 @@ pub fn simulate_agentic_trace_workload_with_router_mode_and_runtime_observers( agentic_lanes: Option, sla: SlaThresholds, observers: super::ReplayRuntimeObservers, +) -> Result { + simulate_agentic_trace_workload_with_options( + args, + router_config, + prefill_load_estimator, + trace, + num_workers, + router_mode, + record_per_request, + max_sim_time_ms, + agentic_lanes, + sla, + observers, + super::AgenticReplayOptions::default(), + ) +} + +#[allow(clippy::too_many_arguments)] +pub fn simulate_agentic_trace_workload_with_options( + args: MockEngineArgs, + router_config: Option, + prefill_load_estimator: Option, + trace: AgenticTrace, + num_workers: usize, + router_mode: ReplayRouterMode, + record_per_request: bool, + max_sim_time_ms: Option, + agentic_lanes: Option, + sla: SlaThresholds, + observers: super::ReplayRuntimeObservers, + agentic_options: super::AgenticReplayOptions, ) -> Result { let super::ReplayRuntimeObservers { scaling_policy, @@ -2259,6 +2296,7 @@ pub fn simulate_agentic_trace_workload_with_router_mode_and_runtime_observers( record_per_request, max_sim_time_ms, agentic_lanes, + agentic_options, sla, scaling_policy, telemetry, @@ -2335,6 +2373,35 @@ pub fn simulate_agentic_trace_workload_disagg_with_router_mode_and_runtime_obser agentic_lanes: Option, sla: SlaThresholds, observers: super::ReplayRuntimeObservers, +) -> Result { + simulate_agentic_trace_workload_disagg_with_options( + config, + router_config, + prefill_load_estimator, + trace, + router_mode, + record_per_request, + max_sim_time_ms, + agentic_lanes, + sla, + observers, + super::AgenticReplayOptions::default(), + ) +} + +#[allow(clippy::too_many_arguments)] +pub fn simulate_agentic_trace_workload_disagg_with_options( + config: OfflineDisaggReplayConfig, + router_config: Option, + prefill_load_estimator: Option, + trace: AgenticTrace, + router_mode: ReplayRouterMode, + record_per_request: bool, + max_sim_time_ms: Option, + agentic_lanes: Option, + sla: SlaThresholds, + observers: super::ReplayRuntimeObservers, + agentic_options: super::AgenticReplayOptions, ) -> Result { let super::ReplayRuntimeObservers { scaling_policy, @@ -2351,6 +2418,7 @@ pub fn simulate_agentic_trace_workload_disagg_with_router_mode_and_runtime_obser record_per_request, max_sim_time_ms, agentic_lanes, + agentic_options, sla, scaling_policy, telemetry, @@ -2369,6 +2437,10 @@ pub fn simulate_agentic_trace_live_workload_with_router_mode_and_options( agentic_lanes: Option, sla: SlaThresholds, ) -> Result { + anyhow::ensure!( + args.native_host_offload.is_none(), + "agentic host offload requires offline backend=vllm" + ); let args = args.normalized()?; validate_online_replay_args(&args, num_workers)?; online::simulate_agentic_trace_workload( diff --git a/lib/mocker/src/replay/mod.rs b/lib/mocker/src/replay/mod.rs index 8bfd9ca8fc91..f78b04cc0a96 100644 --- a/lib/mocker/src/replay/mod.rs +++ b/lib/mocker/src/replay/mod.rs @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +mod agentic; mod artifacts; mod entrypoints; pub(crate) mod offline; @@ -14,6 +15,7 @@ use std::sync::Arc; use crate::common::protocols::{DirectRequest, MockEngineArgs}; use dynamo_kv_router::PrefillLoadEstimator; +pub use agentic::AgenticReplayOptions; /// Backward-compatible Dynamo Mocker name for [`aisimulate_core::ReplayReport`]. pub use aisimulate_core::ReplayReport as TraceSimulationReport; pub(crate) use aisimulate_core::replay::TraceCollector; @@ -92,10 +94,11 @@ pub use entrypoints::{ ReplayKvEventVisibility, generate_trace_worker_artifacts_offline, generate_trace_worker_artifacts_offline_with_kv_event_visibility, simulate_agentic_trace_live_workload_with_router_mode_and_options, + simulate_agentic_trace_workload_disagg_with_options, simulate_agentic_trace_workload_disagg_with_router_mode, simulate_agentic_trace_workload_disagg_with_router_mode_and_runtime_observers, simulate_agentic_trace_workload_disagg_with_router_mode_and_telemetry, - simulate_agentic_trace_workload_with_router_mode, + simulate_agentic_trace_workload_with_options, simulate_agentic_trace_workload_with_router_mode, simulate_agentic_trace_workload_with_router_mode_and_runtime_observers, simulate_agentic_trace_workload_with_router_mode_and_telemetry, simulate_concurrency_file, simulate_concurrency_file_disagg_with_router_mode, diff --git a/lib/mocker/src/replay/offline/entrypoints.rs b/lib/mocker/src/replay/offline/entrypoints.rs index 89366d835c2b..3db0deecc90c 100644 --- a/lib/mocker/src/replay/offline/entrypoints.rs +++ b/lib/mocker/src/replay/offline/entrypoints.rs @@ -44,6 +44,58 @@ fn worker_pool(initial_workers: usize, args: &MockEngineArgs) -> WorkerPoolSpec } } +fn validate_agentic_host_offload( + input: &ReplayRuntimeInput, + roles: &[(usize, &MockEngineArgs)], + scaling_enabled: bool, +) -> Result<()> { + let ReplayRuntimeInput::Workload(driver) = input else { + return Ok(()); + }; + if !driver.is_agentic() + || !roles + .iter() + .any(|(_, args)| args.native_host_offload.is_some()) + { + return Ok(()); + } + + // These entrypoints are offline. Apply the AgentX G2 deployment contract + // to every active role, including a P/D role whose own G2 is disabled. + anyhow::ensure!( + roles.iter().all(|(workers, _)| *workers == 1), + "agentic host offload requires one aggregated worker or one prefill and one decode worker" + ); + anyhow::ensure!( + !scaling_enabled, + "agentic host offload requires static worker pools without a scaling policy" + ); + for (_, args) in roles { + anyhow::ensure!( + args.engine_type == EngineType::Vllm + && args + .ais_backend + .as_deref() + .is_none_or(|backend| backend == "vllm"), + "agentic host offload requires backend=vllm on every role" + ); + anyhow::ensure!( + args.dp_size == 1 && args.ais_attention_dp_size.is_none_or(|dp| dp == 1), + "agentic host offload requires attention DP=1 on every role" + ); + anyhow::ensure!( + args.ais_nextn.is_none() + && args.ais_perf_config.as_ref().is_none_or(|config| { + config + .get("speculation") + .is_none_or(serde_json::Value::is_null) + }), + "agentic replay requires speculative decoding disabled" + ); + } + Ok(()) +} + fn with_telemetry( replayer: Replayer, telemetry: Option, @@ -147,6 +199,7 @@ fn run_aggregated_with_capture_options( telemetry: Option, ) -> Result { let args = args.normalized()?; + validate_agentic_host_offload(&input, &[(num_workers, &args)], scaling_policy.is_some())?; let (engine, factory) = aggregated_replay_setup(&args)?; let spec = replay_spec( ReplayTopology::Aggregated { @@ -242,6 +295,14 @@ fn run_disaggregated_with_capture_options( telemetry: Option, ) -> Result { let config = config.normalized()?; + validate_agentic_host_offload( + &input, + &[ + (config.num_prefill_workers, &config.prefill_args), + (config.num_decode_workers, &config.decode_args), + ], + scaling_policy.is_some(), + )?; let (engine, factory) = disaggregated_replay_setup(&config.prefill_args, &config.decode_args)?; let spec = replay_spec( ReplayTopology::Disaggregated { @@ -581,16 +642,23 @@ pub(crate) fn simulate_agentic_trace_workload( record_per_request: bool, max_sim_time_ms: Option, agentic_lanes: Option, + agentic_options: crate::replay::AgenticReplayOptions, sla: SlaThresholds, scaling_policy: Option>, telemetry: Option, ) -> Result { let args = args.normalized()?; - let agentic_lanes = effective_agentic_lanes(agentic_lanes, trace.play_count()); - let driver = trace.into_trace_driver_with_options( + let agentic_lanes = if agentic_options.snapshot.is_some() { + agentic_lanes + } else { + effective_agentic_lanes(agentic_lanes, trace.play_count()) + }; + let driver = agentic_options.into_driver( + trace, args.block_size, router_mode == ReplayRouterMode::KvRouter, agentic_lanes, + max_sim_time_ms, )?; run_aggregated( args, @@ -618,16 +686,23 @@ pub(crate) fn simulate_agentic_trace_workload_disagg( record_per_request: bool, max_sim_time_ms: Option, agentic_lanes: Option, + agentic_options: crate::replay::AgenticReplayOptions, sla: SlaThresholds, scaling_policy: Option>, telemetry: Option, ) -> Result { let config = config.normalized()?; - let agentic_lanes = effective_agentic_lanes(agentic_lanes, trace.play_count()); - let driver = trace.into_trace_driver_with_options( + let agentic_lanes = if agentic_options.snapshot.is_some() { + agentic_lanes + } else { + effective_agentic_lanes(agentic_lanes, trace.play_count()) + }; + let driver = agentic_options.into_driver( + trace, config.prefill_args.block_size, router_mode == ReplayRouterMode::KvRouter, agentic_lanes, + max_sim_time_ms, )?; run_disaggregated( config, @@ -813,3 +888,207 @@ pub(crate) fn simulate_concurrency_workload_disagg_with_scaling_policy( telemetry, ) } + +#[cfg(test)] +mod agentic_host_offload_tests { + use super::*; + use crate::loadgen::{ + AGENTIC_MOONCAKE_SCHEMA, AGENTIC_MOONCAKE_VERSION, AgenticHashIdScope, + AgenticMooncakeHeader, AgenticMooncakeRow, AgenticSourceProvenance, + }; + + fn agentic_input() -> ReplayRuntimeInput { + let trace = AgenticTrace::from_agentic_mooncake_rows( + AgenticMooncakeHeader { + schema: AGENTIC_MOONCAKE_SCHEMA.to_string(), + version: AGENTIC_MOONCAKE_VERSION, + block_size: 4, + hash_id_scope: AgenticHashIdScope::Local, + source: AgenticSourceProvenance { + format: "test".to_string(), + digest: "agentic-g2-deployment-guard".to_string(), + }, + }, + vec![AgenticMooncakeRow { + request_id: "request".to_string(), + play_id: "play".to_string(), + session_id: "session".to_string(), + model: "model".to_string(), + input_length: Some(4), + output_length: Some(1), + hash_ids: Some(vec![1]), + ..Default::default() + }], + ) + .unwrap(); + ReplayRuntimeInput::Workload(trace.into_trace_driver_with_options(4, true, None).unwrap()) + } + + fn args(scope: Option<&str>) -> MockEngineArgs { + let mut args = MockEngineArgs::builder().build().unwrap(); + args.native_host_offload = scope.map(|scope| { + serde_json::from_value(serde_json::json!({ + "scope": scope, + "num_host_blocks": 16, + "kv_layout_id": "agentic-g2-test", + })) + .unwrap() + }); + args.normalized().unwrap() + } + + #[test] + fn accepts_single_worker_and_independent_pd_g2_scopes() { + let input = agentic_input(); + let scopes = [None, Some("dp_rank_local"), Some("cluster_shared")]; + for prefill_scope in scopes { + let mut prefill = args(prefill_scope); + prefill.ais_tp_size = Some(4); + validate_agentic_host_offload(&input, &[(1, &prefill)], false).unwrap(); + for decode_scope in scopes { + let decode = args(decode_scope); + validate_agentic_host_offload(&input, &[(1, &prefill), (1, &decode)], false) + .unwrap(); + } + } + } + + #[test] + fn rejects_unsupported_deployment_on_every_active_role() { + let input = agentic_input(); + let offloaded = args(Some("dp_rank_local")); + // A role without G2 is still part of the restricted AgentX deployment. + let mut invalid_roles = Vec::new(); + let mut invalid = args(None); + invalid.dp_size = 2; + invalid_roles.push(( + invalid, + "agentic host offload requires attention DP=1 on every role", + )); + let mut invalid = args(None); + invalid.ais_attention_dp_size = Some(2); + invalid_roles.push(( + invalid, + "agentic host offload requires attention DP=1 on every role", + )); + for backend in [EngineType::Sglang, EngineType::Trtllm] { + let mut invalid = args(None); + invalid.engine_type = backend; + invalid_roles.push(( + invalid, + "agentic host offload requires backend=vllm on every role", + )); + } + let mut invalid = args(None); + invalid.ais_nextn = Some(1); + invalid_roles.push(( + invalid, + "agentic replay requires speculative decoding disabled", + )); + for (invalid, expected) in invalid_roles { + for roles in [ + vec![(1, &offloaded), (1, &invalid)], + vec![(1, &invalid), (1, &offloaded)], + ] { + let error = validate_agentic_host_offload(&input, &roles, false).unwrap_err(); + assert_eq!(error.to_string(), expected); + } + } + for roles in [ + vec![(2, &offloaded)], + vec![(2, &offloaded), (1, &offloaded)], + vec![(1, &offloaded), (2, &offloaded)], + ] { + let error = validate_agentic_host_offload(&input, &roles, false).unwrap_err(); + assert_eq!( + error.to_string(), + "agentic host offload requires one aggregated worker or one prefill and one decode worker" + ); + } + let error = validate_agentic_host_offload(&input, &[(1, &offloaded)], true).unwrap_err(); + assert_eq!( + error.to_string(), + "agentic host offload requires static worker pools without a scaling policy" + ); + } + + #[test] + fn native_execution_entrypoints_apply_the_agentic_g2_guard() { + let error = run_aggregated( + args(Some("cluster_shared")), + None, + None, + agentic_input(), + 2, + None, + ReplayRouterMode::RoundRobin, + false, + None, + SlaThresholds::default(), + None, + None, + ) + .unwrap_err(); + assert_eq!( + error.to_string(), + "agentic host offload requires one aggregated worker or one prefill and one decode worker" + ); + + let mut decode = args(None); + decode.dp_size = 2; + let error = run_disaggregated( + OfflineDisaggReplayConfig { + prefill_args: args(Some("dp_rank_local")), + decode_args: decode, + num_prefill_workers: 1, + num_decode_workers: 1, + }, + None, + None, + agentic_input(), + None, + ReplayRouterMode::RoundRobin, + false, + None, + SlaThresholds::default(), + None, + None, + ) + .unwrap_err(); + assert_eq!( + error.to_string(), + "agentic host offload requires attention DP=1 on every role" + ); + } + + #[test] + fn leaves_hbm_and_non_agentic_g2_deployments_unchanged() { + let mut hbm = args(None); + hbm.engine_type = EngineType::Sglang; + hbm.dp_size = 2; + validate_agentic_host_offload(&agentic_input(), &[(2, &hbm)], true).unwrap(); + + let mut offloaded = args(Some("cluster_shared")); + offloaded.dp_size = 2; + validate_agentic_host_offload( + &ReplayRuntimeInput::Requests(VecDeque::new()), + &[(2, &offloaded)], + true, + ) + .unwrap(); + let trace = Trace::from_mooncake_rows( + vec![crate::loadgen::MooncakeRow { + input_length: Some(4), + output_length: Some(1), + hash_ids: Some(vec![1]), + timestamp: Some(0.0), + ..Default::default() + }], + 4, + ) + .unwrap(); + let input = + ReplayRuntimeInput::Workload(trace.into_trace_driver_with_block_size(4).unwrap()); + validate_agentic_host_offload(&input, &[(2, &offloaded)], true).unwrap(); + } +} diff --git a/lib/mocker/src/replay/offline/extensions/kv_events/mod.rs b/lib/mocker/src/replay/offline/extensions/kv_events/mod.rs index eb8ac2ecf1ac..1ea890ef602d 100644 --- a/lib/mocker/src/replay/offline/extensions/kv_events/mod.rs +++ b/lib/mocker/src/replay/offline/extensions/kv_events/mod.rs @@ -76,6 +76,10 @@ impl ReplayEngineObservation for RouterEventObservation { batch .0 .iter() + // Handoff evidence describes GPU-resident blocks. Host stores + // remain in the batch for G2 routing, but are not completed P/D + // handoff destinations. + .filter(|event| event.storage_tier.is_gpu()) .flat_map(|event| match &event.event.data { dynamo_kv_router::protocols::KvCacheEventData::Stored(store) => { store.blocks.as_slice() @@ -249,3 +253,94 @@ pub(in crate::replay) fn generate_trace_worker_artifacts_with_visibility( .collect(), }) } + +#[cfg(test)] +mod tests { + use super::*; + use aisimulate_core::engine::{KvBlock, KvEventData, KvEventTier, StoredBlocks}; + + fn stored(tier: KvEventTier, event_id: u64) -> KvEvent { + KvEvent { + event_id, + dp_rank: 0, + tier, + data: KvEventData::Stored(StoredBlocks { + parent_hash: None, + start_position: Some(0), + blocks: vec![KvBlock { + block_hash: 71, + tokens_hash: 7, + token_ids: None, + }], + }), + } + } + + fn removed(tier: KvEventTier, event_id: u64) -> KvEvent { + KvEvent { + event_id, + dp_rank: 0, + tier, + data: KvEventData::Removed { + block_hashes: vec![71], + }, + } + } + + #[test] + fn g2_events_remain_routable_but_do_not_prove_gpu_handoff_residency() { + for stage in [WorkerStage::Aggregated, WorkerStage::Prefill] { + let transitions = [ + ( + stored(KvEventTier::Device, 1), + StorageTier::Device, + vec![7], + "G1 store", + ), + ( + removed(KvEventTier::Device, 2), + StorageTier::Device, + vec![], + "G1 eviction", + ), + ( + stored(KvEventTier::HostPinned, 3), + StorageTier::HostPinned, + vec![], + "G2 store", + ), + ( + stored(KvEventTier::Device, 4), + StorageTier::Device, + vec![7], + "G2 to G1 restore", + ), + ( + removed(KvEventTier::HostPinned, 5), + StorageTier::HostPinned, + vec![], + "G2 eviction", + ), + ]; + for (event, expected_tier, expected_hashes, transition) in transitions { + let batch = RouterEventObservation::observe_engine_events(stage, 3, 0, vec![event]); + assert_eq!(batch.0.len(), 1, "{transition}: retain the router event"); + assert_eq!( + batch.0[0].storage_tier, expected_tier, + "{transition}: preserve tier" + ); + assert_eq!(batch.0[0].worker_id, 3, "{transition}: preserve worker"); + assert_eq!( + RouterEventObservation::stored_hashes(&batch), + expected_hashes, + "{transition}: only a Device store proves handoff residency" + ); + assert_eq!( + RouterEventObservation::kv_ingest_event_count(&batch), + Some(1), + "{transition}: retain KV ingest evidence for either tier" + ); + } + } + } +} diff --git a/pyproject.toml b/pyproject.toml index 1239fe60fe17..0df99b0f3878 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,7 @@ license-files = ["LICENSE"] requires-python = ">=3.10" dependencies = [ "ai-dynamo-runtime==1.6.0", - "aisimulate==0.13.0.dev202610040000000065 ; python_version >= '3.11' and python_version < '3.14'", + "aisimulate==0.13.0.dev202610050000000066 ; python_version >= '3.11' and python_version < '3.14'", "aiohttp>=3.14.3,<4.0", "transformers>=4.56.0", "kubernetes>=32.0.1,<33.0.0", diff --git a/tests/dependencies/test_aisimulate_namespaces.py b/tests/dependencies/test_aisimulate_namespaces.py index 9c7dd35bc156..cba617b3a0c7 100644 --- a/tests/dependencies/test_aisimulate_namespaces.py +++ b/tests/dependencies/test_aisimulate_namespaces.py @@ -139,7 +139,7 @@ def test_no_manifest_installs_retired_aic_distributions() -> None: assert features["ais-forward-pass"] == ["dep:aisimulate-core"] assert "aic-forward-pass" not in features assert dependencies["aisimulate-core"] == { - "version": "=0.13.0-dev.202610040000000065", + "version": "=0.13.0-dev.202610050000000066", "optional": True, "features": ["python"], } From 48b68e75be0cc9351b6f99b399a65da7a8374cb9 Mon Sep 17 00:00:00 2001 From: Harry Lee Date: Mon, 5 Oct 2026 14:29:02 -0700 Subject: [PATCH 2/7] fix(replay): complete AgentX module copyright header Signed-off-by: Harry Lee --- lib/mocker/src/replay/agentic.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/mocker/src/replay/agentic.rs b/lib/mocker/src/replay/agentic.rs index 89a893a19113..df0f208415b6 100644 --- a/lib/mocker/src/replay/agentic.rs +++ b/lib/mocker/src/replay/agentic.rs @@ -1,4 +1,4 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 use anyhow::{Context, Result, ensure}; From b784a88005c7aabb25760341e2a6c51a255a2d16 Mon Sep 17 00:00:00 2001 From: Harry Lee Date: Tue, 6 Oct 2026 14:03:23 -0700 Subject: [PATCH 3/7] build(ais): use published post-AgentX merge release Signed-off-by: Harry Lee --- Cargo.lock | 4 ++-- Cargo.toml | 2 +- benchmarks/pyproject.toml | 2 +- container/deps/requirements.aisimulate.txt | 2 +- lib/bindings/kvbm/Cargo.lock | 4 ++-- lib/bindings/python/Cargo.lock | 4 ++-- lib/bindings/python/Cargo.toml | 2 +- pyproject.toml | 2 +- tests/dependencies/test_aisimulate_namespaces.py | 2 +- 9 files changed, 12 insertions(+), 12 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e427240b5cc7..735b1847ab19 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -34,9 +34,9 @@ dependencies = [ [[package]] name = "aisimulate-core" -version = "0.13.0-dev.202610050000000066" +version = "0.13.0-dev.202610060000000067" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c9a5aa5cb79fb04ad0c8361efcc707afeef4a08680ef3d457b5a5d63eaf00ac8" +checksum = "baa4fe3c1314053a9e439cd7a7595aec5fb2b22af0ae43d48687179846e7ed32" dependencies = [ "anyhow", "bincode 1.3.3", diff --git a/Cargo.toml b/Cargo.toml index 29c01f96c6a8..1f679eb7a541 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -59,7 +59,7 @@ keywords = ["llm", "genai", "inference", "nvidia", "distributed"] [workspace.dependencies] dynamo-custom-policy-builtin = { path = "lib/router-plugins/builtin", version = "1.6.0" } # Published separately by https://github.com/ai-dynamo/aisimulate. -aisimulate-core = { version = "=0.13.0-dev.202610050000000066" } +aisimulate-core = { version = "=0.13.0-dev.202610060000000067" } # Local crates dynamo-runtime = { path = "lib/runtime", version = "1.6.0" } diff --git a/benchmarks/pyproject.toml b/benchmarks/pyproject.toml index 6938ff370ad6..560b2c95f7d0 100644 --- a/benchmarks/pyproject.toml +++ b/benchmarks/pyproject.toml @@ -39,7 +39,7 @@ classifiers = [ ] dependencies = [ - "aisimulate==0.13.0.dev202610050000000066 ; python_version >= '3.11' and python_version < '3.14'", + "aisimulate==0.13.0.dev202610060000000067 ; python_version >= '3.11' and python_version < '3.14'", "aiperf==0.12.0", "matplotlib", "networkx", diff --git a/container/deps/requirements.aisimulate.txt b/container/deps/requirements.aisimulate.txt index 152e62993105..70dfa9ca9400 100644 --- a/container/deps/requirements.aisimulate.txt +++ b/container/deps/requirements.aisimulate.txt @@ -2,4 +2,4 @@ # SPDX-License-Identifier: Apache-2.0 # Published AISimulate wheel staged into planner and framework runtime wheelhouses. -aisimulate==0.13.0.dev202610050000000066 +aisimulate==0.13.0.dev202610060000000067 diff --git a/lib/bindings/kvbm/Cargo.lock b/lib/bindings/kvbm/Cargo.lock index f034dbb95db6..6fa5424d9698 100644 --- a/lib/bindings/kvbm/Cargo.lock +++ b/lib/bindings/kvbm/Cargo.lock @@ -34,9 +34,9 @@ dependencies = [ [[package]] name = "aisimulate-core" -version = "0.13.0-dev.202610050000000066" +version = "0.13.0-dev.202610060000000067" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c9a5aa5cb79fb04ad0c8361efcc707afeef4a08680ef3d457b5a5d63eaf00ac8" +checksum = "baa4fe3c1314053a9e439cd7a7595aec5fb2b22af0ae43d48687179846e7ed32" dependencies = [ "anyhow", "bincode 1.3.3", diff --git a/lib/bindings/python/Cargo.lock b/lib/bindings/python/Cargo.lock index 7961cd1d52a4..83cf8e620ab1 100644 --- a/lib/bindings/python/Cargo.lock +++ b/lib/bindings/python/Cargo.lock @@ -34,9 +34,9 @@ dependencies = [ [[package]] name = "aisimulate-core" -version = "0.13.0-dev.202610050000000066" +version = "0.13.0-dev.202610060000000067" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c9a5aa5cb79fb04ad0c8361efcc707afeef4a08680ef3d457b5a5d63eaf00ac8" +checksum = "baa4fe3c1314053a9e439cd7a7595aec5fb2b22af0ae43d48687179846e7ed32" dependencies = [ "anyhow", "bincode 1.3.3", diff --git a/lib/bindings/python/Cargo.toml b/lib/bindings/python/Cargo.toml index d18a46664312..113ffb1c4e02 100644 --- a/lib/bindings/python/Cargo.toml +++ b/lib/bindings/python/Cargo.toml @@ -59,7 +59,7 @@ request-trace-s3 = ["dynamo-llm/request-trace-s3"] # former AIC crate surface. The `python` feature is required for # AicEngineBuilder; this extension supplies the interpreter, so the dependency # must not enable `extension-module`. -aisimulate-core = { version = "=0.13.0-dev.202610050000000066", optional = true, features = ["python"] } +aisimulate-core = { version = "=0.13.0-dev.202610060000000067", optional = true, features = ["python"] } dynamo-runtime = { path = "../../runtime" } dynamo-mocker = { path = "../../mocker" } dynamo-parsers = { version = "9.0.0" } diff --git a/pyproject.toml b/pyproject.toml index 0df99b0f3878..2bf0baec3460 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,7 @@ license-files = ["LICENSE"] requires-python = ">=3.10" dependencies = [ "ai-dynamo-runtime==1.6.0", - "aisimulate==0.13.0.dev202610050000000066 ; python_version >= '3.11' and python_version < '3.14'", + "aisimulate==0.13.0.dev202610060000000067 ; python_version >= '3.11' and python_version < '3.14'", "aiohttp>=3.14.3,<4.0", "transformers>=4.56.0", "kubernetes>=32.0.1,<33.0.0", diff --git a/tests/dependencies/test_aisimulate_namespaces.py b/tests/dependencies/test_aisimulate_namespaces.py index cba617b3a0c7..7611f2feb875 100644 --- a/tests/dependencies/test_aisimulate_namespaces.py +++ b/tests/dependencies/test_aisimulate_namespaces.py @@ -139,7 +139,7 @@ def test_no_manifest_installs_retired_aic_distributions() -> None: assert features["ais-forward-pass"] == ["dep:aisimulate-core"] assert "aic-forward-pass" not in features assert dependencies["aisimulate-core"] == { - "version": "=0.13.0-dev.202610050000000066", + "version": "=0.13.0-dev.202610060000000067", "optional": True, "features": ["python"], } From 428da55b3e954fb62025142711c72ea299acd2c3 Mon Sep 17 00:00:00 2001 From: Harry Lee Date: Tue, 6 Oct 2026 14:36:08 -0700 Subject: [PATCH 4/7] fix(replay): align AgentX capabilities and layout validation Signed-off-by: Harry Lee --- .../dynamo/mocker/tests/unit/test_config.py | 110 +++++++++++++++ components/src/dynamo/replay/config.py | 47 +++++-- components/src/dynamo/replay/simulation.py | 14 +- .../dynamo/replay/tests/test_simulation.py | 115 +++++++++++++++ .../dynosim-replay-cli-reference.mdx | 12 +- lib/bindings/python/rust/llm/replay.rs | 4 +- .../replay/test_replay_agentic_timing.py | 132 ++++++++++++++++++ lib/mocker/src/replay/agentic.rs | 13 +- lib/mocker/src/replay/offline/entrypoints.rs | 35 +++-- .../test_aisimulate_namespaces.py | 41 ++++++ 10 files changed, 490 insertions(+), 33 deletions(-) create mode 100644 lib/bindings/python/tests/replay/test_replay_agentic_timing.py diff --git a/components/src/dynamo/mocker/tests/unit/test_config.py b/components/src/dynamo/mocker/tests/unit/test_config.py index 2c2bcdedba89..3624c2e30e90 100644 --- a/components/src/dynamo/mocker/tests/unit/test_config.py +++ b/components/src/dynamo/mocker/tests/unit/test_config.py @@ -813,3 +813,113 @@ def test_ais_sdk_accepts_only_canonical_identity(): MockEngineArgs(aic_backend="vllm") with pytest.raises(TypeError): MockEngineArgs(ais_backend="vllm") + + +@pytest.mark.parametrize("backend", ["vllm", "sglang", "trtllm"]) +@pytest.mark.parametrize("configured_block_size", [None, 128]) +def test_shared_layout_uses_resolved_engine_geometry(backend, configured_block_size): + replay = _load_replay_main() + raw = { + "engine_type": backend, + "block_size": configured_block_size, + "kv_cache_bytes_per_token": 1024, + "kv_bytes_per_token": 4096, + "native_host_offload": {"scope": "cluster_shared", "num_host_blocks": 8}, + } + + lowered = replay.lower_upstream_engine_args(raw) + args = replay.load_engine_args(lowered) + layout = json.loads(lowered["native_host_offload"]["kv_layout_id"]) + + assert args.block_size == layout["block_size"] == lowered["block_size"] + native_geometry = {"engine_type": backend} + if configured_block_size is not None: + native_geometry["block_size"] = configured_block_size + assert ( + args.block_size + == MockEngineArgs.from_json(json.dumps(native_geometry)).block_size + ) + assert args.kv_cache_bytes_per_token == layout["bytes_per_token"] == 1024 + assert args.kv_bytes_per_token == 4096 + + +def test_shared_layout_respects_sglang_page_size(): + replay = _load_replay_main() + lowered = replay.lower_upstream_engine_args( + { + "engine_type": "sglang", + "sglang": {"page_size": 16}, + "kv_cache_bytes_per_token": 1024, + "native_host_offload": {"scope": "cluster_shared", "num_host_blocks": 8}, + } + ) + assert replay.load_engine_args(lowered).block_size == 16 + assert ( + json.loads(lowered["native_host_offload"]["kv_layout_id"])["block_size"] == 16 + ) + + +@pytest.mark.parametrize("explicit_null", [False, True]) +def test_shared_layout_materializes_documented_host_geometry_fallback(explicit_null): + replay = _load_replay_main() + raw = { + "kv_bytes_per_token": 2048, + "native_host_offload": {"scope": "cluster_shared", "num_host_blocks": 8}, + } + if explicit_null: + raw["kv_cache_bytes_per_token"] = None + lowered = replay.lower_upstream_engine_args(raw) + assert lowered["kv_cache_bytes_per_token"] == 2048 + assert replay.load_engine_args(lowered).kv_cache_bytes_per_token == 2048 + assert ( + json.loads(lowered["native_host_offload"]["kv_layout_id"])["bytes_per_token"] + == 2048 + ) + + +@pytest.mark.parametrize( + "geometry", + [ + {}, + {"kv_cache_bytes_per_token": None}, + {"kv_cache_bytes_per_token": 0}, + {"kv_cache_bytes_per_token": -1}, + {"kv_cache_bytes_per_token": True}, + ], +) +def test_shared_layout_rejects_unresolved_or_invalid_cache_geometry(geometry): + replay = _load_replay_main() + with pytest.raises( + ValueError, match="requires resolved positive kv_cache_bytes_per_token" + ): + replay.lower_upstream_engine_args( + { + **geometry, + "native_host_offload": { + "scope": "cluster_shared", + "num_host_blocks": 8, + }, + } + ) + + +def test_nested_replay_geometry_keeps_handoff_separate_from_cache(): + replay = _load_replay_main() + lowered = replay.lower_upstream_engine_args( + { + "rank": { + "backend": "vllm", + "kv_transfer_bytes_per_token": 4096, + "kv_cache_bytes_per_token": 1024, + "native_host_offload": { + "scope": "cluster_shared", + "num_host_blocks": 8, + "kv_layout_id": "fixture-quantized-kv-layout", + }, + } + } + ) + args = replay.load_engine_args(lowered) + assert "kv_transfer_bytes_per_token" not in lowered + assert lowered["kv_bytes_per_token"] == args.kv_bytes_per_token == 4096 + assert lowered["kv_cache_bytes_per_token"] == args.kv_cache_bytes_per_token == 1024 diff --git a/components/src/dynamo/replay/config.py b/components/src/dynamo/replay/config.py index d03ae50a4844..a8847438ebe1 100644 --- a/components/src/dynamo/replay/config.py +++ b/components/src/dynamo/replay/config.py @@ -12,6 +12,8 @@ from typing import Any, Protocol from aisimulate.capacity import materialize_aic_num_gpu_blocks +from aisimulate.config.engine import resolve_block_size +from aisimulate.runner import _kv_layout_id from dynamo.mocker import MockEngineArgs from dynamo.mocker.args import ( @@ -126,6 +128,8 @@ def lower_upstream_engine_args(payload: Mapping[str, Any]) -> dict[str, Any]: raw.update(nested) for source, target in ( ("backend", "engine_type"), + # MockEngineArgs uses this name for P/D transfer bytes; retain the + # independent physical kv_cache_bytes_per_token field unchanged. ("kv_transfer_bytes_per_token", "kv_bytes_per_token"), ): if source in raw: @@ -171,21 +175,44 @@ def lower_upstream_engine_args(payload: Mapping[str, Any]) -> dict[str, Any]: and host.get("scope") == "cluster_shared" and host.get("kv_layout_id") is None ): - # Reuse AISimulate's resolved KV identity rather than inventing a - # connector-specific layout ID. Native nested ranks require an explicit ID. - from aisimulate.runner import _kv_layout_id - + # AISimulate's pinned release owns the layout identity; its private + # helper is covered by the installed-package compatibility test. Native + # nested ranks require an explicit ID. backend = raw.get("engine_type", identity.get("backend", "vllm")) if not isinstance(backend, str) or backend not in {"vllm", "sglang", "trtllm"}: raise ValueError(f"unsupported engine_type for host offload: {backend!r}") + block_size = raw.get("block_size") + if block_size is None and backend == "sglang": + block_size = (raw.get("sglang") or {}).get("page_size") + if block_size is not None and ( + not isinstance(block_size, int) + or isinstance(block_size, bool) + or block_size <= 0 + ): + raise ValueError( + "automatic cluster_shared layout requires positive block_size" + ) + cache_bytes = raw.get("kv_cache_bytes_per_token") + if cache_bytes is None: + # Match MockEngineArgs' documented host-offload fallback. An + # explicit cache size always wins over the P/D transfer geometry. + cache_bytes = raw.get("kv_bytes_per_token") + if ( + not isinstance(cache_bytes, int) + or isinstance(cache_bytes, bool) + or cache_bytes <= 0 + ): + raise ValueError( + "automatic cluster_shared layout requires resolved positive " + "kv_cache_bytes_per_token (or kv_bytes_per_token host fallback)" + ) + # Use identical resolved geometry in both the hash and actual engine. + raw["block_size"] = resolve_block_size(backend, block_size) + raw["kv_cache_bytes_per_token"] = cache_bytes rank = { "backend": backend, - "block_size": raw.get( - "block_size", {"vllm": 64, "sglang": 1, "trtllm": 32}[backend] - ), - "kv_cache_bytes_per_token": raw.get( - "kv_cache_bytes_per_token", raw.get("kv_bytes_per_token") - ), + "block_size": raw["block_size"], + "kv_cache_bytes_per_token": cache_bytes, "timing_model": {"config": raw.get("ais_perf_config", {})}, } raw["native_host_offload"] = { diff --git a/components/src/dynamo/replay/simulation.py b/components/src/dynamo/replay/simulation.py index 024d97d1c127..20a53edc378d 100644 --- a/components/src/dynamo/replay/simulation.py +++ b/components/src/dynamo/replay/simulation.py @@ -85,11 +85,15 @@ def capabilities(self) -> RunnerCapabilities: "weka", ), supports_agentic_lanes=engine_capabilities.supports_agentic_lanes, - supports_agentic_host_offload=True, - supports_agentic_speculative_decoding=False, - supports_agentic_snapshots=True, - supports_agentic_warmup=True, - supports_agentic_profile=True, + supports_agentic_host_offload=( + engine_capabilities.supports_agentic_host_offload + ), + # The composed mocker path preserves HBM-only MTP and decode-speedup + # support. The shared G2 gate and native guard reject that combination. + supports_agentic_speculative_decoding=True, + supports_agentic_snapshots=engine_capabilities.supports_agentic_snapshots, + supports_agentic_warmup=engine_capabilities.supports_agentic_warmup, + supports_agentic_profile=engine_capabilities.supports_agentic_profile, # Full AgentX runtime conformance remains a separate checkpoint. # Keep the public runner honest about the narrower integration here. supported_agentic_topologies=tuple( diff --git a/components/src/dynamo/replay/tests/test_simulation.py b/components/src/dynamo/replay/tests/test_simulation.py index c134e2434ad6..466a9d4bb38f 100644 --- a/components/src/dynamo/replay/tests/test_simulation.py +++ b/components/src/dynamo/replay/tests/test_simulation.py @@ -698,6 +698,121 @@ def __init__( assert seen["supports_disaggregated_attention_dp"] is True +@pytest.mark.parametrize("supported", [False, True]) +def test_agentic_capabilities_follow_engine_support_in_offline_mode( + monkeypatch, supported +): + native = replace( + simulation.EngineReplayRunnerFactory().capabilities(), + supports_agentic_host_offload=supported, + supports_agentic_snapshots=supported, + supports_agentic_warmup=supported, + supports_agentic_profile=supported, + ) + monkeypatch.setattr( + simulation.EngineReplayRunnerFactory, "capabilities", lambda self: native + ) + capabilities = simulation.DynamoReplayRunnerFactory().capabilities() + assert capabilities.supported_execution_modes == ("offline",) + assert capabilities.supports_agentic_host_offload is supported + assert capabilities.supports_agentic_snapshots is supported + assert capabilities.supports_agentic_warmup is supported + assert capabilities.supports_agentic_profile is supported + # Dynamo's composed runtime retains its pre-existing HBM speculation support. + assert capabilities.supports_agentic_speculative_decoding + + +@pytest.mark.parametrize("nested", [False, True]) +@pytest.mark.parametrize("alias", ["aic_nextn", "nextn", "speculation"]) +def test_agentic_speculation_capability_does_not_bypass_g2_scope(nested, alias): + rank = { + "engine_type": "vllm", + "native_host_offload": {"num_host_blocks": 8}, + alias: {"nextn": 2} if alias == "speculation" else 2, + } + deployment = replace( + _agg_deployment(), + num_workers=1, + agg_engine_args={"rank": rank} if nested else rank, + ) + spec = ReplaySpec( + backend_deployment=deployment, + workload={"trace_format": "weka", "agentic_lanes": 1}, + goal={"target": "throughput"}, + ) + capabilities = simulation.DynamoReplayRunnerFactory().capabilities() + with pytest.raises(ValueError, match="requires speculative decoding disabled"): + capabilities.require_compatible(spec) + rank.pop("native_host_offload") + capabilities.require_compatible(spec) + + +@pytest.mark.timeout(30) +@pytest.mark.parametrize( + "speculation", [{"aic_nextn": 2}, {"decode_speedup_ratio": 2.0}] +) +def test_hbm_agentic_speculation_runs_through_native_replay(tmp_path, speculation): + trace = tmp_path / "hbm-agentic.jsonl" + trace.write_text( + json.dumps( + { + "id": "hbm-play", + "models": ["target-model"], + "block_size": 4, + "hash_id_scope": "local", + "requests": [ + { + "t": 0, + "type": "s", + "model": "target-model", + "in": 4, + "out": 7, + "hash_ids": [1], + "api_time": 0.1, + } + ], + } + ) + + "\n" + ) + base_args = { + "engine_type": "vllm", + "block_size": 4, + "num_gpu_blocks": 8, + "max_num_batched_tokens": 16, + "max_num_seqs": 1, + "timing_model": {"type": "fixed", "prefill_ms": 2, "decode_ms": 8}, + } + spec = ReplaySpec( + backend_deployment=replace( + _agg_deployment(), num_workers=1, agg_engine_args=base_args + ), + workload={ + "source_type": "trace", + "load_type": "trace_timestamps", + "trace_format": "weka", + "trace_path": str(trace), + "trace_block_size": 4, + "agentic_lanes": 1, + }, + goal={"target": "throughput"}, + ) + runner = simulation.DynamoReplayRunnerFactory().create(0) + baseline = runner.run(spec) + accelerated = runner.run( + replace( + spec, + backend_deployment=replace( + spec.backend_deployment, agg_engine_args={**base_args, **speculation} + ), + ) + ) + assert baseline.metrics["completed_requests"] == 1 + assert accelerated.metrics["completed_requests"] == 1 + assert accelerated.metrics["total_output_tokens"] == 7 + assert accelerated.metrics["duration_ms"] < baseline.metrics["duration_ms"] + + def test_goodput_goal_fails_closed_when_replay_omits_metric(monkeypatch) -> None: monkeypatch.setattr(simulation, "MockEngineArgs", _FakeEngineArgs) monkeypatch.setattr( diff --git a/docs/fern/pages/reference/components/dynosim-replay-cli-reference.mdx b/docs/fern/pages/reference/components/dynosim-replay-cli-reference.mdx index 2cee6e9d64fa..96479996753b 100644 --- a/docs/fern/pages/reference/components/dynosim-replay-cli-reference.mdx +++ b/docs/fern/pages/reference/components/dynosim-replay-cli-reference.mdx @@ -252,7 +252,8 @@ Mooncake-delta traces. No connector configuration is required. For AgentX graphs, G2 requires vLLM, one aggregated worker or one prefill plus one decode worker, attention DP=1 on every role, static worker pools, and no -speculative decoding. Tensor parallelism may exceed one. These AgentX limits +speculative decoding (`ais_nextn` and `ais_perf_config.speculation` unset, +`decode_speedup_ratio=1.0`). Tensor parallelism may exceed one. These AgentX limits also apply when a Dynamo trace contains agent context but `agentic_lanes` is omitted. Ordinary traces retain their existing topology support; a standard Dynamo trace with two workers and attention DP=2 is covered by the G2 tests. @@ -266,10 +267,17 @@ must supply a compatible `kv_layout_id` for a shared pool. Host-tier events remain distinct from GPU residency, while the KV router retains its configured host-tier weighting. P-to-D handoff uses GPU-resident hashes. +`kv_bytes_per_token` describes P-to-D transfer size; +`kv_cache_bytes_per_token` describes the physical cache footprint used by G2. +Set both when the geometries differ. With host offload enabled, an omitted cache +size falls back to the transfer size. Automatic shared-layout identity requires +a resolved, positive cache size. + The Python `run_trace_replay` API also accepts `agentic_snapshot` (a mapping containing `seed`), `agentic_warmup` (default `False`), and `agentic_profile` (AISimulate's duration, response-grace, and cancellation-drain options). -Snapshots require positive `agentic_lanes`; warmup and profiling require a +These phase controls require offline replay. Snapshots require positive +`agentic_lanes`; warmup and profiling require a snapshot. Profiling cannot be combined with `max_sim_time_ms`. Warmup retains G1/G2 state across the measurement barrier. Results retain `agentic_snapshots`, `agentic_phases`, and `agentic_profile` when those phases are enabled. diff --git a/lib/bindings/python/rust/llm/replay.rs b/lib/bindings/python/rust/llm/replay.rs index cf8a5a55eea8..d5c93f116227 100644 --- a/lib/bindings/python/rust/llm/replay.rs +++ b/lib/bindings/python/rust/llm/replay.rs @@ -957,7 +957,7 @@ pub fn run_mocker_trace_replay( } if replay_concurrency.is_some() { return Err(PyValueError::new_err( - "agentic_snapshot requires agentic trace input without replay_concurrency", + "agentic phases require agentic trace input without replay_concurrency", )); } if !matches!( @@ -1430,6 +1430,8 @@ fn run_loaded_dynamo_request_trace( ); } let trace = if agentic_options.snapshot.is_some() { + // Snapshot sampling uses recorded source time; normalizing + // here would change the sampled frontier and cache state. trace } else { trace.normalize_starts() diff --git a/lib/bindings/python/tests/replay/test_replay_agentic_timing.py b/lib/bindings/python/tests/replay/test_replay_agentic_timing.py new file mode 100644 index 000000000000..58c0fbcf8b2a --- /dev/null +++ b/lib/bindings/python/tests/replay/test_replay_agentic_timing.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json + +import pytest + +from dynamo.replay import run_trace_replay + +from .replay_utils import _report_summary, _vllm_args + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.parallel, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.core, + pytest.mark.timeout(120), +] + + +def _write_timing_trace(path, *, start_ms, root_gap_ms, dependency_delay_ms): + """Original fixture with roots and two turns sharing a completion trigger.""" + records = [ + { + "schema": "dynamo.agentic_mooncake", + "version": 2, + "block_size": 64, + "hash_id_scope": "local", + "source": {"format": "self-authored-test", "digest": "finite-timing"}, + } + ] + for index, (request_id, ready_at_ms, dependencies) in enumerate( + [ + ("root", start_ms, []), + ("independent", start_ms + root_gap_ms, []), + ( + "anchor", + start_ms, + [ + { + "request_id": "root", + "trigger": "completion", + "delay_ms": 0.0, + "relation": "sequence", + } + ], + ), + ( + "dependent", + start_ms, + [ + { + "request_id": "root", + "trigger": "completion", + "delay_ms": dependency_delay_ms, + "relation": "sequence", + } + ], + ), + ], + start=1, + ): + records.append( + { + "request_id": request_id, + "play_id": "play", + "session_id": request_id, + "model": "fixture-model", + "input_length": 64, + "output_length": 2, + "hash_ids": [index], + "not_before_ms": ready_at_ms, + "dependencies": dependencies, + } + ) + path.write_text( + "\n".join(json.dumps(record) for record in records) + "\n", encoding="utf-8" + ) + return path + + +@pytest.mark.parametrize("replay_mode", ["offline", "online"]) +@pytest.mark.parametrize("agentic_lanes", [None, 1]) +def test_agentic_mooncake_finite_binding_preserves_normalized_scaled_timing( + tmp_path, replay_mode, agentic_lanes +): + source = _write_timing_trace( + tmp_path / "source.jsonl", + start_ms=100.0, + root_gap_ms=48.0, + dependency_delay_ms=32.0, + ) + # The former file-loading path normalized starts before scaling timers. + # Author that result independently as a control for the binding dispatch. + prepared = _write_timing_trace( + tmp_path / "prepared.jsonl", + start_ms=0.0, + root_gap_ms=12.0, + dependency_delay_ms=8.0, + ) + timing = [] + for trace, speedup in [(source, 4.0), (prepared, 1.0)]: + report_path = tmp_path / f"{trace.stem}-report.jsonl" + report = run_trace_replay( + trace, + extra_engine_args=_vllm_args(), + replay_mode=replay_mode, + trace_format="agentic_mooncake", + execution_model="fixture-model", + agentic_lanes=agentic_lanes, + arrival_speedup_ratio=speedup, + report_jsonl_path=report_path, + ) + assert _report_summary(report)["completed_requests"] == 4 + records = [json.loads(line) for line in report_path.read_text().splitlines()] + by_id = {record["request_id"]: record for record in records} + root, independent, anchor, dependent = ( + by_id[name] for name in ("root", "independent", "anchor", "dependent") + ) + # Online completion notification can lag terminal delivery. Both turns + # share that notification, so their scheduled arrival gap isolates the + # authored delay without any wall-clock tolerance. + timing.append( + ( + root["arrival_time_ms"], + independent["arrival_time_ms"], + dependent["arrival_time_ms"] - anchor["arrival_time_ms"], + ) + ) + assert timing[-1] == pytest.approx((0.0, 12.0, 8.0)) + assert timing[0] == pytest.approx(timing[1]) diff --git a/lib/mocker/src/replay/agentic.rs b/lib/mocker/src/replay/agentic.rs index df0f208415b6..2ba65e694a0d 100644 --- a/lib/mocker/src/replay/agentic.rs +++ b/lib/mocker/src/replay/agentic.rs @@ -65,6 +65,13 @@ impl AgenticReplayOptions { lanes: Option, max_sim_time_ms: Option, ) -> Result { + // Snapshot lanes independently sample the corpus and may outnumber + // source plays; only finite replay caps lanes to the play count. + let lanes = if self.snapshot.is_some() { + lanes + } else { + super::effective_agentic_lanes(lanes, trace.play_count()) + }; self.validate(lanes, max_sim_time_ms)?; let Some(snapshot) = self.snapshot else { return trace @@ -168,7 +175,11 @@ mod tests { #[test] fn finite_workload_preserves_timestamps_and_scales_timers_once() { - for (lanes, first_ms, second_ms) in [(None, 100.0, 220.0), (Some(1), 0.0, 120.0)] { + for (lanes, first_ms, second_ms) in [ + (None, 100.0, 220.0), + (Some(1), 0.0, 120.0), + (Some(3), 0.0, 120.0), + ] { let mut regular = AgenticReplayOptions::default() .into_driver(recorded_graph(), 4, true, lanes, None) .unwrap(); diff --git a/lib/mocker/src/replay/offline/entrypoints.rs b/lib/mocker/src/replay/offline/entrypoints.rs index 3db0deecc90c..17b709d0f2d6 100644 --- a/lib/mocker/src/replay/offline/entrypoints.rs +++ b/lib/mocker/src/replay/offline/entrypoints.rs @@ -27,7 +27,6 @@ use crate::loadgen::{AgenticTrace, Trace, WorkloadDriver}; use crate::replay::{ OfflineDisaggReplayConfig, ReplayPrefillLoadEstimator, ReplayRouterMode, ReplayTelemetryOptions, ReplayWorkerArtifacts, SlaThresholds, TraceSimulationReport, - effective_agentic_lanes, }; use crate::scheduler::RouterEventVisibility; @@ -84,13 +83,14 @@ fn validate_agentic_host_offload( "agentic host offload requires attention DP=1 on every role" ); anyhow::ensure!( - args.ais_nextn.is_none() + args.decode_speedup_ratio == 1.0 + && args.ais_nextn.is_none() && args.ais_perf_config.as_ref().is_none_or(|config| { config .get("speculation") .is_none_or(serde_json::Value::is_null) }), - "agentic replay requires speculative decoding disabled" + "agentic host offload requires speculative decoding disabled" ); } Ok(()) @@ -648,11 +648,6 @@ pub(crate) fn simulate_agentic_trace_workload( telemetry: Option, ) -> Result { let args = args.normalized()?; - let agentic_lanes = if agentic_options.snapshot.is_some() { - agentic_lanes - } else { - effective_agentic_lanes(agentic_lanes, trace.play_count()) - }; let driver = agentic_options.into_driver( trace, args.block_size, @@ -692,11 +687,6 @@ pub(crate) fn simulate_agentic_trace_workload_disagg( telemetry: Option, ) -> Result { let config = config.normalized()?; - let agentic_lanes = if agentic_options.snapshot.is_some() { - agentic_lanes - } else { - effective_agentic_lanes(agentic_lanes, trace.play_count()) - }; let driver = agentic_options.into_driver( trace, config.prefill_args.block_size, @@ -983,7 +973,19 @@ mod agentic_host_offload_tests { invalid.ais_nextn = Some(1); invalid_roles.push(( invalid, - "agentic replay requires speculative decoding disabled", + "agentic host offload requires speculative decoding disabled", + )); + let mut invalid = args(None); + invalid.decode_speedup_ratio = 2.0; + invalid_roles.push(( + invalid, + "agentic host offload requires speculative decoding disabled", + )); + let mut invalid = args(None); + invalid.ais_perf_config = Some(serde_json::json!({"speculation": {"nextn": 2}})); + invalid_roles.push(( + invalid, + "agentic host offload requires speculative decoding disabled", )); for (invalid, expected) in invalid_roles { for roles in [ @@ -1066,10 +1068,15 @@ mod agentic_host_offload_tests { let mut hbm = args(None); hbm.engine_type = EngineType::Sglang; hbm.dp_size = 2; + hbm.decode_speedup_ratio = 2.0; + validate_agentic_host_offload(&agentic_input(), &[(2, &hbm)], true).unwrap(); + hbm.decode_speedup_ratio = 1.0; + hbm.ais_nextn = Some(2); validate_agentic_host_offload(&agentic_input(), &[(2, &hbm)], true).unwrap(); let mut offloaded = args(Some("cluster_shared")); offloaded.dp_size = 2; + offloaded.decode_speedup_ratio = 2.0; validate_agentic_host_offload( &ReplayRuntimeInput::Requests(VecDeque::new()), &[(2, &offloaded)], diff --git a/tests/dependencies/test_aisimulate_namespaces.py b/tests/dependencies/test_aisimulate_namespaces.py index 7611f2feb875..d973a4b0e05d 100644 --- a/tests/dependencies/test_aisimulate_namespaces.py +++ b/tests/dependencies/test_aisimulate_namespaces.py @@ -5,6 +5,7 @@ from __future__ import annotations +import json import subprocess import sys import textwrap @@ -173,3 +174,43 @@ def test_aisimulate_wheel_uses_canonical_import_namespaces() -> None: entry for entry in release.entry_points if entry.name == "aiconfigurator" ) assert legacy_cli.value == "aisimulate.legacy_cli.entrypoint:main" + + +def test_pinned_aisimulate_exports_replay_layout_helpers() -> None: + if sys.version_info < (3, 11) or sys.version_info >= (3, 14): + pytest.skip("AISimulate supports Python 3.11 through 3.13") + + # The exact dependency pin owns this private layout protocol. Fail this + # installed-package check when an upgrade removes or changes its contract. + from aisimulate.config.engine import resolve_block_size + from aisimulate.runner import _kv_layout_id + + block_size = resolve_block_size("vllm", None) + rank = { + "backend": "vllm", + "block_size": block_size, + "kv_cache_bytes_per_token": 1024, + "timing_model": { + "config": { + "model": "resolved-model", + "tp": 2, + "backend_version": "fixture-version", + "pp": 1, + "dcp": 2, + "kvcache_quant_mode": "fp8", + "attention_backend": "fixture-attention", + } + }, + } + assert json.loads(_kv_layout_id(rank, "fallback-model", 1)) == { + "model": "resolved-model", + "backend": "vllm", + "tp": 2, + "block_size": block_size, + "bytes_per_token": 1024, + "backend_version": "fixture-version", + "pp": 1, + "dcp": 2, + "kvcache_quant_mode": "fp8", + "attention_backend": "fixture-attention", + } From 3f5f5460ad3ea18ff1826d9feb8796ba90e0bc78 Mon Sep 17 00:00:00 2001 From: Harry Lee Date: Wed, 7 Oct 2026 03:43:40 -0700 Subject: [PATCH 5/7] fix(ci): isolate generic image builds from shared workers Signed-off-by: Harry Lee --- .github/actions/builder-refresher/action.yml | 14 ++++++++++---- .github/workflows/pr.yaml | 16 ++++++++++++++++ .github/workflows/shared-build-image.yml | 3 +++ 3 files changed, 29 insertions(+), 4 deletions(-) diff --git a/.github/actions/builder-refresher/action.yml b/.github/actions/builder-refresher/action.yml index 3400ee7c51a7..5086940db949 100644 --- a/.github/actions/builder-refresher/action.yml +++ b/.github/actions/builder-refresher/action.yml @@ -10,9 +10,8 @@ description: 'Ensure a BuildKit builder is healthy; re-initialize it if the conn # still reachable. If it succeeds the action exits immediately — no extra # work is done. # 2. If the inspect/bootstrap fails the builder is considered stale and -# init-dynamo-builder is called with the same flavor/arch/cuda_version that -# were used when the builder was first created, effectively re-routing -# workers and re-registering the builder from scratch. +# init-dynamo-builder is called with the same routing and isolation options +# that were used when the builder was first created. inputs: builder_name: @@ -29,6 +28,10 @@ inputs: description: 'CUDA version used on re-init (e.g. 13.0, 13.1). Leave empty for general flavor.' required: false default: '' + fresh_builder: + description: 'Preserve an isolated Kubernetes builder when re-initializing' + required: false + default: 'false' # Kubernetes fallback passthrough inputs (forwarded to init-dynamo-builder) ephemeral_storage: @@ -76,7 +79,9 @@ runs: shell: bash run: | echo "::warning::Builder '${{ inputs.builder_name }}' is unhealthy. Removing and re-initializing..." - docker buildx rm "${{ inputs.builder_name }}" || true + # Other image jobs may still use the same run-scoped Kubernetes pods. + # Only the workflow cleanup job should delete those shared daemons. + docker buildx rm --keep-daemon "${{ inputs.builder_name }}" || true - name: Re-initialize builder (if unhealthy) if: steps.check-health.outcome == 'failure' @@ -86,6 +91,7 @@ runs: flavor: ${{ inputs.flavor }} arch: ${{ inputs.arch }} cuda_version: ${{ inputs.cuda_version }} + fresh_builder: ${{ inputs.fresh_builder }} ephemeral_storage: ${{ inputs.ephemeral_storage }} namespace: ${{ inputs.namespace }} replicas: ${{ inputs.replicas }} diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 376b6690d76a..d6b3369879ef 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -14,6 +14,8 @@ concurrency: cancel-in-progress: true env: + # Generic image builds share one run-scoped K8s builder pair instead of + # competing with other workflows on the remote general worker pool. BUILDER_NAME: b-${{ github.run_id }}-${{ github.run_attempt }} jobs: @@ -594,6 +596,7 @@ jobs: cuda_version: '[""]' platform: 'linux/amd64' # @whoisj to enable ARM64 support in a future patch. builder_name: ${{ needs.changed-files.outputs.builder_name }} + fresh_builder: true build_timeout_minutes: 60 secrets: inherit @@ -699,6 +702,7 @@ jobs: cuda_version: '[""]' platform: 'linux/amd64,linux/arm64' builder_name: ${{ needs.changed-files.outputs.builder_name }} + fresh_builder: true inline_compliance: true artifact_retention_days: '7' build_timeout_minutes: 60 @@ -771,6 +775,7 @@ jobs: cuda_version: '[""]' platform: 'linux/amd64,linux/arm64' builder_name: ${{ needs.changed-files.outputs.builder_name }} + fresh_builder: true inline_compliance: true artifact_retention_days: '7' # Include cold multi-arch compilation and the later test-image/compliance stages. @@ -1057,6 +1062,7 @@ jobs: uses: $/.github/workflows/dynamo-pipeline.yml with: builder_name: ${{ needs.changed-files.outputs.builder_name }} + fresh_builder: true build_timeout_minutes: 60 # `defaulted` is applied by the repository-root conftest.py to any test that # declared no lifecycle or hardware marker. It is routed here, not to the @@ -1957,6 +1963,8 @@ jobs: name: Clean K8s builder if exists runs-on: prod-default-small-v2 if: always() + # Wait for every builder consumer, including jobs that can fall back to + # the same K8s pair. Test-only jobs use published images, not the builder. needs: - changed-files - operator @@ -1964,12 +1972,20 @@ jobs: - sidecar-build - planner-test - planner-copy-to-acr + - frontend-build + - triton-build + - vllm-dev-build + - vllm-local-dev-build - vllm-copy-to-acr - vllm-multi-gpu-test - vllm-efa-build + - sglang-dev-build + - sglang-local-dev-build - sglang-copy-to-acr - sglang-multi-gpu-test - sglang-efa-build + - trtllm-dev-build + - trtllm-local-dev-build - trtllm-copy-to-acr - trtllm-multi-gpu-test - trtllm-efa-build diff --git a/.github/workflows/shared-build-image.yml b/.github/workflows/shared-build-image.yml index 24265a2031d2..f86014b94bb7 100644 --- a/.github/workflows/shared-build-image.yml +++ b/.github/workflows/shared-build-image.yml @@ -439,6 +439,7 @@ jobs: builder_name: ${{ inputs.builder_name }} flavor: ${{ steps.calculate-builder-flavor.outputs.builder_flavor }} arch: ${{ inputs.platform }} + fresh_builder: ${{ inputs.fresh_builder }} namespace: ${{ vars.BUILDKIT_NAMESPACE }} - name: Generate Dockerfile id: gen-dockerfile @@ -570,6 +571,7 @@ jobs: flavor: ${{ steps.calculate-builder-flavor.outputs.builder_flavor }} arch: ${{ inputs.platform }} cuda_version: ${{ matrix.cuda_version }} + fresh_builder: ${{ inputs.fresh_builder }} namespace: ${{ vars.BUILDKIT_NAMESPACE }} - name: Compliance extract if: ${{ (inputs.compliance_only || inputs.inline_compliance) && inputs.target != 'dev' && inputs.target != 'local-dev' }} @@ -614,6 +616,7 @@ jobs: flavor: ${{ steps.calculate-builder-flavor.outputs.builder_flavor }} arch: ${{ inputs.platform }} cuda_version: ${{ matrix.cuda_version }} + fresh_builder: ${{ inputs.fresh_builder }} namespace: ${{ vars.BUILDKIT_NAMESPACE }} - name: Build and Push Test Image if: ${{ !inputs.compliance_only && inputs.target != 'dev' && inputs.target != 'local-dev' && inputs.target != 'frontend' }} From f1e26141ed1161ff397b33c54cc2a7f07b09777b Mon Sep 17 00:00:00 2001 From: Harry Lee Date: Wed, 7 Oct 2026 10:56:49 -0700 Subject: [PATCH 6/7] fix(ci): pass builder names through environment variables Signed-off-by: Harry Lee --- .github/actions/bootstrap-buildkit/action.yml | 18 ++++++++++++------ .github/actions/builder-refresher/action.yml | 12 ++++++++---- 2 files changed, 20 insertions(+), 10 deletions(-) diff --git a/.github/actions/bootstrap-buildkit/action.yml b/.github/actions/bootstrap-buildkit/action.yml index 486cb2ea81b2..744bb0882d5e 100644 --- a/.github/actions/bootstrap-buildkit/action.yml +++ b/.github/actions/bootstrap-buildkit/action.yml @@ -69,26 +69,30 @@ runs: - name: Define remote buildkit builders if: inputs.buildkit_worker_addresses != '' shell: bash + env: + BUILDER_NAME: ${{ inputs.builder_name }} run: | ADDRS="${{ inputs.buildkit_worker_addresses }}" IFS=',' read -ra ADDR_LIST <<< "$ADDRS" FIRST=true for addr in "${ADDR_LIST[@]}"; do if $FIRST; then - docker buildx create --use --name ${{ inputs.builder_name }} --driver remote "$addr" + docker buildx create --use --name "$BUILDER_NAME" --driver remote "$addr" FIRST=false else - docker buildx create --append --name ${{ inputs.builder_name }} --driver remote "$addr" + docker buildx create --append --name "$BUILDER_NAME" --driver remote "$addr" fi done - name: Create Kubernetes builder for both platforms if: inputs.buildkit_worker_addresses == '' shell: bash + env: + BUILDER_NAME: ${{ inputs.builder_name }} run: | # If the builder already exists, these commands will just set the right configurations to use it - echo "🔨 Creating K8s builder '${{ inputs.builder_name }}'." - docker buildx create --use --name ${{ inputs.builder_name }} --driver kubernetes --platform=linux/amd64 \ + echo "🔨 Creating K8s builder." + docker buildx create --use --name "$BUILDER_NAME" --driver kubernetes --platform=linux/amd64 \ '--driver-opt=requests.ephemeral-storage=${{ inputs.ephemeral_storage }}' \ '--driver-opt=namespace=${{ inputs.namespace }}' \ '--driver-opt=loadbalance=sticky' \ @@ -99,7 +103,7 @@ runs: '--driver-opt="nodeselector=kubernetes.io/arch=amd64,role=dynamo-builder-fallback"' \ '--driver-opt="tolerations=${{ inputs.tolerations }}"' - docker buildx create --append --name ${{ inputs.builder_name }} --driver kubernetes --platform=linux/arm64 \ + docker buildx create --append --name "$BUILDER_NAME" --driver kubernetes --platform=linux/arm64 \ '--driver-opt=requests.ephemeral-storage=${{ inputs.ephemeral_storage }}' \ '--driver-opt=namespace=${{ inputs.namespace }}' \ '--driver-opt=loadbalance=sticky' \ @@ -124,10 +128,12 @@ runs: - name: Bootstrap buildkit if: inputs.skip_bootstrap != 'true' shell: bash + env: + BUILDER_NAME: ${{ inputs.builder_name }} run: | echo "Bootstrapping buildkit..." for i in 1 2 3; do - if docker buildx inspect ${{ inputs.builder_name }} --bootstrap; then + if docker buildx inspect --bootstrap -- "$BUILDER_NAME"; then echo "Bootstrap succeeded on attempt $i" break fi diff --git a/.github/actions/builder-refresher/action.yml b/.github/actions/builder-refresher/action.yml index 5086940db949..6721a963aefd 100644 --- a/.github/actions/builder-refresher/action.yml +++ b/.github/actions/builder-refresher/action.yml @@ -69,19 +69,23 @@ runs: id: check-health continue-on-error: true shell: bash + env: + BUILDER_NAME: ${{ inputs.builder_name }} run: | - echo "Checking BuildKit builder '${{ inputs.builder_name }}'..." - docker buildx inspect "${{ inputs.builder_name }}" --bootstrap + echo "Checking BuildKit builder health..." + docker buildx inspect --bootstrap -- "$BUILDER_NAME" echo "Builder is healthy." - name: Remove stale builder (if unhealthy) if: steps.check-health.outcome == 'failure' shell: bash + env: + BUILDER_NAME: ${{ inputs.builder_name }} run: | - echo "::warning::Builder '${{ inputs.builder_name }}' is unhealthy. Removing and re-initializing..." + echo "::warning::Builder is unhealthy. Removing and re-initializing..." # Other image jobs may still use the same run-scoped Kubernetes pods. # Only the workflow cleanup job should delete those shared daemons. - docker buildx rm --keep-daemon "${{ inputs.builder_name }}" || true + docker buildx rm --keep-daemon -- "$BUILDER_NAME" || true - name: Re-initialize builder (if unhealthy) if: steps.check-health.outcome == 'failure' From b016a57bcafed569c88d8ee109f19f436faa6045 Mon Sep 17 00:00:00 2001 From: Harry Lee Date: Wed, 7 Oct 2026 13:59:04 -0700 Subject: [PATCH 7/7] ci: restore remote routing for general image builds Signed-off-by: Harry Lee --- .github/workflows/pr.yaml | 6 ------ 1 file changed, 6 deletions(-) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index d6b3369879ef..2705e984f48c 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -14,8 +14,6 @@ concurrency: cancel-in-progress: true env: - # Generic image builds share one run-scoped K8s builder pair instead of - # competing with other workflows on the remote general worker pool. BUILDER_NAME: b-${{ github.run_id }}-${{ github.run_attempt }} jobs: @@ -596,7 +594,6 @@ jobs: cuda_version: '[""]' platform: 'linux/amd64' # @whoisj to enable ARM64 support in a future patch. builder_name: ${{ needs.changed-files.outputs.builder_name }} - fresh_builder: true build_timeout_minutes: 60 secrets: inherit @@ -702,7 +699,6 @@ jobs: cuda_version: '[""]' platform: 'linux/amd64,linux/arm64' builder_name: ${{ needs.changed-files.outputs.builder_name }} - fresh_builder: true inline_compliance: true artifact_retention_days: '7' build_timeout_minutes: 60 @@ -775,7 +771,6 @@ jobs: cuda_version: '[""]' platform: 'linux/amd64,linux/arm64' builder_name: ${{ needs.changed-files.outputs.builder_name }} - fresh_builder: true inline_compliance: true artifact_retention_days: '7' # Include cold multi-arch compilation and the later test-image/compliance stages. @@ -1062,7 +1057,6 @@ jobs: uses: $/.github/workflows/dynamo-pipeline.yml with: builder_name: ${{ needs.changed-files.outputs.builder_name }} - fresh_builder: true build_timeout_minutes: 60 # `defaulted` is applied by the repository-root conftest.py to any test that # declared no lifecycle or hardware marker. It is routed here, not to the