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 3400ee7c51a7..6721a963aefd 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: @@ -66,17 +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..." - docker buildx rm "${{ inputs.builder_name }}" || true + 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 -- "$BUILDER_NAME" || true - name: Re-initialize builder (if unhealthy) if: steps.check-health.outcome == 'failure' @@ -86,6 +95,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..2705e984f48c 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -1957,6 +1957,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 +1966,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' }} diff --git a/components/src/dynamo/replay/api.py b/components/src/dynamo/replay/api.py index 9efb02dc885a..e110469037fe 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/simulation.py b/components/src/dynamo/replay/simulation.py index 5e98eecd7290..2002c172a161 100644 --- a/components/src/dynamo/replay/simulation.py +++ b/components/src/dynamo/replay/simulation.py @@ -87,6 +87,15 @@ def capabilities(self) -> RunnerCapabilities: "weka", ), supports_agentic_lanes=engine_capabilities.supports_agentic_lanes, + 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( @@ -169,8 +178,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, @@ -311,9 +322,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() @@ -366,6 +384,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 @@ -382,6 +405,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, "aggregated"), num_workers=deployment.num_workers, @@ -396,6 +420,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"), decode_engine_args=self._engine_args(deployment, "decode"), @@ -516,9 +541,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..c6b84d3f33c8 --- /dev/null +++ b/components/src/dynamo/replay/tests/test_agentx_g2.py @@ -0,0 +1,457 @@ +# 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 = { + "backend": "vllm", + "worker_type": "aggregated", + "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",) + # ReplaySpec accepts rank descriptors; the runner materializes these into + # native MockerConfig launch inputs with engine nesting and top-level DP. + workers = { + role: {"rank": {**copy.deepcopy(args), "worker_type": role}, "dp_size": 1} + 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["rank"]["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["rank"].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["rank"]["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["rank"]["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["rank"]["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["rank"]["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 68f592eb00b7..4c7c7b3b7cea 100644 --- a/components/src/dynamo/replay/tests/test_simulation.py +++ b/components/src/dynamo/replay/tests/test_simulation.py @@ -17,6 +17,7 @@ reason="AI Simulate is an optional Dynamo simulation dependency", ) +from aisimulate.runner import materialize_engine_launch_config from aisimulate.sweeper.provider import AdapterReplaySpec, RuntimeHookSpec from aisimulate.sweeper.replay import ( BackendDeploymentSpec, @@ -276,6 +277,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( @@ -690,6 +715,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, @@ -852,3 +992,141 @@ def test_compiled_custom_timing_consumes_capacity_only_fields(timing): assert args["engine"]["timing_model"].get("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_launch_config_preserves_layout_and_link_controls(nested): + 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_launch_config("vllm", "", {}, raw, "aggregated") + if nested: + raw = { + "dp_size": materialized["dp_size"], + "tensor_parallel_size": materialized["tensor_parallel_size"], + "rank": materialized["engine"], + } + config = _launch_config(raw) + engine = config["engine"] + assert ( + engine["native_host_offload"] == materialized["engine"]["native_host_offload"] + ) + assert engine["native_host_offload"].items() >= host.items() + assert config["tensor_parallel_size"] == 2 + assert engine["kv_cache_bytes_per_token"] == 250_000 + assert "rank" not in config + + +@pytest.mark.parametrize( + "rank, message", + [ + ({"backend": "vllm", "engine_type": "sglang"}, "unknown field `engine_type`"), + ( + {"kv_transfer_bytes_per_token": 16, "kv_bytes_per_token": 32}, + "duplicate field `kv_transfer_bytes_per_token`", + ), + ], +) +def test_nested_engine_arguments_reject_conflicting_aliases(rank, message): + with pytest.raises(ValueError, match=message): + _launch_config( + { + "rank": { + "num_gpu_blocks": 3, + "timing_model": { + "type": "fixed", + "prefill_ms": 1.0, + "decode_ms": 1.0, + }, + **rank, + } + } + ) + + +@pytest.mark.parametrize("configured_block_size", [None, 128]) +def test_shared_g2_layout_matches_canonical_engine_geometry(configured_block_size): + payload = { + "engine_type": "vllm", + "num_gpu_blocks": 3, + "kv_cache_bytes_per_token": 1024, + "kv_bytes_per_token": 4096, + "native_host_offload": {"scope": "cluster_shared", "num_host_blocks": 8}, + "timing_model": {"type": "fixed", "prefill_ms": 1.0, "decode_ms": 1.0}, + } + if configured_block_size is not None: + payload["block_size"] = configured_block_size + engine = _launch_config(payload)["engine"] + layout = json.loads(engine["native_host_offload"]["kv_layout_id"]) + + assert ( + engine["block_size"] + == layout["block_size"] + == (64 if configured_block_size is None else configured_block_size) + ) + # Physical G2 storage and P/D transfer sizes remain independent. + assert engine["kv_cache_bytes_per_token"] == layout["bytes_per_token"] == 1024 + assert engine["kv_transfer_bytes_per_token"] == 4096 + + +@pytest.mark.parametrize( + "geometry, message", + [ + ({}, "requires kv_cache_bytes_per_token"), + ({"kv_cache_bytes_per_token": None}, "requires kv_cache_bytes_per_token"), + ({"kv_cache_bytes_per_token": 0}, "kv_cache_bytes_per_token must be positive"), + ({"kv_cache_bytes_per_token": -1}, "expected usize"), + ({"kv_cache_bytes_per_token": True}, "expected usize"), + ], +) +def test_shared_g2_launch_requires_resolved_physical_cache_geometry(geometry, message): + # The CLI resolves model geometry first. Canonical replay inputs must supply + # physical G2 bytes explicitly, independently of optional handoff bytes. + with pytest.raises(ValueError, match=message): + _launch_config( + { + "num_gpu_blocks": 3, + "kv_bytes_per_token": 4096, + "native_host_offload": { + "scope": "cluster_shared", + "num_host_blocks": 8, + }, + "timing_model": {"type": "fixed", "prefill_ms": 1.0, "decode_ms": 1.0}, + **geometry, + } + ) + + +@pytest.mark.parametrize("backend", ["sglang", "trtllm"]) +def test_shared_g2_launch_rejects_unsupported_backend(backend): + with pytest.raises( + ValueError, match="native_host_offload is supported only for backend=vllm" + ): + _launch_config( + { + "engine_type": backend, + "num_gpu_blocks": 3, + "kv_cache_bytes_per_token": 1024, + "native_host_offload": { + "scope": "cluster_shared", + "num_host_blocks": 8, + }, + "timing_model": {"type": "fixed", "prefill_ms": 1.0, "decode_ms": 1.0}, + } + ) diff --git a/container/deps/requirements.aisimulate.txt b/container/deps/requirements.aisimulate.txt index d09517bf37f3..eb2e5bd5593c 100644 --- a/container/deps/requirements.aisimulate.txt +++ b/container/deps/requirements.aisimulate.txt @@ -1,5 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# Published AISimulate wheel staged into planner and framework runtime wheelhouses. # Keep in sync with the aisimulate-core pins in Cargo.toml and lib/bindings/python/Cargo.toml. aisimulate==0.13.0.dev202610060000000067 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 8c892b3aaf2e..83d598957fab 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. ### Canonical Python Engine Configuration @@ -282,6 +282,54 @@ 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 `engine.native_host_offload` in the canonical engine configuration passed to +the Python replay API. 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 (`engine.aic_nextn` and +`engine.timing_model.config.speculation` unset, +`engine.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. +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. AISimulate's public engine materializer derives the layout identity for +CLI and ReplaySpec inputs; direct canonical configuration callers must supply +a compatible `kv_layout_id` in `engine.native_host_offload` 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. + +`engine.kv_transfer_bytes_per_token` describes P-to-D transfer size; +`engine.kv_cache_bytes_per_token` describes the physical cache footprint used by G2. +G2 requires a resolved, positive physical cache size. Set the transfer size +separately when the geometries differ. + +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). +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. + +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`. @@ -299,7 +347,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`. Enabled settings use the production `PlannerConfig` defaults and normalization. The default target is `throughput`, which enables diff --git a/lib/bindings/python/rust/llm/replay.rs b/lib/bindings/python/rust/llm/replay.rs index e55166ae4138..73e02a59e44d 100644 --- a/lib/bindings/python/rust/llm/replay.rs +++ b/lib/bindings/python/rust/llm/replay.rs @@ -15,8 +15,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; @@ -241,7 +241,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<'_>, @@ -279,6 +279,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| { @@ -313,6 +316,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 phases require 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 { @@ -416,7 +451,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( @@ -425,6 +460,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)?, @@ -442,6 +486,7 @@ pub fn run_mocker_trace_replay( num_workers, replay_concurrency, agentic_lanes, + agentic_options, &replay_mode, arrival_speedup_ratio, router_mode, @@ -648,6 +693,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, @@ -659,6 +705,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" @@ -758,11 +808,15 @@ 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() { + // Snapshot sampling uses recorded source time; normalizing + // here would change the sampled frontier and cache state. + 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, @@ -774,19 +828,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, @@ -797,6 +852,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 35a891c8c7e2..c67b80b54495 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -2342,6 +2342,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/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..09a68bf570cb --- /dev/null +++ b/lib/bindings/python/tests/replay/test_replay_agentic_timing.py @@ -0,0 +1,175 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json + +import pytest + +from dynamo._core import run_mocker_trace_replay +from dynamo.mocker.config import normalize_mocker_config +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]) + + +@pytest.mark.parametrize("g3_scope", ["worker_local", "cluster_shared"]) +@pytest.mark.parametrize("agentic_lanes", [None, 1]) +def test_native_agentic_replay_rejects_g3_after_canonical_config_validation( + tmp_path, g3_scope, agentic_lanes +): + trace = _write_timing_trace( + tmp_path / "g3.jsonl", + start_ms=0.0, + root_gap_ms=48.0, + dependency_delay_ms=32.0, + ) + # G3 is valid in the canonical engine schema, including its required G2 + # staging tier. AgentX must still reject it at the native replay boundary. + args = normalize_mocker_config( + { + "engine": { + "backend": "vllm", + "block_size": 64, + "kv_cache_bytes_per_token": 1024, + "native_host_offload": { + "scope": "dp_rank_local", + "num_host_blocks": 8, + }, + "g3_offload": {"scope": g3_scope, "num_g3_blocks": 16}, + }, + "dp_size": 1, + } + ) + # The binding exposes native replay errors as PyException. The exact error + # proves both explicit and trace-inferred AgentX reach the deployment guard. + with pytest.raises(Exception, match="^agentic replay does not support G3$"): + run_mocker_trace_replay( + [trace], + extra_engine_args=args, + replay_mode="offline", + trace_format="agentic_mooncake", + execution_model="fixture-model", + agentic_lanes=agentic_lanes, + ) diff --git a/lib/mocker/src/replay/agentic.rs b/lib/mocker/src/replay/agentic.rs new file mode 100644 index 000000000000..2ba65e694a0d --- /dev/null +++ b/lib/mocker/src/replay/agentic.rs @@ -0,0 +1,298 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// 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 { + // 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 + .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), + (Some(3), 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 eb80d53442c1..7bd5cd15fafc 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,14 @@ 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.g3_offload.is_none(), + "agentic replay does not support G3" + ); + 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 +2252,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: MockerConfig, + 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 +2300,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 +2377,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 +2422,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 +2441,14 @@ pub fn simulate_agentic_trace_live_workload_with_router_mode_and_options( agentic_lanes: Option, sla: SlaThresholds, ) -> Result { + anyhow::ensure!( + args.g3_offload.is_none(), + "agentic replay does not support G3" + ); + 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 eace36883067..2cd239f1a6bc 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, MockerConfig}; 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 b32ad86b88b4..56214306ba35 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; @@ -44,6 +43,66 @@ fn worker_pool(initial_workers: usize, args: &MockerConfig) -> WorkerPoolSpec { } } +fn validate_agentic_host_offload( + input: &ReplayRuntimeInput, + roles: &[(usize, &MockerConfig)], + scaling_enabled: bool, +) -> Result<()> { + let ReplayRuntimeInput::Workload(driver) = input else { + return Ok(()); + }; + if !driver.is_agentic() { + return Ok(()); + } + anyhow::ensure!( + roles.iter().all(|(_, args)| args.g3_offload.is_none()), + "agentic replay does not support G3" + ); + if !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.backend == EngineType::Vllm + && args.ais_perf_config().is_none_or(|config| { + config + .get("backend") + .is_none_or(|backend| backend == "vllm") + }), + "agentic host offload requires backend=vllm on every role" + ); + anyhow::ensure!( + args.dp_size == 1, + "agentic host offload requires attention DP=1 on every role" + ); + anyhow::ensure!( + args.decode_speedup_ratio == 1.0 + && args.aic_nextn.is_none() + && args.ais_perf_config().is_none_or(|config| { + config + .get("speculation") + .is_none_or(serde_json::Value::is_null) + }), + "agentic host offload requires speculative decoding disabled" + ); + } + Ok(()) +} + fn with_telemetry( replayer: Replayer, telemetry: Option, @@ -147,6 +206,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 +302,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 { @@ -577,16 +645,18 @@ 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 driver = agentic_options.into_driver( + trace, args.block_size, router_mode == ReplayRouterMode::KvRouter, agentic_lanes, + max_sim_time_ms, )?; run_aggregated( args, @@ -614,16 +684,18 @@ 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 driver = agentic_options.into_driver( + trace, config.prefill_args.block_size, router_mode == ReplayRouterMode::KvRouter, agentic_lanes, + max_sim_time_ms, )?; run_disaggregated( config, @@ -809,3 +881,270 @@ 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>) -> MockerConfig { + MockerConfig::from_value(serde_json::json!({ + "engine": { + "kv_cache_bytes_per_token": 1024, + "native_host_offload": scope.map(|scope| serde_json::json!({ + "scope": scope, + "num_host_blocks": 16, + "kv_layout_id": "agentic-g2-test", + })), + }, + })) + .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.tensor_parallel_size = 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", + )); + for backend in [EngineType::Sglang, EngineType::Trtllm] { + let mut invalid = args(None); + invalid.backend = backend; + invalid_roles.push(( + invalid, + "agentic host offload requires backend=vllm on every role", + )); + } + let mut invalid = args(None); + invalid.timing_model = aisimulate_core::engine::TimingModelConfig::External { + provider: "ais".to_string(), + config: serde_json::json!({"backend": "sglang"}), + }; + invalid_roles.push(( + invalid, + "agentic host offload requires backend=vllm on every role", + )); + let mut invalid = args(None); + invalid.aic_nextn = Some(1); + invalid_roles.push(( + invalid, + "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.timing_model = aisimulate_core::engine::TimingModelConfig::External { + provider: "ais".to_string(), + config: 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 [ + 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 native_execution_rejects_canonical_g3_for_agentic_replay() { + for scope in ["worker_local", "cluster_shared"] { + let mut offloaded = args(Some("cluster_shared")); + offloaded.g3_offload = Some( + serde_json::from_value(serde_json::json!({ + "scope": scope, + "num_g3_blocks": 8, + })) + .unwrap(), + ); + // G3 is valid in the canonical engine schema, but remains outside + // the qualified AgentX support boundary. + let offloaded = offloaded.normalized().unwrap(); + validate_agentic_host_offload( + &ReplayRuntimeInput::Requests(VecDeque::new()), + &[(1, &offloaded)], + false, + ) + .unwrap(); + let error = run_aggregated( + offloaded, + None, + None, + agentic_input(), + 1, + None, + ReplayRouterMode::RoundRobin, + false, + None, + SlaThresholds::default(), + None, + None, + ) + .unwrap_err(); + assert_eq!(error.to_string(), "agentic replay does not support G3"); + } + } + + #[test] + fn leaves_hbm_and_non_agentic_g2_deployments_unchanged() { + let mut hbm = args(None); + hbm.backend = 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.aic_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)], + 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 e4a35c5b3540..878ff082b700 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,9 @@ 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 == StorageTier::Device) .flat_map(|event| match &event.event.data { dynamo_kv_router::protocols::KvCacheEventData::Stored(store) => { @@ -250,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/tests/dependencies/test_aisimulate_namespaces.py b/tests/dependencies/test_aisimulate_namespaces.py index f8228c42e011..bd711267c9e6 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 @@ -192,3 +193,47 @@ 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_materializes_shared_host_offload() -> None: + if sys.version_info < (3, 11) or sys.version_info >= (3, 14): + pytest.skip("AISimulate supports Python 3.11 through 3.13") + + # Adapters consume the public launch contract; layout identity remains + # owned by AISimulate rather than a private helper imported by Dynamo. + from aisimulate.runner import materialize_engine_launch_config + + config = materialize_engine_launch_config( + "vllm", + "", + {}, + { + "aic_model_path": "resolved-model", + "num_gpu_blocks": 4, + "tensor_parallel_size": 2, + "kv_cache_bytes_per_token": 1024, + "kv_bytes_per_token": 4096, + "native_host_offload": { + "scope": "cluster_shared", + "num_host_blocks": 8, + "h2d_bandwidth_gbps": 7.0, + }, + "timing_model": {"type": "fixed", "prefill_ms": 1.0, "decode_ms": 1.0}, + }, + "aggregated", + ) + engine = config["engine"] + host = engine["native_host_offload"] + assert config["tensor_parallel_size"] == 2 + assert engine["kv_cache_bytes_per_token"] == 1024 + assert engine["kv_transfer_bytes_per_token"] == 4096 + assert host["scope"] == "cluster_shared" + assert host["num_host_blocks"] == 8 + assert host["h2d_bandwidth_gbps"] == 7.0 + assert json.loads(host["kv_layout_id"]) == { + "model": "resolved-model", + "backend": "vllm", + "tp": 2, + "block_size": engine["block_size"], + "bytes_per_token": 1024, + }