From 58028d5ba1cb33b0e135e18f1cc8833556a3639e Mon Sep 17 00:00:00 2001 From: Anthony Casagrande Date: Fri, 17 Jul 2026 11:35:37 -0700 Subject: [PATCH] feat(graph)!: native Dynamo trace replay -- Graph IR agentic workload lane with unified segment store and session routing Add a first-class agentic benchmarking lane built on Graph IR: a typed, validated intermediate representation of multi-turn, multi-agent conversation graphs that AIPerf ingests, lowers, and replays against an inference endpoint with recorded structure and timing preserved. Ingest and lowering - Graph IR schema and structural validation (node/edge typing, replay-output and branch rules, trie prompt convention) with actionable gate errors. - Adapters lower weka, dynamo, and dag_jsonl traces into Graph IR; parallel adapter variants (dynamo/weka trace_parallel) scale large-corpus parses. - Native lowering builds graphs directly without a source trace. Unified segment store - A single interned segment store per build (content pool + per-node manifests) is the sole graph store shape; the worker opens it lazily from the dataset broadcast and materializes each node's request payload from interned content, layering run-level endpoint options while keeping per-node dispatch overrides and stream settings winning. - Trie-based content interning deduplicates shared prefixes; a theoretical-prefix-cache accumulator emits the infinite-cache prefix hit rate without carrying hash ids through the request path. Runtime, timing, and routing - Async dataflow graph runtime dispatches nodes over the credit system with a cooperative duration deadline; replay timing honors recorded per-node delays and synthesis scaling. - Session routing keeps a trace's turns sticky to one worker; recorded dynamo session-identity headers are stripped when routing is active or uniquified per replay instance so concurrent instances never share a server session. - Graph first-token anchoring emits a per-credit FirstToken for post-TTFT observation independent of prefill-limit gating. Records pipeline and metrics - The records/post-processor pipeline routes records by record type through a per-request RecordsMessage envelope (producers emit typed records, observers act on them), reconciled with origin/main's route-by-record-type refactor. - Context-overflow records bypass the perf accumulators and error tracker but still advance the success counter and forward only the context_overflow_count metric so the submission-rate gate stays correct. Surface and tests - CLI/schema surface: graph_format auto-detection/override, synthesis config (--synthesis-max-osl), and phase autodefaults for graph/dag corpora. - GraphIRReplayStrategy with a phase-teardown hook that detaches observers and closes sticky trace lifecycles between phases. - Reference documentation for the schema, ingest/build pipeline, runtime, segment store, worker materialization, and troubleshooting, plus unit/integration/component test suites and fidelity tooling. Signed-off-by: Anthony Casagrande Co-Authored-By: Claude Opus 4.8 (1M context) --- .cursor/rules/python.mdc | 2 +- .github/copilot-instructions.md | 2 +- .github/workflows/run-unit-tests.yml | 4 +- AGENTS.md | 2 +- CLAUDE.md | 2 +- Makefile | 4 +- README.md | 3 +- docs/architecture.md | 4 +- docs/benchmark-datasets.md | 4 +- docs/benchmark-modes/agentic.md | 435 +++ docs/benchmark-modes/dag.md | 210 +- .../benchmark-modes/timing-modes-reference.md | 2 +- docs/cli-options.md | 176 +- docs/environment-variables.md | 45 + docs/index.yml | 26 +- docs/metrics-reference.md | 44 +- docs/plugins/plugin-system.md | 58 + .../reference/graph-async-dataflow-runtime.md | 532 ++++ docs/reference/graph-ingest-build-pipeline.md | 1093 ++++++++ docs/reference/graph-ir-schema.md | 287 ++ docs/reference/graph-ir-validation.md | 158 ++ .../graph-runtime-troubleshooting.md | 327 +++ docs/reference/graph-segment-unified-store.md | 305 +++ docs/reference/graph-structural-handoff.md | 222 ++ .../reference/graph-trie-prompt-convention.md | 56 + .../reference/graph-worker-materialization.md | 453 ++++ docs/reference/yaml-config-roadmap.md | 16 +- fern/docs.yml | 4 + pyproject.toml | 2 +- src/aiperf/api/api_service.py | 12 +- src/aiperf/cli.py | 1 + src/aiperf/cli_commands/_chat_stats.py | 3 + src/aiperf/cli_commands/dynamo.py | 25 + .../cli_commands/dynamo_trace_report.py | 475 ++++ src/aiperf/cli_runner/_aggregate.py | 212 +- src/aiperf/common/base_component_service.py | 8 +- src/aiperf/common/base_service.py | 4 +- src/aiperf/common/clock.py | 364 +++ src/aiperf/common/constants.py | 10 + src/aiperf/common/enums/__init__.py | 4 + src/aiperf/common/enums/enums.py | 37 +- src/aiperf/common/environment.py | 393 ++- src/aiperf/common/hash_id_random_generator.py | 89 + .../common/mixins/command_handler_mixin.py | 28 +- src/aiperf/common/mixins/hooks_mixin.py | 9 +- src/aiperf/common/mixins/message_bus_mixin.py | 4 +- .../common/mixins/progress_tracker_mixin.py | 4 +- src/aiperf/common/mixins/pull_client_mixin.py | 4 +- .../common/mixins/reply_client_mixin.py | 6 +- src/aiperf/common/models/dataset_models.py | 115 + src/aiperf/common/models/export_models.py | 32 + .../common/models/model_endpoint_info.py | 23 + src/aiperf/common/models/record_models.py | 49 + src/aiperf/common/random_generator.py | 11 + src/aiperf/common/scenario/__init__.py | 34 + src/aiperf/common/scenario/_env_locks.py | 251 ++ src/aiperf/common/scenario/base.py | 194 ++ .../common/scenario/context_overflow.py | 105 + .../common/scenario/inferencex_agentx_mvp.py | 37 + src/aiperf/common/scenario/registry.py | 22 + .../common/scenario/submission_outcome.py | 90 + src/aiperf/common/scenario/validator.py | 402 +++ src/aiperf/config/cli_parameter.py | 1 + src/aiperf/config/comm/base.py | 24 +- src/aiperf/config/config.py | 98 +- src/aiperf/config/dataset/config.py | 35 +- src/aiperf/config/dataset/resolver.py | 55 +- src/aiperf/config/dataset/trace.py | 61 +- src/aiperf/config/endpoint.py | 71 + src/aiperf/config/flags/_converter_dataset.py | 31 +- .../config/flags/_converter_endpoint.py | 19 + .../config/flags/_converter_profiling.py | 83 +- src/aiperf/config/flags/_section_fields.py | 8 + src/aiperf/config/flags/cli_config.py | 255 +- src/aiperf/config/flags/converter.py | 37 +- src/aiperf/config/flags/resolver.py | 3 + src/aiperf/config/loader/parsing.py | 33 +- src/aiperf/config/phases.py | 23 +- src/aiperf/config/resolution/__init__.py | 2 + .../resolution/graph_dispatch_resolver.py | 79 + src/aiperf/config/resolution/plan.py | 59 + src/aiperf/config/resolution/predicates.py | 44 + src/aiperf/config/resolution/resolvers.py | 23 + .../config/schema/aiperf-config.schema.json | 270 +- .../multiprocess_service_manager.py | 9 +- src/aiperf/credit/callback_handler.py | 135 +- src/aiperf/credit/issuer.py | 225 +- src/aiperf/credit/messages.py | 48 +- src/aiperf/credit/sticky_router.py | 106 +- src/aiperf/credit/structs.py | 93 +- src/aiperf/dataset/_mp_context.py | 151 ++ src/aiperf/dataset/_tokenizer_preload.py | 135 + src/aiperf/dataset/composer/base.py | 11 +- .../dataset/composer/synthetic_rankings.py | 4 +- src/aiperf/dataset/dataset_manager.py | 140 +- .../dataset/generator/_coding_cicd_docs.py | 439 +++ .../generator/_coding_conversations.py | 253 ++ .../_coding_conversations_advanced.py | 337 +++ .../dataset/generator/_coding_errors_diff.py | 382 +++ src/aiperf/dataset/generator/_coding_go.py | 231 ++ src/aiperf/dataset/generator/_coding_json.py | 122 + src/aiperf/dataset/generator/_coding_ml.py | 261 ++ .../dataset/generator/_coding_prompts_conv.py | 135 + .../dataset/generator/_coding_python.py | 280 ++ src/aiperf/dataset/generator/_coding_rust.py | 235 ++ src/aiperf/dataset/generator/_coding_sql.py | 135 + src/aiperf/dataset/generator/_coding_text.py | 329 +++ src/aiperf/dataset/generator/_coding_tool.py | 368 +++ .../dataset/generator/_coding_tool_long.py | 376 +++ .../dataset/generator/_coding_typescript.py | 221 ++ src/aiperf/dataset/generator/_coding_vocab.py | 345 +++ .../dataset/generator/coding_content.py | 243 ++ src/aiperf/dataset/generator/image.py | 4 +- src/aiperf/dataset/generator/prompt.py | 58 +- src/aiperf/dataset/graph/__init__.py | 8 + src/aiperf/dataset/graph/adapters/__init__.py | 20 + .../graph/adapters/dag_jsonl/__init__.py | 10 + .../graph/adapters/dag_jsonl/lowering.py | 247 ++ .../dataset/graph/adapters/dag_jsonl/trace.py | 183 ++ .../dataset/graph/adapters/dag_jsonl/tree.py | 325 +++ .../dataset/graph/adapters/dynamo/__init__.py | 17 + .../adapters/dynamo/store_backed_pool.py | 171 ++ .../dataset/graph/adapters/dynamo/trace.py | 1090 ++++++++ .../graph/adapters/dynamo/trace_parallel.py | 750 ++++++ .../graph/adapters/dynamo/trace_reader.py | 444 ++++ .../graph/adapters/dynamo/trie_lowering.py | 465 ++++ src/aiperf/dataset/graph/adapters/native.py | 41 + .../dataset/graph/adapters/protocols.py | 41 + .../dataset/graph/adapters/shared/__init__.py | 12 + .../dataset/graph/adapters/shared/content.py | 557 ++++ .../dataset/graph/adapters/shared/idle_gap.py | 37 + .../graph/adapters/shared/output_cap.py | 36 + .../graph/adapters/shared/peak_context.py | 106 + .../graph/adapters/shared/selection.py | 129 + .../dataset/graph/adapters/weka/__init__.py | 15 + .../dataset/graph/adapters/weka/trace.py | 913 +++++++ .../graph/adapters/weka/trace_models.py | 248 ++ .../graph/adapters/weka/trace_parallel.py | 629 +++++ .../dataset/graph/adapters/weka/trie_build.py | 458 ++++ src/aiperf/dataset/graph/auto_derive.py | 165 ++ src/aiperf/dataset/graph/codecs.py | 102 + src/aiperf/dataset/graph/decode.py | 184 ++ .../dataset/graph/graph_meta_sidecar.py | 166 ++ .../dataset/graph/graph_path_catalog.py | 77 + src/aiperf/dataset/graph/merge.py | 113 + src/aiperf/dataset/graph/models.py | 383 +++ src/aiperf/dataset/graph/native_lowering.py | 745 ++++++ src/aiperf/dataset/graph/parse_context.py | 110 + src/aiperf/dataset/graph/parser.py | 367 +++ .../dataset/graph/segment_ir/__init__.py | 10 + .../dataset/graph/segment_ir/envelope.py | 68 + .../graph/segment_ir/interval_order.py | 179 ++ src/aiperf/dataset/graph/segment_ir/pool.py | 168 ++ .../dataset/graph/segment_ir/prefix_cache.py | 121 + .../dataset/graph/segment_ir/store_builder.py | 521 ++++ .../dataset/graph/segment_ir/trie_content.py | 1149 ++++++++ src/aiperf/dataset/graph/store_build.py | 524 ++++ src/aiperf/dataset/graph/validator.py | 581 ++++ src/aiperf/dataset/graph/workload_detect.py | 494 ++++ src/aiperf/dataset/graph_client_store.py | 39 + .../dataset/graph_segment_unified_store.py | 534 ++++ src/aiperf/dataset/loader/base_hf_dataset.py | 4 +- src/aiperf/dataset/loader/exgentic.py | 4 +- src/aiperf/dataset/loader/hf_asr.py | 12 +- src/aiperf/dataset/loader/sharegpt.py | 4 +- src/aiperf/dataset/memory_map_utils.py | 20 +- .../aggregate/aggregate_base_exporter.py | 49 + .../aggregate_confidence_csv_exporter.py | 15 +- .../aggregate_confidence_json_exporter.py | 149 +- src/aiperf/exporters/metrics_csv_exporter.py | 4 +- src/aiperf/exporters/metrics_json_exporter.py | 64 +- .../timeslice_metrics_csv_exporter.py | 8 +- .../timeslice_metrics_json_exporter.py | 8 +- src/aiperf/graph/__init__.py | 8 + src/aiperf/graph/analysis/__init__.py | 23 + src/aiperf/graph/analysis/snapshot.py | 117 + src/aiperf/graph/analysis/timeline.py | 193 ++ src/aiperf/graph/channel_store.py | 431 +++ src/aiperf/graph/channels.py | 45 + src/aiperf/graph/context.py | 109 + src/aiperf/graph/credit_dispatch_adapter.py | 545 ++++ src/aiperf/graph/dispatch/__init__.py | 39 + src/aiperf/graph/dispatch/llm.py | 104 + src/aiperf/graph/dynamic_pool.py | 256 ++ src/aiperf/graph/executor.py | 603 +++++ src/aiperf/graph/placement.py | 41 + src/aiperf/graph/reducers.py | 115 + src/aiperf/graph/scheduler.py | 194 ++ src/aiperf/graph/worker_materialize.py | 533 ++++ .../types/context_overflow_count_metric.py | 44 + .../types/inter_chunk_latency_metric.py | 10 +- .../metrics/types/stream_latency_metrics.py | 5 +- .../types/streamed_request_count_metric.py | 50 + .../metrics/types/streamed_request_metric.py | 69 + .../types/theoretical_prefix_cache_metric.py | 42 + .../time_to_first_output_token_metric.py | 11 +- src/aiperf/metrics/types/ttft_metric.py | 7 +- src/aiperf/metrics/types/ttst_metric.py | 7 +- src/aiperf/network_latency/probe.py | 10 +- src/aiperf/orchestrator/local_executor.py | 87 +- src/aiperf/orchestrator/models.py | 17 + src/aiperf/plugin/categories.yaml | 33 + src/aiperf/plugin/enums.py | 16 +- src/aiperf/plugin/plugins.py | 12 +- src/aiperf/plugin/plugins.yaml | 116 +- src/aiperf/plugin/schema/plugins.schema.json | 100 +- src/aiperf/plugin/schema/schemas.py | 19 +- src/aiperf/plugin/types.py | 2 +- .../post_processors/base_metrics_processor.py | 22 +- .../theoretical_prefix_cache.py | 193 ++ src/aiperf/records/inference_result_parser.py | 57 +- .../records/record_processor_service.py | 33 + src/aiperf/records/records_manager.py | 72 +- .../_max_concurrency_under_sla.py | 2 +- src/aiperf/server_metrics/csv_exporter.py | 9 +- src/aiperf/server_metrics/data_collector.py | 4 +- src/aiperf/server_metrics/json_exporter.py | 8 +- src/aiperf/server_metrics/manager.py | 26 +- src/aiperf/server_metrics/parquet_exporter.py | 17 +- src/aiperf/server_metrics/storage.py | 4 +- .../timing/_branch_orchestrator_spawn.py | 53 +- .../timing/_branch_orchestrator_state.py | 10 + src/aiperf/timing/branch_orchestrator.py | 105 +- src/aiperf/timing/concurrency.py | 12 + src/aiperf/timing/config.py | 532 +++- src/aiperf/timing/conversation_source.py | 28 + src/aiperf/timing/graph_channel.py | 30 + src/aiperf/timing/graph_ir_source.py | 150 ++ src/aiperf/timing/graph_ir_trace_view.py | 56 + src/aiperf/timing/graph_warmup_handoff.py | 76 + src/aiperf/timing/manager.py | 189 +- src/aiperf/timing/phase/credit_counter.py | 69 +- src/aiperf/timing/phase/runner.py | 210 +- src/aiperf/timing/phase_orchestrator.py | 108 +- src/aiperf/timing/session_tree.py | 299 +++ src/aiperf/timing/snapshot_chop.py | 303 +++ src/aiperf/timing/strategies/cache_bust.py | 170 ++ src/aiperf/timing/strategies/core.py | 48 +- .../timing/strategies/fixed_schedule.py | 9 +- .../timing/strategies/graph_ir_replay.py | 2354 +++++++++++++++++ .../timing/strategies/user_centric_rate.py | 6 +- src/aiperf/transports/aiohttp_transport.py | 21 +- src/aiperf/transports/base_transports.py | 15 + .../dashboard/realtime_metrics_dashboard.py | 33 +- src/aiperf/workers/inference_client.py | 124 +- src/aiperf/workers/session_manager.py | 64 +- src/aiperf/workers/session_routing.py | 240 ++ src/aiperf/workers/worker.py | 713 ++++- src/aiperf/workers/worker_manager.py | 4 +- src/aiperf/zmq/dealer_request_client.py | 4 +- src/aiperf/zmq/router_reply_client.py | 9 +- src/aiperf/zmq/streaming_dealer_client.py | 15 +- src/aiperf/zmq/streaming_router_client.py | 14 +- src/aiperf/zmq/sub_client.py | 8 +- src/aiperf/zmq/zmq_base_client.py | 16 +- src/aiperf/zmq/zmq_proxy_base.py | 20 +- src/aiperf/zmq/zmq_proxy_sockets.py | 4 +- tests/component_integration/graph/conftest.py | 18 + .../graph/test_bare_run_single_pass.py | 211 ++ .../graph/test_dag_jsonl_byte_parity.py | 587 ++++ .../graph/test_dag_jsonl_e2e.py | 559 ++++ .../graph/test_dag_jsonl_fidelity.py | 417 +++ .../graph/test_dispatch_plumbing.py | 111 + .../graph/test_duplication_report.py | 350 +++ .../graph/test_dynamo_build_wiring.py | 95 + .../graph/test_dynamo_e2e_materialize.py | 292 ++ .../test_dynamo_subagent_only_dispatch.py | 158 ++ ...test_graph_ir_completion_and_accounting.py | 322 +++ .../graph/test_graph_ir_replay_strategy.py | 355 +++ .../graph/test_graph_sampling.py | 234 ++ .../graph/test_lane_ramp.py | 207 ++ .../graph/test_weka_trace_fidelity.py | 1200 +++++++++ .../graph/test_weka_trie_dispatch.py | 154 ++ .../graph/test_weka_trie_e2e_materialize.py | 149 ++ .../graph/test_weka_trie_hf_streaming.py | 215 ++ .../graph/test_wrap_guard.py | 336 +++ .../timing/test_dag_v1_adversarial.py | 10 +- .../dag/graph_parity/delayed_turn.dag.jsonl | 1 + .../dag/graph_parity/fork_minimal.dag.jsonl | 3 + .../dag/graph_parity/fork_toolcalls.dag.jsonl | 2 + .../dag/graph_parity/mixed_full.dag.jsonl | 3 + .../dag/graph_parity/prespawn.dag.jsonl | 2 + .../dag/graph_parity/spawn_join.dag.jsonl | 2 + .../dag/graph_parity/spawn_repeat.dag.jsonl | 2 + tests/fixtures/dag/two_spawn_trees.dag.jsonl | 4 + tests/harness/dynamo_synth_corpus.py | 216 ++ .../fixtures/dynamo_dir/_build_fixture.py | 131 + .../fixtures/dynamo_dir/trace.000000.jsonl.gz | Bin 0 -> 454 bytes .../native_dynamic/accumulate_chain.yaml | 47 + .../native_dynamic/planner_reviewer.yaml | 33 + .../weka_multigraph_dir/00_subagent.json | 27 + .../weka_multigraph_dir/01_plain.json | 11 + .../graph/test_dag_jsonl_fidelity_real.py | 212 ++ .../graph/test_dag_multiworker_placement.py | 109 + .../graph/test_dynamic_slots_mock_e2e.py | 173 ++ tests/integration/graph/test_dynamo_e2e.py | 108 + .../graph/test_first_token_live_timing.py | 584 ++++ .../graph/test_issue_1106_selection.py | 308 +++ .../graph/test_start_anchor_live_timing.py | 327 +++ .../graph/test_weka_cachebust_raw_export.py | 123 + .../graph/test_weka_isl_endpoint_osl_e2e.py | 132 + .../graph/test_weka_sticky_placement.py | 120 + .../graph/test_weka_trace_fidelity_real.py | 138 + .../graph/test_weka_unified_store_ab.py | 75 + .../test_session_routing_raw_export.py | 188 ++ .../cli_commands/test_dynamo_trace_report.py | 260 ++ .../test_aggregate_submission_writer.py | 398 +++ .../cli_runner/test_aggregation_round3.py | 36 +- .../cli_runner/test_request_count_override.py | 6 +- tests/unit/common/test_aiperf_logger.py | 24 +- tests/unit/common/test_clock.py | 228 ++ tests/unit/common/test_dynamo_settings.py | 16 + tests/unit/common/test_finite.py | 6 +- .../common/test_hash_id_random_generator.py | 45 + tests/unit/common/test_random_generator.py | 24 + tests/unit/config/test_camel_case_config.py | 4 +- .../config/test_cli_flags_graph_selection.py | 185 ++ .../config/test_converter_dataset_entries.py | 90 + .../test_converter_graph_single_pass.py | 236 ++ .../config/test_distribution_validators.py | 6 +- .../unit/config/test_dump_config_roundtrip.py | 6 +- .../config/test_graph_dispatch_resolver.py | 124 + tests/unit/config/test_graph_format_flag.py | 18 + .../config/test_graph_workload_resolution.py | 130 + tests/unit/config/test_plot_envelope.py | 10 +- tests/unit/config/test_resolver_edge_cases.py | 80 + tests/unit/config/test_resolvers.py | 15 +- tests/unit/config/test_scenario_cli_flags.py | 145 + tests/unit/config/test_scenario_resolver.py | 75 + .../config/test_session_routing_config.py | 89 + .../config/test_sweep_round2_adversarial.py | 12 +- .../config/test_sweep_round3_adversarial.py | 4 +- tests/unit/config/test_sweep_round3_fixes.py | 26 +- .../config/test_sweep_round4_adversarial.py | 10 +- .../credit/test_first_token_event_flag.py | 271 ++ .../credit/test_graph_sticky_lifecycle.py | 284 ++ tests/unit/credit/test_issuer_finality.py | 232 ++ .../generator/test_coding_content_hash_ids.py | 33 + tests/unit/dataset/graph/__init__.py | 2 + tests/unit/dataset/graph/adapters/__init__.py | 3 + .../graph/adapters/dag_jsonl/__init__.py | 2 + .../graph/adapters/dag_jsonl/test_lowering.py | 796 ++++++ .../graph/adapters/dag_jsonl/test_trace.py | 391 +++ .../graph/adapters/dag_jsonl/test_tree.py | 569 ++++ .../fixtures/dynamo_nested/_build_fixtures.py | 400 +++ .../dynamo_nested/cycle_AB_A.jsonl.gz | Bin 0 -> 273 bytes .../dynamo_nested/mixed_turn.jsonl.gz | Bin 0 -> 371 bytes .../dynamo_nested/nested_2_level.jsonl.gz | Bin 0 -> 269 bytes .../dynamo_nested/nested_3_level.jsonl.gz | Bin 0 -> 288 bytes .../dynamo_nested/parallel_subagents.jsonl.gz | Bin 0 -> 295 bytes .../dynamo_nested/parallel_two_root.jsonl.gz | Bin 0 -> 258 bytes .../tool_call_id_linkage.jsonl.gz | Bin 0 -> 382 bytes .../adapters/test_dynamo_causal_stamping.py | 114 + .../test_dynamo_corpus_scale_memory.py | 928 +++++++ .../graph/adapters/test_dynamo_gates.py | 110 + .../graph/adapters/test_dynamo_multigraph.py | 277 ++ .../adapters/test_dynamo_parallel_build.py | 470 ++++ .../adapters/test_dynamo_reader_schema.py | 148 ++ .../test_dynamo_recorded_streaming.py | 69 + .../adapters/test_dynamo_seed_alignment.py | 86 + .../graph/adapters/test_dynamo_trace.py | 809 ++++++ .../adapters/test_dynamo_trace_fidelity.py | 106 + .../test_dynamo_trace_nested_fixtures.py | 205 ++ .../adapters/test_dynamo_trace_reader.py | 478 ++++ .../graph/adapters/test_dynamo_tree_build.py | 346 +++ .../graph/adapters/test_dynamo_trie_ir.py | 257 ++ .../adapters/test_dynamo_trie_lowering.py | 570 ++++ .../adapters/test_dynamo_validate_gate.py | 154 ++ .../graph/adapters/test_peak_context.py | 134 + .../test_selection_filter_then_cap.py | 327 +++ .../graph/adapters/test_weka_dynamo_parity.py | 747 ++++++ .../unit/dataset/graph/segment_ir/__init__.py | 2 + .../graph/segment_ir/test_interval_order.py | 50 + .../dataset/graph/segment_ir/test_pool.py | 77 + .../graph/segment_ir/test_prefix_cache.py | 263 ++ .../graph/segment_ir/test_store_builder.py | 182 ++ .../graph/segment_ir/test_trie_content.py | 275 ++ .../graph/test_catalog_presence_predicate.py | 32 + .../graph/test_edge_delay_exclusivity.py | 75 + .../graph/test_first_token_anchor_rule.py | 106 + .../dataset/graph/test_graph_meta_sidecar.py | 62 + .../dataset/graph/test_graph_meta_writer.py | 91 + .../dataset/graph/test_graph_parse_context.py | 64 + .../dataset/graph/test_native_lowering.py | 644 +++++ .../graph/test_native_lowering_slots.py | 537 ++++ .../graph/test_parser_decode_strictness.py | 341 +++ .../unit/dataset/graph/test_parser_detect.py | 281 ++ .../dataset/graph/test_validator_rules.py | 252 ++ .../test_workload_detect_path_predicate.py | 31 + tests/unit/dataset/loader/test_exgentic.py | 8 - .../test_dag_jsonl_streaming_store_parity.py | 273 ++ .../dataset/test_dynamo_direct_store_route.py | 387 +++ .../test_dynamo_store_golden_digest.py | 178 ++ .../test_dynamo_streaming_store_parity.py | 464 ++++ .../dataset/test_graph_client_metadata.py | 75 + .../test_graph_segment_unified_store.py | 311 +++ .../dataset/test_graph_store_build_offload.py | 150 ++ .../dataset/test_graph_store_build_stats.py | 266 ++ .../unit/dataset/test_graph_store_builder.py | 102 + tests/unit/dataset/test_loader_mp_context.py | 257 ++ .../dataset/test_nonweka_interned_route.py | 334 +++ .../dataset/test_streaming_trie_sidecar.py | 78 + tests/unit/dataset/test_trie_unified_build.py | 128 + .../unit/dataset/test_unified_flag_wiring.py | 208 ++ .../unit/dataset/test_weka_streaming_gate.py | 66 + .../test_aggregate_submission_outcome.py | 223 ++ .../aggregate/test_sweep_exporters.py | 8 +- tests/unit/exporters/test_console_exporter.py | 43 + .../exporters/test_metrics_json_exporter.py | 232 ++ .../fixtures/weka_first_token_anchor.json | 16 + tests/unit/graph/fixtures/weka_min.json | 11 + .../graph/fixtures/weka_start_anchor.json | 16 + tests/unit/graph/fixtures/weka_subagent.json | 27 + tests/unit/graph/test_accelerated_warmup.py | 179 ++ tests/unit/graph/test_analysis_core.py | 173 ++ .../graph/test_build_strategy_graph_ir.py | 164 ++ tests/unit/graph/test_cache_bust_marker.py | 212 ++ .../graph/test_channel_store_reducer_enum.py | 56 + .../test_collapse_leading_start_offsets.py | 147 + .../graph/test_content_seed_determinism.py | 206 ++ .../graph/test_corpus_content_synthesizer.py | 44 + .../graph/test_credit_counter_graph_bypass.py | 180 ++ .../graph/test_credit_dispatch_adapter.py | 696 +++++ tests/unit/graph/test_credit_struct_fields.py | 39 + .../graph/test_default_pool_model_override.py | 116 + tests/unit/graph/test_dynamic_pool.py | 229 ++ tests/unit/graph/test_dynamic_slots_e2e.py | 500 ++++ tests/unit/graph/test_effective_root_seed.py | 41 + tests/unit/graph/test_executor_runs_weka.py | 434 +++ tests/unit/graph/test_executor_watchdog.py | 113 + .../graph/test_first_token_anchor_lowering.py | 215 ++ .../graph/test_first_token_fidelity_tool.py | 237 ++ tests/unit/graph/test_first_token_routing.py | 368 +++ tests/unit/graph/test_first_token_runtime.py | 443 ++++ tests/unit/graph/test_graph_endpoint_guard.py | 55 + .../graph/test_graph_instance_stop_paths.py | 231 ++ tests/unit/graph/test_graph_parse_context.py | 378 +++ tests/unit/graph/test_graph_path_catalog.py | 40 + tests/unit/graph/test_graph_return_bridge.py | 111 + tests/unit/graph/test_graph_sticky_stamp.py | 133 + tests/unit/graph/test_hf_sidecar_scope_out.py | 32 + .../graph/test_hf_streaming_trie_stores.py | 405 +++ tests/unit/graph/test_issuer_graph_path.py | 296 +++ tests/unit/graph/test_lane_fanout_recycle.py | 314 +++ tests/unit/graph/test_mixed_anchor_gate.py | 219 ++ .../unit/graph/test_pressure_frontier_chop.py | 288 ++ tests/unit/graph/test_pressure_stage.py | 662 +++++ tests/unit/graph/test_r6_wiring.py | 68 + tests/unit/graph/test_real_content_guard.py | 64 + .../test_recorded_mode_stream_override.py | 111 + .../graph/test_registry_run_parse_parity.py | 97 + .../graph/test_runtime_hardening_round3.py | 130 + .../unit/graph/test_slot_graph_eager_drain.py | 174 ++ tests/unit/graph/test_start_anchor_edges.py | 95 + .../graph/test_start_anchor_fidelity_tool.py | 264 ++ tests/unit/graph/test_start_anchor_runtime.py | 291 ++ .../unit/graph/test_start_anchor_snapshot.py | 184 ++ .../graph/test_start_anchor_weka_stamping.py | 185 ++ .../graph/test_trie_prefix_reuse_oracle.py | 669 +++++ tests/unit/graph/test_tstar_activation.py | 134 + tests/unit/graph/test_tstar_selection.py | 110 + .../graph/test_typed_channel_placeholders.py | 154 ++ .../test_unified_interned_materialize.py | 139 + .../graph/test_warmup_boundary_rewrite.py | 236 ++ .../unit/graph/test_warmup_handoff_consume.py | 448 ++++ .../graph/test_warmup_routing_and_seed.py | 372 +++ tests/unit/graph/test_warmup_variants.py | 90 + .../unit/graph/test_weka_block_size_decode.py | 103 + .../graph/test_weka_content_knob_wiring.py | 114 + .../graph/test_weka_fidelity_tool_gates.py | 121 + tests/unit/graph/test_weka_hf_real_ingest.py | 222 ++ tests/unit/graph/test_weka_ingest.py | 130 + .../test_weka_isl_endpoint_osl_fidelity.py | 166 ++ .../graph/test_weka_memory_regressions.py | 240 ++ .../graph/test_weka_parallel_hardening.py | 198 ++ .../test_weka_parse_kwargs_unification.py | 68 + .../test_weka_segment_pool_merge_codec.py | 139 + tests/unit/graph/test_weka_segments.py | 36 + .../graph/test_weka_trace_models_finite.py | 73 + tests/unit/graph/test_weka_trie_build.py | 482 ++++ .../graph/test_weka_trie_build_resolution.py | 374 +++ tests/unit/graph/test_weka_trie_hash_scope.py | 181 ++ .../unit/graph/test_weka_trie_idle_gap_cap.py | 330 +++ .../graph/test_weka_trie_interval_order.py | 540 ++++ tests/unit/graph/test_weka_trie_route.py | 53 + tests/unit/graph/test_weka_trie_snapshot.py | 315 +++ .../unit/graph/test_weka_unified_dispatch.py | 144 + tests/unit/graph/test_worker_bytes_path.py | 485 ++++ tests/unit/graph/test_worker_graph_branch.py | 256 ++ tests/unit/graph/test_worker_graph_capture.py | 377 +++ tests/unit/graph/test_worker_materialize.py | 84 + .../graph/test_worker_payload_features.py | 700 +++++ .../graph/test_workload_detect_predicates.py | 318 +++ tests/unit/metrics/conftest.py | 17 + .../test_inter_chunk_latency_metric.py | 17 +- .../metrics/test_streamed_request_count.py | 201 ++ .../test_theoretical_prefix_cache_metric.py | 50 + .../test_time_to_first_output_metric.py | 10 +- tests/unit/metrics/test_ttst_metric.py | 11 +- .../search_planner/test_monotonic.py | 2 +- .../test_smooth_isotonic_v2_wiring.py | 4 +- tests/unit/orchestrator/test_cell_callback.py | 4 +- .../unit/orchestrator/test_local_executor.py | 147 +- tests/unit/plot/test_dashboard_utils.py | 4 +- .../plugin/test_graph_adapter_category.py | 35 + .../plugin/test_session_routing_registry.py | 30 + .../post_processors/test_streaming_gate.py | 123 + .../test_theoretical_prefix_cache.py | 222 ++ .../property/_numeric_bounds_baseline.txt | 21 + tests/unit/property/test_finite_invariants.py | 7 +- .../unit/property/test_pydantic_field_fuzz.py | 2 +- tests/unit/records/test_records_manager.py | 107 + tests/unit/scenario/test_context_overflow.py | 242 ++ tests/unit/scenario/test_env_locks.py | 168 ++ tests/unit/scenario/test_scenario_registry.py | 75 + .../unit/scenario/test_scenario_validator.py | 429 +++ .../search_recipes/test_recipes_round3.py | 4 +- .../test_recipes_round3_followup.py | 9 +- .../test_server_metrics_data_collector.py | 2 +- tests/unit/test_cli_runner_sweep_helpers.py | 6 +- tests/unit/timing/conftest.py | 8 + tests/unit/timing/phase/test_runner.py | 306 +++ .../test_graph_observer_teardown.py | 145 + .../strategies/test_user_centric_rate.py | 71 + tests/unit/timing/test_graph_phase_channel.py | 94 + .../unit/timing/test_session_tree_finality.py | 160 ++ tests/unit/timing/test_session_tree_wiring.py | 487 ++++ ...est_timing_manager_first_token_advisory.py | 149 ++ .../test_timing_manager_sidecar_load.py | 141 + .../unit/transports/test_aiohttp_transport.py | 79 +- ...est_aiohttp_transport_bytes_passthrough.py | 74 + .../ui/test_realtime_metrics_dashboard.py | 70 + tests/unit/workers/test_inference_client.py | 350 +++ .../workers/test_session_fork_refcount.py | 87 + tests/unit/workers/test_session_routing.py | 205 ++ tests/unit/workers/test_streamed_stamp.py | 128 + tests/unit/workers/test_worker.py | 183 +- .../test_worker_graph_predispatch_errors.py | 269 ++ .../test_worker_graph_store_discovery.py | 55 + tools/dag_jsonl_fidelity.py | 497 ++++ tools/ergonomics_baseline.json | 55 + tools/generate_cli_docs.py | 12 +- tools/ruff_baseline.json | 40 + tools/weka_trace_fidelity.py | 1270 +++++++++ tools/weka_trie_timing_sim.py | 304 +++ 545 files changed, 92023 insertions(+), 632 deletions(-) create mode 100644 docs/benchmark-modes/agentic.md create mode 100644 docs/reference/graph-async-dataflow-runtime.md create mode 100644 docs/reference/graph-ingest-build-pipeline.md create mode 100644 docs/reference/graph-ir-schema.md create mode 100644 docs/reference/graph-ir-validation.md create mode 100644 docs/reference/graph-runtime-troubleshooting.md create mode 100644 docs/reference/graph-segment-unified-store.md create mode 100644 docs/reference/graph-structural-handoff.md create mode 100644 docs/reference/graph-trie-prompt-convention.md create mode 100644 docs/reference/graph-worker-materialization.md create mode 100644 src/aiperf/cli_commands/dynamo.py create mode 100644 src/aiperf/cli_commands/dynamo_trace_report.py create mode 100644 src/aiperf/common/clock.py create mode 100644 src/aiperf/common/hash_id_random_generator.py create mode 100644 src/aiperf/common/scenario/__init__.py create mode 100644 src/aiperf/common/scenario/_env_locks.py create mode 100644 src/aiperf/common/scenario/base.py create mode 100644 src/aiperf/common/scenario/context_overflow.py create mode 100644 src/aiperf/common/scenario/inferencex_agentx_mvp.py create mode 100644 src/aiperf/common/scenario/registry.py create mode 100644 src/aiperf/common/scenario/submission_outcome.py create mode 100644 src/aiperf/common/scenario/validator.py create mode 100644 src/aiperf/config/resolution/graph_dispatch_resolver.py create mode 100644 src/aiperf/dataset/_mp_context.py create mode 100644 src/aiperf/dataset/_tokenizer_preload.py create mode 100644 src/aiperf/dataset/generator/_coding_cicd_docs.py create mode 100644 src/aiperf/dataset/generator/_coding_conversations.py create mode 100644 src/aiperf/dataset/generator/_coding_conversations_advanced.py create mode 100644 src/aiperf/dataset/generator/_coding_errors_diff.py create mode 100644 src/aiperf/dataset/generator/_coding_go.py create mode 100644 src/aiperf/dataset/generator/_coding_json.py create mode 100644 src/aiperf/dataset/generator/_coding_ml.py create mode 100644 src/aiperf/dataset/generator/_coding_prompts_conv.py create mode 100644 src/aiperf/dataset/generator/_coding_python.py create mode 100644 src/aiperf/dataset/generator/_coding_rust.py create mode 100644 src/aiperf/dataset/generator/_coding_sql.py create mode 100644 src/aiperf/dataset/generator/_coding_text.py create mode 100644 src/aiperf/dataset/generator/_coding_tool.py create mode 100644 src/aiperf/dataset/generator/_coding_tool_long.py create mode 100644 src/aiperf/dataset/generator/_coding_typescript.py create mode 100644 src/aiperf/dataset/generator/_coding_vocab.py create mode 100644 src/aiperf/dataset/generator/coding_content.py create mode 100644 src/aiperf/dataset/graph/__init__.py create mode 100644 src/aiperf/dataset/graph/adapters/__init__.py create mode 100644 src/aiperf/dataset/graph/adapters/dag_jsonl/__init__.py create mode 100644 src/aiperf/dataset/graph/adapters/dag_jsonl/lowering.py create mode 100644 src/aiperf/dataset/graph/adapters/dag_jsonl/trace.py create mode 100644 src/aiperf/dataset/graph/adapters/dag_jsonl/tree.py create mode 100644 src/aiperf/dataset/graph/adapters/dynamo/__init__.py create mode 100644 src/aiperf/dataset/graph/adapters/dynamo/store_backed_pool.py create mode 100644 src/aiperf/dataset/graph/adapters/dynamo/trace.py create mode 100644 src/aiperf/dataset/graph/adapters/dynamo/trace_parallel.py create mode 100644 src/aiperf/dataset/graph/adapters/dynamo/trace_reader.py create mode 100644 src/aiperf/dataset/graph/adapters/dynamo/trie_lowering.py create mode 100644 src/aiperf/dataset/graph/adapters/native.py create mode 100644 src/aiperf/dataset/graph/adapters/protocols.py create mode 100644 src/aiperf/dataset/graph/adapters/shared/__init__.py create mode 100644 src/aiperf/dataset/graph/adapters/shared/content.py create mode 100644 src/aiperf/dataset/graph/adapters/shared/idle_gap.py create mode 100644 src/aiperf/dataset/graph/adapters/shared/output_cap.py create mode 100644 src/aiperf/dataset/graph/adapters/shared/peak_context.py create mode 100644 src/aiperf/dataset/graph/adapters/shared/selection.py create mode 100644 src/aiperf/dataset/graph/adapters/weka/__init__.py create mode 100644 src/aiperf/dataset/graph/adapters/weka/trace.py create mode 100644 src/aiperf/dataset/graph/adapters/weka/trace_models.py create mode 100644 src/aiperf/dataset/graph/adapters/weka/trace_parallel.py create mode 100644 src/aiperf/dataset/graph/adapters/weka/trie_build.py create mode 100644 src/aiperf/dataset/graph/auto_derive.py create mode 100644 src/aiperf/dataset/graph/codecs.py create mode 100644 src/aiperf/dataset/graph/decode.py create mode 100644 src/aiperf/dataset/graph/graph_meta_sidecar.py create mode 100644 src/aiperf/dataset/graph/graph_path_catalog.py create mode 100644 src/aiperf/dataset/graph/merge.py create mode 100644 src/aiperf/dataset/graph/models.py create mode 100644 src/aiperf/dataset/graph/native_lowering.py create mode 100644 src/aiperf/dataset/graph/parse_context.py create mode 100644 src/aiperf/dataset/graph/parser.py create mode 100644 src/aiperf/dataset/graph/segment_ir/__init__.py create mode 100644 src/aiperf/dataset/graph/segment_ir/envelope.py create mode 100644 src/aiperf/dataset/graph/segment_ir/interval_order.py create mode 100644 src/aiperf/dataset/graph/segment_ir/pool.py create mode 100644 src/aiperf/dataset/graph/segment_ir/prefix_cache.py create mode 100644 src/aiperf/dataset/graph/segment_ir/store_builder.py create mode 100644 src/aiperf/dataset/graph/segment_ir/trie_content.py create mode 100644 src/aiperf/dataset/graph/store_build.py create mode 100644 src/aiperf/dataset/graph/validator.py create mode 100644 src/aiperf/dataset/graph/workload_detect.py create mode 100644 src/aiperf/dataset/graph_client_store.py create mode 100644 src/aiperf/dataset/graph_segment_unified_store.py create mode 100644 src/aiperf/graph/__init__.py create mode 100644 src/aiperf/graph/analysis/__init__.py create mode 100644 src/aiperf/graph/analysis/snapshot.py create mode 100644 src/aiperf/graph/analysis/timeline.py create mode 100644 src/aiperf/graph/channel_store.py create mode 100644 src/aiperf/graph/channels.py create mode 100644 src/aiperf/graph/context.py create mode 100644 src/aiperf/graph/credit_dispatch_adapter.py create mode 100644 src/aiperf/graph/dispatch/__init__.py create mode 100644 src/aiperf/graph/dispatch/llm.py create mode 100644 src/aiperf/graph/dynamic_pool.py create mode 100644 src/aiperf/graph/executor.py create mode 100644 src/aiperf/graph/placement.py create mode 100644 src/aiperf/graph/reducers.py create mode 100644 src/aiperf/graph/scheduler.py create mode 100644 src/aiperf/graph/worker_materialize.py create mode 100644 src/aiperf/metrics/types/context_overflow_count_metric.py create mode 100644 src/aiperf/metrics/types/streamed_request_count_metric.py create mode 100644 src/aiperf/metrics/types/streamed_request_metric.py create mode 100644 src/aiperf/metrics/types/theoretical_prefix_cache_metric.py create mode 100644 src/aiperf/post_processors/theoretical_prefix_cache.py create mode 100644 src/aiperf/timing/graph_channel.py create mode 100644 src/aiperf/timing/graph_ir_source.py create mode 100644 src/aiperf/timing/graph_ir_trace_view.py create mode 100644 src/aiperf/timing/graph_warmup_handoff.py create mode 100644 src/aiperf/timing/session_tree.py create mode 100644 src/aiperf/timing/snapshot_chop.py create mode 100644 src/aiperf/timing/strategies/cache_bust.py create mode 100644 src/aiperf/timing/strategies/graph_ir_replay.py create mode 100644 src/aiperf/workers/session_routing.py create mode 100644 tests/component_integration/graph/conftest.py create mode 100644 tests/component_integration/graph/test_bare_run_single_pass.py create mode 100644 tests/component_integration/graph/test_dag_jsonl_byte_parity.py create mode 100644 tests/component_integration/graph/test_dag_jsonl_e2e.py create mode 100644 tests/component_integration/graph/test_dag_jsonl_fidelity.py create mode 100644 tests/component_integration/graph/test_dispatch_plumbing.py create mode 100644 tests/component_integration/graph/test_duplication_report.py create mode 100644 tests/component_integration/graph/test_dynamo_build_wiring.py create mode 100644 tests/component_integration/graph/test_dynamo_e2e_materialize.py create mode 100644 tests/component_integration/graph/test_dynamo_subagent_only_dispatch.py create mode 100644 tests/component_integration/graph/test_graph_ir_completion_and_accounting.py create mode 100644 tests/component_integration/graph/test_graph_ir_replay_strategy.py create mode 100644 tests/component_integration/graph/test_graph_sampling.py create mode 100644 tests/component_integration/graph/test_lane_ramp.py create mode 100644 tests/component_integration/graph/test_weka_trace_fidelity.py create mode 100644 tests/component_integration/graph/test_weka_trie_dispatch.py create mode 100644 tests/component_integration/graph/test_weka_trie_e2e_materialize.py create mode 100644 tests/component_integration/graph/test_weka_trie_hf_streaming.py create mode 100644 tests/component_integration/graph/test_wrap_guard.py create mode 100644 tests/fixtures/dag/graph_parity/delayed_turn.dag.jsonl create mode 100644 tests/fixtures/dag/graph_parity/fork_minimal.dag.jsonl create mode 100644 tests/fixtures/dag/graph_parity/fork_toolcalls.dag.jsonl create mode 100644 tests/fixtures/dag/graph_parity/mixed_full.dag.jsonl create mode 100644 tests/fixtures/dag/graph_parity/prespawn.dag.jsonl create mode 100644 tests/fixtures/dag/graph_parity/spawn_join.dag.jsonl create mode 100644 tests/fixtures/dag/graph_parity/spawn_repeat.dag.jsonl create mode 100644 tests/fixtures/dag/two_spawn_trees.dag.jsonl create mode 100644 tests/harness/dynamo_synth_corpus.py create mode 100644 tests/integration/graph/fixtures/dynamo_dir/_build_fixture.py create mode 100644 tests/integration/graph/fixtures/dynamo_dir/trace.000000.jsonl.gz create mode 100644 tests/integration/graph/fixtures/native_dynamic/accumulate_chain.yaml create mode 100644 tests/integration/graph/fixtures/native_dynamic/planner_reviewer.yaml create mode 100644 tests/integration/graph/fixtures/weka_multigraph_dir/00_subagent.json create mode 100644 tests/integration/graph/fixtures/weka_multigraph_dir/01_plain.json create mode 100644 tests/integration/graph/test_dag_jsonl_fidelity_real.py create mode 100644 tests/integration/graph/test_dag_multiworker_placement.py create mode 100644 tests/integration/graph/test_dynamic_slots_mock_e2e.py create mode 100644 tests/integration/graph/test_dynamo_e2e.py create mode 100644 tests/integration/graph/test_first_token_live_timing.py create mode 100644 tests/integration/graph/test_issue_1106_selection.py create mode 100644 tests/integration/graph/test_start_anchor_live_timing.py create mode 100644 tests/integration/graph/test_weka_cachebust_raw_export.py create mode 100644 tests/integration/graph/test_weka_isl_endpoint_osl_e2e.py create mode 100644 tests/integration/graph/test_weka_sticky_placement.py create mode 100644 tests/integration/graph/test_weka_trace_fidelity_real.py create mode 100644 tests/integration/graph/test_weka_unified_store_ab.py create mode 100644 tests/integration/test_session_routing_raw_export.py create mode 100644 tests/unit/cli_commands/test_dynamo_trace_report.py create mode 100644 tests/unit/cli_runner/test_aggregate_submission_writer.py create mode 100644 tests/unit/common/test_clock.py create mode 100644 tests/unit/common/test_dynamo_settings.py create mode 100644 tests/unit/common/test_hash_id_random_generator.py create mode 100644 tests/unit/config/test_cli_flags_graph_selection.py create mode 100644 tests/unit/config/test_converter_dataset_entries.py create mode 100644 tests/unit/config/test_converter_graph_single_pass.py create mode 100644 tests/unit/config/test_graph_dispatch_resolver.py create mode 100644 tests/unit/config/test_graph_format_flag.py create mode 100644 tests/unit/config/test_graph_workload_resolution.py create mode 100644 tests/unit/config/test_scenario_cli_flags.py create mode 100644 tests/unit/config/test_scenario_resolver.py create mode 100644 tests/unit/config/test_session_routing_config.py create mode 100644 tests/unit/credit/test_first_token_event_flag.py create mode 100644 tests/unit/credit/test_graph_sticky_lifecycle.py create mode 100644 tests/unit/credit/test_issuer_finality.py create mode 100644 tests/unit/dataset/generator/test_coding_content_hash_ids.py create mode 100644 tests/unit/dataset/graph/__init__.py create mode 100644 tests/unit/dataset/graph/adapters/__init__.py create mode 100644 tests/unit/dataset/graph/adapters/dag_jsonl/__init__.py create mode 100644 tests/unit/dataset/graph/adapters/dag_jsonl/test_lowering.py create mode 100644 tests/unit/dataset/graph/adapters/dag_jsonl/test_trace.py create mode 100644 tests/unit/dataset/graph/adapters/dag_jsonl/test_tree.py create mode 100644 tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/_build_fixtures.py create mode 100644 tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/cycle_AB_A.jsonl.gz create mode 100644 tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/mixed_turn.jsonl.gz create mode 100644 tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/nested_2_level.jsonl.gz create mode 100644 tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/nested_3_level.jsonl.gz create mode 100644 tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/parallel_subagents.jsonl.gz create mode 100644 tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/parallel_two_root.jsonl.gz create mode 100644 tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/tool_call_id_linkage.jsonl.gz create mode 100644 tests/unit/dataset/graph/adapters/test_dynamo_causal_stamping.py create mode 100644 tests/unit/dataset/graph/adapters/test_dynamo_corpus_scale_memory.py create mode 100644 tests/unit/dataset/graph/adapters/test_dynamo_gates.py create mode 100644 tests/unit/dataset/graph/adapters/test_dynamo_multigraph.py create mode 100644 tests/unit/dataset/graph/adapters/test_dynamo_parallel_build.py create mode 100644 tests/unit/dataset/graph/adapters/test_dynamo_reader_schema.py create mode 100644 tests/unit/dataset/graph/adapters/test_dynamo_recorded_streaming.py create mode 100644 tests/unit/dataset/graph/adapters/test_dynamo_seed_alignment.py create mode 100644 tests/unit/dataset/graph/adapters/test_dynamo_trace.py create mode 100644 tests/unit/dataset/graph/adapters/test_dynamo_trace_fidelity.py create mode 100644 tests/unit/dataset/graph/adapters/test_dynamo_trace_nested_fixtures.py create mode 100644 tests/unit/dataset/graph/adapters/test_dynamo_trace_reader.py create mode 100644 tests/unit/dataset/graph/adapters/test_dynamo_tree_build.py create mode 100644 tests/unit/dataset/graph/adapters/test_dynamo_trie_ir.py create mode 100644 tests/unit/dataset/graph/adapters/test_dynamo_trie_lowering.py create mode 100644 tests/unit/dataset/graph/adapters/test_dynamo_validate_gate.py create mode 100644 tests/unit/dataset/graph/adapters/test_peak_context.py create mode 100644 tests/unit/dataset/graph/adapters/test_selection_filter_then_cap.py create mode 100644 tests/unit/dataset/graph/adapters/test_weka_dynamo_parity.py create mode 100644 tests/unit/dataset/graph/segment_ir/__init__.py create mode 100644 tests/unit/dataset/graph/segment_ir/test_interval_order.py create mode 100644 tests/unit/dataset/graph/segment_ir/test_pool.py create mode 100644 tests/unit/dataset/graph/segment_ir/test_prefix_cache.py create mode 100644 tests/unit/dataset/graph/segment_ir/test_store_builder.py create mode 100644 tests/unit/dataset/graph/segment_ir/test_trie_content.py create mode 100644 tests/unit/dataset/graph/test_catalog_presence_predicate.py create mode 100644 tests/unit/dataset/graph/test_edge_delay_exclusivity.py create mode 100644 tests/unit/dataset/graph/test_first_token_anchor_rule.py create mode 100644 tests/unit/dataset/graph/test_graph_meta_sidecar.py create mode 100644 tests/unit/dataset/graph/test_graph_meta_writer.py create mode 100644 tests/unit/dataset/graph/test_graph_parse_context.py create mode 100644 tests/unit/dataset/graph/test_native_lowering.py create mode 100644 tests/unit/dataset/graph/test_native_lowering_slots.py create mode 100644 tests/unit/dataset/graph/test_parser_decode_strictness.py create mode 100644 tests/unit/dataset/graph/test_parser_detect.py create mode 100644 tests/unit/dataset/graph/test_validator_rules.py create mode 100644 tests/unit/dataset/graph/test_workload_detect_path_predicate.py create mode 100644 tests/unit/dataset/test_dag_jsonl_streaming_store_parity.py create mode 100644 tests/unit/dataset/test_dynamo_direct_store_route.py create mode 100644 tests/unit/dataset/test_dynamo_store_golden_digest.py create mode 100644 tests/unit/dataset/test_dynamo_streaming_store_parity.py create mode 100644 tests/unit/dataset/test_graph_client_metadata.py create mode 100644 tests/unit/dataset/test_graph_segment_unified_store.py create mode 100644 tests/unit/dataset/test_graph_store_build_offload.py create mode 100644 tests/unit/dataset/test_graph_store_build_stats.py create mode 100644 tests/unit/dataset/test_graph_store_builder.py create mode 100644 tests/unit/dataset/test_loader_mp_context.py create mode 100644 tests/unit/dataset/test_nonweka_interned_route.py create mode 100644 tests/unit/dataset/test_streaming_trie_sidecar.py create mode 100644 tests/unit/dataset/test_trie_unified_build.py create mode 100644 tests/unit/dataset/test_unified_flag_wiring.py create mode 100644 tests/unit/dataset/test_weka_streaming_gate.py create mode 100644 tests/unit/exporters/aggregate/test_aggregate_submission_outcome.py create mode 100644 tests/unit/graph/fixtures/weka_first_token_anchor.json create mode 100644 tests/unit/graph/fixtures/weka_min.json create mode 100644 tests/unit/graph/fixtures/weka_start_anchor.json create mode 100644 tests/unit/graph/fixtures/weka_subagent.json create mode 100644 tests/unit/graph/test_accelerated_warmup.py create mode 100644 tests/unit/graph/test_analysis_core.py create mode 100644 tests/unit/graph/test_build_strategy_graph_ir.py create mode 100644 tests/unit/graph/test_cache_bust_marker.py create mode 100644 tests/unit/graph/test_channel_store_reducer_enum.py create mode 100644 tests/unit/graph/test_collapse_leading_start_offsets.py create mode 100644 tests/unit/graph/test_content_seed_determinism.py create mode 100644 tests/unit/graph/test_corpus_content_synthesizer.py create mode 100644 tests/unit/graph/test_credit_counter_graph_bypass.py create mode 100644 tests/unit/graph/test_credit_dispatch_adapter.py create mode 100644 tests/unit/graph/test_credit_struct_fields.py create mode 100644 tests/unit/graph/test_default_pool_model_override.py create mode 100644 tests/unit/graph/test_dynamic_pool.py create mode 100644 tests/unit/graph/test_dynamic_slots_e2e.py create mode 100644 tests/unit/graph/test_effective_root_seed.py create mode 100644 tests/unit/graph/test_executor_runs_weka.py create mode 100644 tests/unit/graph/test_executor_watchdog.py create mode 100644 tests/unit/graph/test_first_token_anchor_lowering.py create mode 100644 tests/unit/graph/test_first_token_fidelity_tool.py create mode 100644 tests/unit/graph/test_first_token_routing.py create mode 100644 tests/unit/graph/test_first_token_runtime.py create mode 100644 tests/unit/graph/test_graph_endpoint_guard.py create mode 100644 tests/unit/graph/test_graph_instance_stop_paths.py create mode 100644 tests/unit/graph/test_graph_parse_context.py create mode 100644 tests/unit/graph/test_graph_path_catalog.py create mode 100644 tests/unit/graph/test_graph_return_bridge.py create mode 100644 tests/unit/graph/test_graph_sticky_stamp.py create mode 100644 tests/unit/graph/test_hf_sidecar_scope_out.py create mode 100644 tests/unit/graph/test_hf_streaming_trie_stores.py create mode 100644 tests/unit/graph/test_issuer_graph_path.py create mode 100644 tests/unit/graph/test_lane_fanout_recycle.py create mode 100644 tests/unit/graph/test_mixed_anchor_gate.py create mode 100644 tests/unit/graph/test_pressure_frontier_chop.py create mode 100644 tests/unit/graph/test_pressure_stage.py create mode 100644 tests/unit/graph/test_r6_wiring.py create mode 100644 tests/unit/graph/test_real_content_guard.py create mode 100644 tests/unit/graph/test_recorded_mode_stream_override.py create mode 100644 tests/unit/graph/test_registry_run_parse_parity.py create mode 100644 tests/unit/graph/test_runtime_hardening_round3.py create mode 100644 tests/unit/graph/test_slot_graph_eager_drain.py create mode 100644 tests/unit/graph/test_start_anchor_edges.py create mode 100644 tests/unit/graph/test_start_anchor_fidelity_tool.py create mode 100644 tests/unit/graph/test_start_anchor_runtime.py create mode 100644 tests/unit/graph/test_start_anchor_snapshot.py create mode 100644 tests/unit/graph/test_start_anchor_weka_stamping.py create mode 100644 tests/unit/graph/test_trie_prefix_reuse_oracle.py create mode 100644 tests/unit/graph/test_tstar_activation.py create mode 100644 tests/unit/graph/test_tstar_selection.py create mode 100644 tests/unit/graph/test_typed_channel_placeholders.py create mode 100644 tests/unit/graph/test_unified_interned_materialize.py create mode 100644 tests/unit/graph/test_warmup_boundary_rewrite.py create mode 100644 tests/unit/graph/test_warmup_handoff_consume.py create mode 100644 tests/unit/graph/test_warmup_routing_and_seed.py create mode 100644 tests/unit/graph/test_warmup_variants.py create mode 100644 tests/unit/graph/test_weka_block_size_decode.py create mode 100644 tests/unit/graph/test_weka_content_knob_wiring.py create mode 100644 tests/unit/graph/test_weka_fidelity_tool_gates.py create mode 100644 tests/unit/graph/test_weka_hf_real_ingest.py create mode 100644 tests/unit/graph/test_weka_ingest.py create mode 100644 tests/unit/graph/test_weka_isl_endpoint_osl_fidelity.py create mode 100644 tests/unit/graph/test_weka_memory_regressions.py create mode 100644 tests/unit/graph/test_weka_parallel_hardening.py create mode 100644 tests/unit/graph/test_weka_parse_kwargs_unification.py create mode 100644 tests/unit/graph/test_weka_segment_pool_merge_codec.py create mode 100644 tests/unit/graph/test_weka_segments.py create mode 100644 tests/unit/graph/test_weka_trace_models_finite.py create mode 100644 tests/unit/graph/test_weka_trie_build.py create mode 100644 tests/unit/graph/test_weka_trie_build_resolution.py create mode 100644 tests/unit/graph/test_weka_trie_hash_scope.py create mode 100644 tests/unit/graph/test_weka_trie_idle_gap_cap.py create mode 100644 tests/unit/graph/test_weka_trie_interval_order.py create mode 100644 tests/unit/graph/test_weka_trie_route.py create mode 100644 tests/unit/graph/test_weka_trie_snapshot.py create mode 100644 tests/unit/graph/test_weka_unified_dispatch.py create mode 100644 tests/unit/graph/test_worker_bytes_path.py create mode 100644 tests/unit/graph/test_worker_graph_branch.py create mode 100644 tests/unit/graph/test_worker_graph_capture.py create mode 100644 tests/unit/graph/test_worker_materialize.py create mode 100644 tests/unit/graph/test_worker_payload_features.py create mode 100644 tests/unit/graph/test_workload_detect_predicates.py create mode 100644 tests/unit/metrics/test_streamed_request_count.py create mode 100644 tests/unit/metrics/test_theoretical_prefix_cache_metric.py create mode 100644 tests/unit/plugin/test_graph_adapter_category.py create mode 100644 tests/unit/plugin/test_session_routing_registry.py create mode 100644 tests/unit/post_processors/test_streaming_gate.py create mode 100644 tests/unit/post_processors/test_theoretical_prefix_cache.py create mode 100644 tests/unit/scenario/test_context_overflow.py create mode 100644 tests/unit/scenario/test_env_locks.py create mode 100644 tests/unit/scenario/test_scenario_registry.py create mode 100644 tests/unit/scenario/test_scenario_validator.py create mode 100644 tests/unit/timing/strategies/test_graph_observer_teardown.py create mode 100644 tests/unit/timing/test_graph_phase_channel.py create mode 100644 tests/unit/timing/test_session_tree_finality.py create mode 100644 tests/unit/timing/test_session_tree_wiring.py create mode 100644 tests/unit/timing/test_timing_manager_first_token_advisory.py create mode 100644 tests/unit/timing/test_timing_manager_sidecar_load.py create mode 100644 tests/unit/transports/test_aiohttp_transport_bytes_passthrough.py create mode 100644 tests/unit/workers/test_session_routing.py create mode 100644 tests/unit/workers/test_streamed_stamp.py create mode 100644 tests/unit/workers/test_worker_graph_predispatch_errors.py create mode 100644 tests/unit/workers/test_worker_graph_store_discovery.py create mode 100644 tools/dag_jsonl_fidelity.py create mode 100644 tools/weka_trace_fidelity.py create mode 100755 tools/weka_trie_timing_sim.py diff --git a/.cursor/rules/python.mdc b/.cursor/rules/python.mdc index 5ac091b04a..09cfac2726 100644 --- a/.cursor/rules/python.mdc +++ b/.cursor/rules/python.mdc @@ -134,7 +134,7 @@ Feature branches use `/feature-name` format, forked from `main`. One P - Decorators: `@on_init`, `@on_start`, `@on_stop`, `@on_message`, `@on_command`, `@background_task`, `@on_pull_message`, `@on_request`. - Communication: `publish()` for broadcast, `@on_message` to subscribe, `send_command_and_wait_for_response()` for sync. - `AIPerfLifecycleMixin` for standalone components: `CREATED` -> `INITIALIZING` -> `INITIALIZED` -> `STARTING` -> `RUNNING` -> `STOPPING` -> `STOPPED`; `FAILED` terminal. -- `dag_jsonl` input type: conversation DAG benchmarks (fork + spawn modes). See `docs/benchmark-modes/dag.md` for abstractions and authoring. +- `dag_jsonl` input type: legacy conversation DAG benchmarks (fork + spawn modes). See `docs/benchmark-modes/dag.md`. For new agentic workloads use Graph IR — see `docs/benchmark-modes/agentic.md`. - Validator gate convention: unsupported constructs raise `NotImplementedError` with a leading `": "` prefix where `` identifies the conversation/turn (e.g. `"conversation 'foo' turn 3: ..."`). New validators must follow this shape. - Per-turn payload contract: `extra_body` / `max_tokens` / `model` are dispatch-turn only; `raw_tools` is the lone field that walks history (system-prompt-like). Dataset rows author `extra`, not `extra_body`. See `docs/dev/patterns.md` "Per-turn dataset `extra`". diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md index 2c46c26d2a..a53a3fb1a1 100644 --- a/.github/copilot-instructions.md +++ b/.github/copilot-instructions.md @@ -129,7 +129,7 @@ Feature branches use `/feature-name` format, forked from `main`. One P - Decorators: `@on_init`, `@on_start`, `@on_stop`, `@on_message`, `@on_command`, `@background_task`, `@on_pull_message`, `@on_request`. - Communication: `publish()` for broadcast, `@on_message` to subscribe, `send_command_and_wait_for_response()` for sync. - `AIPerfLifecycleMixin` for standalone components: `CREATED` -> `INITIALIZING` -> `INITIALIZED` -> `STARTING` -> `RUNNING` -> `STOPPING` -> `STOPPED`; `FAILED` terminal. -- `dag_jsonl` input type: conversation DAG benchmarks (fork + spawn modes). See `docs/benchmark-modes/dag.md` for abstractions and authoring. +- `dag_jsonl` input type: legacy conversation DAG benchmarks (fork + spawn modes). See `docs/benchmark-modes/dag.md`. For new agentic workloads use Graph IR — see `docs/benchmark-modes/agentic.md`. - Validator gate convention: unsupported constructs raise `NotImplementedError` with a leading `": "` prefix where `` identifies the conversation/turn (e.g. `"conversation 'foo' turn 3: ..."`). New validators must follow this shape. - Per-turn payload contract: `extra_body` / `max_tokens` / `model` are dispatch-turn only; `raw_tools` is the lone field that walks history (system-prompt-like). Dataset rows author `extra`, not `extra_body`. See `docs/dev/patterns.md` "Per-turn dataset `extra`". diff --git a/.github/workflows/run-unit-tests.yml b/.github/workflows/run-unit-tests.yml index 7b07af12be..66bc5656d4 100644 --- a/.github/workflows/run-unit-tests.yml +++ b/.github/workflows/run-unit-tests.yml @@ -115,7 +115,7 @@ jobs: - name: Run unit tests (Windows) if: runner.os == 'Windows' run: | - uv run pytest tests/unit -n auto --tb=short -m 'not performance and not stress and not slow' + uv run pytest tests/unit -n auto --timeout 600 --timeout-method thread --tb=short -m 'not performance and not stress and not slow' - name: Run zmq real-transport tests (Windows) # Real-socket libzmq tests (no looptime). ``!cancelled()`` lets them run # even if the unit step failed, matching the POSIX ``make test-ci`` @@ -131,7 +131,7 @@ jobs: # finishing after a manual stop. if: ${{ !cancelled() && runner.os == 'Windows' }} run: | - uv run pytest tests/component_integration -n auto --tb=short -v -m 'not performance and not stress and not slow' + uv run pytest tests/component_integration -n auto --timeout 600 --timeout-method thread --tb=short -v -m 'not performance and not stress and not slow' - name: Upload results to Codecov if: matrix.os == 'linux' && matrix.python-version == '3.12' uses: codecov/codecov-action@v5 diff --git a/AGENTS.md b/AGENTS.md index 3160ed5798..d2a6d2e3ba 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -128,7 +128,7 @@ Feature branches use `/feature-name` format, forked from `main`. One P - Decorators: `@on_init`, `@on_start`, `@on_stop`, `@on_message`, `@on_command`, `@background_task`, `@on_pull_message`, `@on_request`. - Communication: `publish()` for broadcast, `@on_message` to subscribe, `send_command_and_wait_for_response()` for sync. - `AIPerfLifecycleMixin` for standalone components: `CREATED` -> `INITIALIZING` -> `INITIALIZED` -> `STARTING` -> `RUNNING` -> `STOPPING` -> `STOPPED`; `FAILED` terminal. -- `dag_jsonl` input type: conversation DAG benchmarks (fork + spawn modes). See `docs/benchmark-modes/dag.md` for abstractions and authoring. +- `dag_jsonl` input type: legacy conversation DAG benchmarks (fork + spawn modes). See `docs/benchmark-modes/dag.md`. For new agentic workloads use Graph IR — see `docs/benchmark-modes/agentic.md`. - Validator gate convention: unsupported constructs raise `NotImplementedError` with a leading `": "` prefix where `` identifies the conversation/turn (e.g. `"conversation 'foo' turn 3: ..."`). New validators must follow this shape. - Per-turn payload contract: `extra_body` / `max_tokens` / `model` are dispatch-turn only; `raw_tools` is the lone field that walks history (system-prompt-like). Dataset rows author `extra`, not `extra_body`. See `docs/dev/patterns.md` "Per-turn dataset `extra`". diff --git a/CLAUDE.md b/CLAUDE.md index e4c300f7b0..47292fa0de 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -128,7 +128,7 @@ Feature branches use `/feature-name` format, forked from `main`. One P - Decorators: `@on_init`, `@on_start`, `@on_stop`, `@on_message`, `@on_command`, `@background_task`, `@on_pull_message`, `@on_request`. - Communication: `publish()` for broadcast, `@on_message` to subscribe, `send_command_and_wait_for_response()` for sync. - `AIPerfLifecycleMixin` for standalone components: `CREATED` -> `INITIALIZING` -> `INITIALIZED` -> `STARTING` -> `RUNNING` -> `STOPPING` -> `STOPPED`; `FAILED` terminal. -- `dag_jsonl` input type: conversation DAG benchmarks (fork + spawn modes). See `docs/benchmark-modes/dag.md` for abstractions and authoring. +- `dag_jsonl` input type: legacy conversation DAG benchmarks (fork + spawn modes). See `docs/benchmark-modes/dag.md`. For new agentic workloads use Graph IR — see `docs/benchmark-modes/agentic.md`. - Validator gate convention: unsupported constructs raise `NotImplementedError` with a leading `": "` prefix where `` identifies the conversation/turn (e.g. `"conversation 'foo' turn 3: ..."`). New validators must follow this shape. - Per-turn payload contract: `extra_body` / `max_tokens` / `model` are dispatch-turn only; `raw_tools` is the lone field that walks history (system-prompt-like). Dataset rows author `extra`, not `extra_body`. See `docs/dev/patterns.md` "Per-turn dataset `extra`". diff --git a/Makefile b/Makefile index bea9176480..51ecaa12a3 100644 --- a/Makefile +++ b/Makefile @@ -247,13 +247,13 @@ test-ci: #? run the tests using pytest-xdist for CI. @printf "$(bold)$(blue)Running unit and component integration tests (CI mode)...$(reset)\n" @# Run unit tests first with coverage @printf "$(bold)$(blue)Running unit tests...$(reset)\n" - @$(activate_venv) && pytest tests/unit -n auto --cov=src/aiperf --cov-branch --cov-report= -m 'not performance and not stress and not slow' --tb=short $(args) || exit_code=$$?; \ + @$(activate_venv) && pytest tests/unit -n auto --timeout 600 --timeout-method thread --cov=src/aiperf --cov-branch --cov-report= -m 'not performance and not stress and not slow' --tb=short $(args) || exit_code=$$?; \ # Run real-socket zmq transport tests (real time + real sockets, no looptime) regardless of unit result \ printf "$(bold)$(blue)Running zmq real-transport tests...$(reset)\n"; \ $(activate_venv) && pytest tests/zmq --cov=src/aiperf --cov-branch --cov-append --cov-report= -m 'not performance and not stress and not slow' --no-looptime --tb=short $(args) || exit_code=$$((exit_code + $$?)); \ # Run component integration tests with coverage append regardless of unit test result \ printf "$(bold)$(blue)Running component integration tests...$(reset)\n"; \ - $(activate_venv) && MALLOC_ARENA_MAX=2 pytest tests/component_integration -n auto --cov=src/aiperf --cov-branch --cov-append --cov-report=html --cov-report=xml --cov-report=term -m 'not performance and not stress and not slow' -v --tb=short $(args) || exit_code=$$((exit_code + $$?)); \ + $(activate_venv) && MALLOC_ARENA_MAX=2 pytest tests/component_integration -n auto --timeout 600 --timeout-method thread --cov=src/aiperf --cov-branch --cov-append --cov-report=html --cov-report=xml --cov-report=term -m 'not performance and not stress and not slow' -v --tb=short $(args) || exit_code=$$((exit_code + $$?)); \ if [[ $$exit_code -eq 0 ]]; then \ printf "$(bold)$(green)AIPerf unit and component integration tests (CI mode) passed!$(reset)\n"; \ else \ diff --git a/README.md b/README.md index ebc046324b..02f70870da 100644 --- a/README.md +++ b/README.md @@ -119,7 +119,7 @@ Log File: /home/user/Code/aiperf/artifacts/granite4:350m-openai-chat-concurrency - Scalable multiprocess architecture with 10 services communicating via ZMQ - 3 UI modes: `dashboard` (real-time TUI), `simple` (progress bars), `none` (headless) -- Multiple benchmarking modes: concurrency, request-rate, [request-rate with max concurrency](docs/tutorials/request-rate-concurrency.md), [trace replay](docs/benchmark-modes/trace-replay.md) +- Multiple benchmarking modes: concurrency, request-rate, [request-rate with max concurrency](docs/tutorials/request-rate-concurrency.md), [trace replay](docs/benchmark-modes/trace-replay.md), [agentic workloads](docs/benchmark-modes/agentic.md) - Extensible plugin system for endpoints, datasets, transports, and metrics - [Public dataset support](docs/benchmark-datasets.md) including ShareGPT and custom formats @@ -153,6 +153,7 @@ Log File: /home/user/Code/aiperf/artifacts/granite4:350m-openai-chat-concurrency ### Workloads and Data - [Trace Benchmarking](docs/benchmark-modes/trace-replay.md) - Deterministic workload replay +- [Agentic Workloads](docs/benchmark-modes/agentic.md) - Benchmark multi-step agent workflows: replay recorded agent traces (Dynamo, Weka) or author them in AIPerf's Graph IR - [Bailian Traces](docs/tutorials/bailian-trace.md) - Bailian production trace replay - [Baseten Traces](docs/tutorials/baseten-trace.md) - Baseten Parquet production trace replay - [BurstGPT Traces](docs/tutorials/burst-gpt-trace.md) - BurstGPT real-world bursty traffic trace replay diff --git a/docs/architecture.md b/docs/architecture.md index fa8917aba1..dc2a6deeca 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -54,6 +54,7 @@ The Dataset Manager handles all aspects of input data management during benchmar - Parsing and validating input data to ensure it matches the expected format - Writing dataset to memory-mapped files, enabling workers to access data directly without message passing - Supporting custom dataset types, such as MoonCake traces, for advanced benchmarking scenarios +- For graph IR workloads, broadcasting a graph-typed `DatasetConfiguredNotification` — `DatasetMetadata.graph` (the trace universe plus per-node prefix-cache map) and `GraphSegmentClientMetadata` (the unified segment store and structural sidecar locations) — instead of conversation entries and a conversation memory-mapped store; the store build itself is owned by `GraphStoreBuilder` (`src/aiperf/dataset/graph/store_build.py`), and the schedule plane and workers read those exact broadcast paths (nothing re-parses the workload) - Managing the lifecycle of datasets, including initialization, iteration, and cleanup ### Timing Manager @@ -61,7 +62,7 @@ The Dataset Manager handles all aspects of input data management during benchmar The Timing Manager controls and coordinates the timing of requests during benchmarking runs through a credit-based system. **Key Responsibilities:** -- Scheduling when each request should be sent based on the selected timing mode (fixed schedule, request-rate, or user-centric rate) +- Scheduling when each request should be sent based on the selected timing mode (fixed schedule, request-rate, user-centric rate, or graph IR for agentic workloads) - Managing precise timing to accurately reproduce real-world or synthetic load patterns - Supporting advanced timing scenarios, such as replaying traces with specific inter-arrival times or simulating bursty traffic - Ensuring that requests are dispatched to workers at the correct intervals for reliable measurement @@ -183,6 +184,7 @@ The Timing Manager uses a **credit-based flow control system** to control when r - **Fixed schedule mode**: Replays conversation traces at precise timestamps from dataset metadata - **Request-rate mode**: Issues credits at a specific rate with configurable arrival patterns (constant, Poisson, gamma, concurrency burst) - **User-centric rate mode**: Each session acts as a separate user with calculated gaps between turns + - **Graph IR mode**: Runs a dataflow executor per agentic-workload trace, issuing a node's credit once its dependencies complete (selected automatically for graph workloads) **Flow Control Benefits:** - Prevents overwhelming the inference server diff --git a/docs/benchmark-datasets.md b/docs/benchmark-datasets.md index 36a48b7d9d..37aabf5989 100644 --- a/docs/benchmark-datasets.md +++ b/docs/benchmark-datasets.md @@ -72,7 +72,9 @@ This document describes datasets that AIPerf can use to generate stimulus. Addit ## Exgentic Agent Trace Replay -The Exgentic loaders stream recorded agent sessions directly from Hugging Face. `exgentic` is pinned to v1 revision `70036b93a04e61b0ea2706a68b962f4f26774587`; `exgentic_v2` is independently pinned to v2 revision `4b8ad4ab198438e5a170f9171c19c6a2cf7c1814`. Each replays successful, positive-token chat call snapshots. Recorded messages, system instructions, tool definitions, output-token limits, request controls, and call start times are preserved. Tools are not executed, and live responses are not added to later requests. Every request carries the source session as `x-dynamo-session-id` for Dynamo agentic tracing while AIPerf retains its own request correlation ID. +The Exgentic loaders stream recorded agent sessions directly from Hugging Face. `exgentic` is pinned to v1 revision `70036b93a04e61b0ea2706a68b962f4f26774587`; `exgentic_v2` is independently pinned to v2 revision `4b8ad4ab198438e5a170f9171c19c6a2cf7c1814`. Each replays successful, positive-token chat call snapshots. Recorded messages, system instructions, tool definitions, output-token limits, request controls, and call start times are preserved. Tools are not executed, and live responses are not added to later requests. The loaders no longer stamp a session header themselves; AIPerf keeps a stable per-conversation correlation ID, and pairing the run with `--session-routing dynamo_headers` sends `X-Dynamo-Session-ID` for Dynamo session affinity and agentic tracing. + +This changes fixed-schedule affinity grouping. `--session-routing dynamo_headers` keys affinity on the live per-conversation correlation ID, and in `--fixed-schedule` replay the loader explodes one recorded agent session into many single-turn span-conversations — so each span now gets its own key and routes independently, and Dynamo prefix-affinity no longer groups a recorded session's spans onto one worker. The removed stamp instead put the shared recorded session ID on every span, grouping them (but that recorded ID also collided across dataset-recycled replays, binding unrelated live conversations to a single router entry). The two behaviors serve different fidelity goals; restoring source-session grouping without the recycling collision would need a dataset-provided affinity key (possible future work). Provide a finite materialization bound through `--num-conversations`, `--num-dataset-entries`, or `--request-count`. `--benchmark-duration` limits request issuance, not dataset setup. diff --git a/docs/benchmark-modes/agentic.md b/docs/benchmark-modes/agentic.md new file mode 100644 index 0000000000..5dbf4252bb --- /dev/null +++ b/docs/benchmark-modes/agentic.md @@ -0,0 +1,435 @@ + + +# Agentic Workloads: Native Graph IR + +Agentic workloads let AIPerf benchmark the request patterns real agents produce — multi-step LLM workflows with dependencies between calls, fan-out/fan-in, and state threaded from one response into the next prompt — instead of a flat list of independent requests or a linear conversation. Replay recorded agent traces (Dynamo, Weka) or author workloads directly; either way the workload is expressed in AIPerf's **Graph IR** and executed on the graph runtime. + +This page is the user guide for the **native Graph IR** file format: hand-authored `.yaml`, `.yml`, or `.jsonl` files that describe graph topology plus one or more traces. Imported Weka and Dynamo agent traces use the same agentic-workload lane and the same unified segment store. Prompts are limited to plain-string chat messages; unsupported constructs fail at parse time with an error naming the offending location (dispatch mechanics are covered under [Runtime behavior](#runtime-behavior)). + +## Graph IR vs. `dag_jsonl` + +Graph IR is AIPerf's native agentic workload representation. The earlier `dag_jsonl` conversation mode ([DAG Benchmarks](./dag.md)) predates it: built on the multi-turn session system, its context inheritance is tree-shaped — a session cannot have two FORK parents (no diamonds); SPAWN join gates provide control-flow fan-in. It is **legacy** — kept for its reactive fork/spawn orchestration, `BranchStats` export, and the legacy DAG child stop-condition rules, which the graph runtime does not replicate — and the graph plane replays the same `dag_jsonl` files with payload-identical profiling wire bodies (canonical order-insensitive comparison; [run-level accounting differs](./dag.md#known-documented-divergences)). New agentic workloads should target Graph IR. + +| Mode | Select with | Status | File shape | +|---|---|---|---| +| **Native Graph IR** | `--graph-format native` | **Native agentic mode.** Explicit dataflow graphs: nodes read and write channels, traces provide initial channel values, and LLM nodes build prompts from graph state. Fan-out and fan-in, including context/dataflow fan-in. | Native graph YAML or graph-record JSONL. | +| **DAG JSONL (graph plane)** | `--graph-format dag_jsonl` | **Migration path.** A DAG JSONL file replayed on the graph runtime (lanes, edge-delay timing, unified-segment KV-prefix dedup), with payload-identical profiling wire bodies to the legacy plane (canonical order-insensitive comparison; [run-level accounting differs](./dag.md#known-documented-divergences)). | Conversation JSONL; see [Running a `dag_jsonl` file on the graph runtime](./dag.md#running-a-dag_jsonl-file-on-the-graph-runtime). | +| **DAG JSONL (legacy plane)** | `--custom-dataset-type dag_jsonl` | **Legacy.** Branching chat conversations where each line is a session and `forks` / `spawns` connect sessions; tree-shaped context inheritance. Sole source of `BranchStats` export, the legacy child stop-condition rules, and reactive spawn orchestration. | Conversation JSONL; see [DAG Benchmarks](./dag.md). | + +Do not combine the two selectors. For a hand-authored native graph file, pass `--input-file ` and `--graph-format native`; to replay a DAG JSONL file on the graph plane, pass `--input-file ` and `--graph-format dag_jsonl`. In both cases do **not** also pass `--custom-dataset-type dag_jsonl` — that selects the separate legacy custom-dataset loader. + +AIPerf auto-detects the imported trace formats (`weka_trace`, `dynamo_trace`) through graph adapters, so `--graph-format` is optional for them. Native graph files and `dag_jsonl` graph-plane runs are intentionally **not** auto-detected: a plain `.yaml` or `.jsonl` might also be a normal custom dataset — and a `dag_jsonl` file is a valid legacy custom dataset — so AIPerf treats a file as a native graph or a graph-plane `dag_jsonl` workload only when you say so explicitly with `--graph-format`. + +## Replaying recorded agent traces (Dynamo, Weka) + +Point `--input-file` at the recorded capture. The `dynamo_trace` and `weka_trace` formats are auto-detected; pass `--graph-format` to force one explicitly: + +```bash +aiperf profile \ + --model Qwen3-0.6B \ + --endpoint-type chat \ + --url http://localhost:8000 \ + --input-file ./captures/trace.jsonl.gz \ + --graph-format dynamo_trace \ + --streaming +``` + +Accepted `--input-file` shapes per format: + +- `dynamo_trace` — a `.jsonl` or `.jsonl.gz` request-trace file (Dynamo's `jsonl` / `jsonl_gz` file sinks), a segmented capture (`trace.000000.jsonl.gz`, `trace.000001.jsonl.gz`, ...), or a directory containing those. +- `weka_trace` — a single `.json` trace file, a directory of `.json` trace files, or a Hugging Face corpus id (e.g. `org/weka-corpus`, loaded directly via `datasets`). + +AIPerf lowers each recorded trace into the unified segment store and replays it. Both formats keep the recorded inter-request delays, warped through the same per-trace idle-gap cap (60s default — see `--synthesis-idle-gap-cap`; set `synthesis.idle_gap_cap_seconds: null` in YAML to replay the raw recorded timeline), so an unbounded run spans the slowest trace's recorded duration — pass `--benchmark-duration ` to bound it. Both formats also pin each call's generation to the recorded output length (the node's `max_tokens`, mapped to the endpoint's token field; a recorded 0 upgrades to 1 with a warning), so replay never over-generates relative to the capture. [Dataset selection](#dataset-selection) documents the knobs that choose which and how many traces run (`--num-dataset-entries`, `--max-context-length`, `--allow-dataset-wrap`); the [warmup sections](#warmup-at-the-t-snapshot-window) cover recorded-replay warmup, including `--agentic-cache-warmup-duration`. + +## Authoring example + +Save this as `hello.graph.yaml`: + +```yaml +graph: + version: "2.0" + system: "You are a concise assistant." + +traces: + - id: hello-1 + messages: + - role: user + content: "Say hello in one short sentence." +``` + +Run it (native files require an explicit `--graph-format native`): + +```bash +aiperf profile \ + --model Qwen3-0.6B \ + --endpoint-type chat \ + --url http://localhost:8000 \ + --input-file hello.graph.yaml \ + --graph-format native +``` + +The YAML example uses the linear-chat shorthand: when the graph has no explicit `nodes`, the native parser derives one LLM node from the trace's `messages` and the graph-level `system` prompt. The lowering interns each trace's messages into the unified segment store at parse time (per-trace graphs, keyed by `trace.graph_ref`), so the shorthand dispatches end-to-end. + +## Minimal JSONL example + +Native JSONL uses one JSON object per line. Each object has a `kind` field. The `graph` record declares topology and must come before `trace` records. + +```jsonl +{"kind":"graph","version":"2.0","nodes":{"ask":{"node_type":"llm","prompt":[{"role":"user","content":["Question: ","@question"]}],"output":"answer"}}} +{"kind":"trace","id":"question-1","initial_state":{"question":"What is Graph IR in one sentence?"}} +``` + +In the JSONL example: + +- `ask` is an LLM node. +- `prompt` is a chat-message array. The string `@question` inside a message `content` list reads the `question` channel from the trace's `initial_state`. +- `output: answer` declares the channel that receives the LLM node result when the graph executes with a dispatch-capable payload path. +- `START -> ask -> END` edges are added automatically because the file has one root node and one leaf node. + +## Native Graph IR schema reference + +A native graph workload is parsed into these logical records: + +| Record | Required? | Purpose | +|---|---:|---| +| `graph` | optional but typical | Declares schema version, channels, nodes, edges, and provenance. If omitted, AIPerf uses an empty graph and derives the linear-chat shorthand from trace messages. | +| `trace` | yes for benchmark input | Supplies one runnable instance: `id`, optional tags, optional `graph_ref`, initial channel values, and replay outputs. | + +YAML can be written as one document with top-level sections: + +```yaml +graph: + version: "2.0" + state: {} + nodes: {} + edges: [] + +traces: [] +``` + +JSONL writes the same content as separate records: + +```jsonl +{"kind":"graph","version":"2.0","nodes":{},"edges":[]} +{"kind":"trace","id":"trace-1"} +``` + +### Graph fields + +Common `graph` fields: + +| Field | Type | Notes | +|---|---|---| +| `version` | string | Current native schema version is `"2.0"`. | +| `system` | string | Optional system prompt used by the linear-chat shorthand. | +| `state` | map | Channel declarations. Missing output channels and common prompt channels are inferred with safe defaults. | +| `nodes` | map | Node id to node spec. Node ids are referenced by `edges`. | +| `edges` | list | Static (unconditional) edges, optionally carrying scheduling-delay/anchor fields. If no explicit `START` or `END` edge is present, AIPerf injects edges from roots and to leaves. | + +### Node fields + +`node_type: llm` is the only node type: + +```yaml +nodes: + summarize: + node_type: llm + prompt: + - role: user + content: + - "Summarize this text: " + - "@document" + output: summary + streaming: true + max_tokens: 128 +``` + +Important LLM fields: + +| Field | Type | Notes | +|---|---|---| +| `prompt` | list | Prompt grammar that resolves to chat messages. | +| `output` | string | Channel that receives the model response. | +| `streaming` | bool | Whether this node should use streaming dispatch when the endpoint supports it. Defaults to `true`. | +| `model` | string | Model name dispatched for this call (same name and meaning as a conversation turn's `model`). Omit to use the run's `--model`. | +| `max_tokens` | int | Generation cap for this call (same name and meaning as a conversation turn's `max_tokens`). The worker maps it to the endpoint's token field (`max_completion_tokens`, or `max_tokens` under `--use-legacy-max-tokens`). Omit to leave generation uncapped. | +| `raw_tools` | list | OpenAI-compatible tool definitions for this call (same name and meaning as a conversation turn's `raw_tools`), sent as the body `tools` field. | +| `extra_headers` | map | Per-call HTTP headers (same name and meaning as a conversation turn's `extra_headers`), attached to the request headers, never the body. | +| `extra_body` | map | Per-call request-body fields such as `temperature`, `top_p`, or provider-specific keys (same name and meaning as a conversation turn's `extra_body`), passed through verbatim. Set the model, stream mode, token cap, and tools via the native fields above, not here. | + +`llm` is the only node type: every live workload — imported traces and hand-authored graphs alike — is a flat graph of LLM nodes wired with static edges. Any other `node_type`/`kind` value fails at parse as unknown. + +### Prompt grammar + +LLM `prompt` resolves to the chat messages sent to the endpoint. + +- A dict item is treated as a chat message and passed through after resolving any channel references in its `content` list. +- A top-level string `@messages_channel` splices a messages-typed channel into the prompt array. +- A string inside a message `content` list becomes a text block unless it starts with `@`, in which case it reads that channel and emits a content block of the channel's type. +- Use `@@literal` when you need a literal string that starts with `@`. + +Example with a messages splice plus a text channel: + +```yaml +graph: + nodes: + continue_chat: + node_type: llm + prompt: + - "@history" + - role: user + content: + - "Now answer this follow-up: " + - "@follow_up" + output: answer + +traces: + - id: chat-1 + initial_state: + history: + - role: user + content: "Explain cache locality." + - role: assistant + content: "Cache locality means reusing nearby data efficiently." + follow_up: "Give one LLM-serving example." +``` + +`history` is inferred as a messages channel because it is referenced at prompt-array level. `follow_up` is inferred as a text channel because it is referenced inside message content. + +### Static vs. dynamic content + +Whether a `@channel` reference is *static* (baked at build time) or *dynamic* (filled at run time from a predecessor's actual response) is inferred from who writes the channel: + +- **Static** — the channel is only seeded by `initial_state` (or self-written by the reading node, which observes pre-write state). Its content is known at build time and interned directly. The examples above are static. +- **Dynamic** — the channel is written by one or more upstream LLM nodes. The reference lowers to a *slot* filled at run time with those nodes' real responses, so a node's prompt can splice what the model actually said upstream. + +```yaml +graph: + nodes: + plan: + prompt: [{role: user, content: "Draft a plan."}] + output: plan_out + review: + prompt: + - role: user + content: ["Review this plan: ", "@plan_out"] # dynamic: plan's real reply + output: review_out + edges: + - {source: START, target: plan} + - {source: plan, target: review} + - {source: review, target: END} +traces: + - id: t1 +``` + +Dynamic composition rules and constraints: + +- **Array-level splices** (`"@history"`) on a messages channel reconstruct the full **user/assistant alternation**: `initial_state` messages first, then, for each upstream writer in completion order, that writer's authored user turn followed by its *actual* reply. Each user turn is authored once (in its own node's prompt) and each assistant turn is the live response — so the naive accumulate chain below yields a well-formed conversation with **no re-stating** of prior user turns. + + ```yaml + nodes: + turn1: {prompt: ["@hist", {role: user, content: "Name a color."}], output: hist} + turn2: {prompt: ["@hist", {role: user, content: "Name an animal."}], output: hist} + turn3: {prompt: ["@hist", {role: user, content: "Combine them."}], output: t3} + edges: + - {source: turn1, target: turn2} + - {source: turn2, target: turn3} + ``` + + `turn3`'s prompt materializes `[user "Name a color.", assistant , user "Name an animal.", assistant , user "Combine them."]`. +- **Chained writers**: when several nodes accumulate into one channel (`A` writes `hist`; `B` reads `@hist` and writes `hist`; `C` reads `@hist`), the writers must form an edge chain — completion ancestry along edges is what orders the writes deterministically; concurrent writers to a spliced channel are rejected at lowering. Every writer after the first must also itself read the channel — that read gates its dispatch behind the prior write and lets its authored turn be placed correctly in the reconstruction. If a channel has `initial_state` content, its first (root) writer must read `@hist` too (so it sees the seed the reconstruction attributes to it). +- **Block-level refs** (`"@plan_out"` inside a `content` list) compose the single writer's response into the containing message, so static instructions and the dynamic value share one message. +- Dynamic content requires per-trace **sticky routing** (automatic) so the node that produced a response and the node that splices it run on the same worker. It is not compatible with the t\* snapshot window (`--trajectory-start-max-ratio` must be `0`), which is off by default (full replay); a slot workload is rejected at load only when the window was explicitly engaged — via `--scenario inferencex-agentx-mvp` or explicit `--trajectory-start-min/max-ratio` flags. +- A producer whose request fails or returns no replayable content is *omitted* from the downstream splice (the assistant turn simply does not appear), matching how a real client would proceed past a failed turn. A tool-calls-only reply is not empty: its recorded assistant message (`tool_calls` included) splices verbatim. + +### Trace fields + +Common `trace` fields: + +| Field | Type | Notes | +|---|---|---| +| `id` | string | Required stable trace id. | +| `tags` | list | Optional labels round-tripped with the trace (provenance); not consumed at runtime. | +| `graph_ref` | string | Optional named-graph reference for multi-graph workloads; omit for the single top-level graph. | +| `messages` | list | Linear-chat shorthand. Lifted into the `messages` channel when the graph has no explicit nodes. | +| `initial_state` | map | Initial channel values available before the first node fires. | +| `replay_outputs` | map | Optional per-node recorded output channel values (`node_id -> {channel: value}`) for replay-style workloads. | + +There is no trace-level arrival-time field. To delay a node relative to trace start, set `min_start_delay_us` on the node (or on its `START` edge); inter-node pacing uses the edge delay fields described above. + +## Endpoint guidance + +Every graph credit dispatches against the run's global endpoint. Use the global endpoint flags: + +```bash +--endpoint-type chat --url http://localhost:8000 --model Qwen3-0.6B +``` + +Guidelines: + +- Put node-specific request-body knobs in `extra_body`; those are carried in the graph payload envelope and applied by the worker at materialization time. +- Use global `--url`, `--model`, `--header`, `--endpoint-type`, and `--custom-endpoint` for profile-time routing. + +There are no graph-level endpoint pools or per-node `endpoint` references: a graph record carrying `endpoints` — or a node carrying `endpoint` — fails at parse as an unknown field. Route every credit through the global `--url`/`--model` instead. + +## Session routing + +`--session-routing ` stamps live per-session identity on every graph request for external-router affinity (see the CLI reference for the four modes). Graph credits carry the full identity facts the routing plugins consume: + +- The **session key** is the trajectory's `x_correlation_id` (one per root chain or subagent chain per trace instance, fresh per recycle), and every request also carries the instance's **root trajectory corr** for tree-level grouping. +- `is_final_turn` is the trajectory's **recorded session-final fact** (the last recorded turn of that chain), so bind/close and session-final semantics track the recording. Tree-level finality stays conservative (`is_tree_final` is always `False` on the graph plane). +- With a routing mode active, the plugin **owns session identity**: recorded `x-dynamo-*` identity headers in Dynamo captures are stripped instead of being replayed (otherwise two conflicting identities would ride one request). Without a routing mode, recorded identity headers are replayed with per-instance uniquification, as before. +- Body-mutating modes (`dynamo_nvext`) cannot rewrite the pre-serialized verbatim node bytes of recorded-trace replay; the body transform is skipped with a one-time warning (header stamping still applies). Prefer header-based modes (`dynamo_headers`, `smg_routing_key`, `session_id_header`) for byte-exact graph replay. + +## Dataset selection + +Imported recorded-trace workloads (the Weka and Dynamo graph adapters) honor the standard dataset-selection knobs, so you choose *which* and *how many* traces run without editing the corpus. This section documents their **graph-plane semantics** — how the graph adapters interpret each knob. `--num-dataset-entries`, `--dataset-sampling-strategy`, and `--concurrency` are general-purpose flags that synthetic and public datasets also honor (with their usual meanings); the table below describes only their graph-adapter behavior. `--max-context-length` and `--allow-dataset-wrap` are graph-adapter-specific and are ignored by synthetic and public datasets. + +| Knob | Default | Effect on the graph plane | +|---|---|---| +| `--num-dataset-entries N` | unset (load all) | Caps the corpus to `N` distinct traces. **Unset loads every eligible trace** — there is no implicit default of 100 on the graph plane. | +| `--max-context-length T` | unset (no filter) | Drops any trace whose **peak context** — `input + output` tokens on its single largest request — exceeds `T`. Computed schema-only at parse time (no build, no tokenization). | +| `--allow-dataset-wrap` / `--no-allow-dataset-wrap` | derived | Whether selection may **wrap** (reuse the finite trace pool) to fill more concurrency/requests than there are distinct traces. Unset defers to a derived default computed at resolution time: **wrapping is enabled only when cache-bust is on** (`--cache-bust != none`), so plain runs default to no wrap. | +| `--dataset-sampling-strategy` | `sequential` | Order freed lanes draw templates in: `sequential` (in-order, byte-identical to the historical cursor draw), `shuffle` (per-pass seeded permutation, without replacement), or `random` (coerced to without-replacement shuffle here, so each corpus pass covers every trace once). | +| `--concurrency` | `1` | Number of trace instances replayed at once (the regular-aiperf default). | +| `--concurrency-ramp-duration` | unset | Ramps **lane admission** 1 → `--concurrency` over the duration: parked lanes start dispatching as the limit rises, spreading trace starts onto a cold server. (Graph credits bypass session slots, so this flag drives the lane limit directly on the graph plane.) | + +Selection is **filter-then-cap**: `--max-context-length` rejects oversized traces *first*, then `--num-dataset-entries` keeps the first `N` of the *eligible* survivors (in the adapter's deterministic scan order — directory files by name, Dynamo trees by root session id). The cap is never applied to the raw prefix, so a rejected trace early in the corpus never eats into the `N` kept. A once-per-build summary logs `scanned`, `rejected_by_maxctx`, `eligible`, and `loaded` counts. + +```mermaid +flowchart LR + C[Corpus traces
deterministic scan order] --> F{peak context
> --max-context-length?} + F -- yes --> R[reject] + F -- no --> K[keep] + K --> N{kept == --num-dataset-entries?} + N -- yes --> STOP[stop scanning] + N -- no --> C + STOP --> D[distinct loaded traces] + K --> D +``` + +### Wrap-guard: over-subscription fails instead of silently cloning + +When the resolved `--concurrency` exceeds the number of **distinct loaded traces** and wrapping is not allowed, AIPerf raises a `ConfigurationError` rather than silently cloning traces to fill the extra lanes (the previous behavior, ai-dynamo/aiperf#1106). This is the common trap after a `--max-context-length` filter shrinks the eligible pool well below the requested concurrency. + +To resolve it, pick one: + +- **Lower `--concurrency`** to at most the distinct loaded count. +- **Pass `--allow-dataset-wrap`** to intentionally reuse the finite pool across lanes/recycles. +- **Enable cache-bust** (e.g. `--cache-bust first_turn_prefix`), which both turns wrapping on by default and gives every cloned instance a distinct prefix marker. + +The default `--concurrency 1` never over-subscribes a non-empty corpus, so a plain run never trips the guard. + +### Duplication report + +Whenever lanes recycle the finite trace pool — to sustain concurrency, satisfy `--request-count`, or satisfy `--num-conversations` — the **dispatch duplication factor** is `total instances started / distinct loaded traces`. A factor above `1` means the same recorded traces were replayed more than once. This is a report, not a failure. AIPerf emits a **WARNING** only when the duplication has no cache-bust antidote (`--cache-bust` off / `none`): identical first-turn prefixes across clones collide in the server's KV cache and inflate prefix-cache-hit metrics. With cache-bust on, every instance mints a distinct marker, so the duplication is safe and the report stays quiet. Warmup phases are exempt (their priming is meant to warm the cache). + +## Runtime behavior + +At runtime, AIPerf executes each trace as an async dataflow graph: + +1. Seed the trace's `initial_state`. +2. Schedule nodes reachable from `START` whose inputs are satisfied. +3. For graph replay LLM nodes, map the fired runtime node to a build-time `node_ordinal`. +4. Issue a graph credit through the normal credit router; the worker materializes the request body from the unified segment store (`GraphSegmentUnifiedClient`). +5. Resolve the parked graph dispatch future when the graph return observer receives the worker return. +6. Publish node outputs, then schedule static successors as predecessor nodes finish. +7. Finish the trace when all reachable paths have reached `END` or have no more runnable successors. + +Graph LLM credits are materialized on workers from the unified segment store keyed by `(trace_id, node_ordinal, phase_variant)`. Native files ride the same store: parsing lowers every LLM node's prompt into content-addressed segments (the same content plane the Weka/Dynamo adapters emit), so hand-authored graphs dispatch end-to-end through `aiperf profile`. The lowering constrains what native prompts may contain: message content must be plain strings (or lists of string blocks, concatenated in order) with `role`/`content` keys only. An `@channel` splice backed by trace `initial_state` is interned as static content; an `@channel` splice that reads a channel written by an ancestor LLM node lowers to a dynamic slot filled at run time from that node's real response (see [Static vs. dynamic content](#static-vs-dynamic-content)). Every graph is a flat set of LLM nodes wired with static edges. Unsupported constructs fail at parse time with a `NotImplementedError` naming the offending location. + +Concurrency comes from two places: AIPerf can run multiple trace instances at once, and a single trace can have multiple ready graph nodes at once. Size graph lane concurrency with both levels in mind when your graph fans out. + +A bare graph-workload run — `aiperf profile --input-file workload.yaml --graph-format native`, or an auto-detected Weka/Dynamo capture via `--input-file` alone — with none of `--request-count` / `--num-conversations` / `--benchmark-duration` does a **single pass over the loaded corpus**: AIPerf loads all eligible traces (`--num-dataset-entries` unset = all) and runs **each trace exactly once**, bounded by the loaded **session count** (the distinct loaded trace count) — exactly the way `dag_jsonl` bounds a bare run by its root-session count. There is no auto-`--request-count 10` truncation for agentic workloads. Two seams cooperate: + +- **Config time** — the corpus size is not yet known (the weka HuggingFace corpus is streamed, and the `--max-context-length` filter runs at parse time), so the CLI-to-config converter (`_converter_profiling.py`) detects the graph workload (an explicit `--graph-format`, or an `--input-file` a graph adapter recognizes) and skips the auto-10, leaving the profiling phase unbounded. The phase then **validates** because the phase×dataset rule (`check_phase_dataset_compatibility`) exempts a no-stop phase whose dataset is a graph workload — its stop is inferred from the loaded corpus, the same way `--fixed-schedule` infers its stop from the trace. A no-stop concurrency phase against a **non-graph** dataset is still rejected. +- **Runtime** — the loaded corpus IS known when `GraphIRReplayStrategy` is built, so it derives an **explicit** session target `expected_num_sessions = len(traces)` (`_resolved_num_sessions`). The bare run therefore takes the same **bounded recycle path** as `--num-conversations`: freed lanes draw sequentially over the corpus and recycle is capped at `N`, so every distinct trace runs exactly once and then the gate closes — giving progress reporting a concrete `N`-session target instead of an implicit lane-clamp. + +Non-graph concurrency runs still get the plain-aiperf `--request-count 10` default and still require a stop condition. + +Set an explicit stop condition and it **overrides** the derived session target: the phase **recycles** fresh trace instances — freed lanes draw round-robin over the corpus, each instance cache-bust-marked — until the stop-condition gate closes, so `--concurrency` is sustained even beyond the corpus size. `--request-count N` caps total LLM-node dispatches (not trace instances), `--num-conversations N` caps distinct root sessions, and `--benchmark-duration D` caps wall time (it cancels still-parked idle nodes and keeps the records dispatched so far). + +Because a bare run is a single pass, `--concurrency` **cannot exceed the distinct loaded traces** when wrapping is disallowed (`--cache-bust` off and `--allow-dataset-wrap` unset): there are too few distinct traces to fill the lanes without cloning, so the setup-phase wrap-guard **fails loudly** with a `ConfigurationError` rather than silently cloning to fill (the ai-dynamo/aiperf #1106 contract). Pass `--cache-bust first_turn_prefix` (or `--allow-dataset-wrap`) to intentionally recycle the corpus, or reduce `--concurrency`. + +Pacing is owned by the recorded replay: node dispatch timing comes from the recorded delays and dataflow readiness, and `--concurrency` bounds how many trace instances replay at once. Flags that select a different pacing model — `--request-rate` (with any `--arrival-pattern`), `--user-centric-rate`, `--fixed-schedule`, `--adaptive-scale`, and their warmup variants — are **rejected up front** for agentic workloads rather than silently ignored. To run a graph-detected file through the linear pipeline instead, pin a loader with `--custom-dataset-type`. + +### Warmup at the t\* snapshot window + +Imported recorded-trace replays (trie graphs) run the **full trace by default** — the t\* window is off (`--trajectory-start-min-ratio`/`--trajectory-start-max-ratio` unset = `0.0`). Under [`--scenario inferencex-agentx-mvp`](../cli-options.md#scenario) (a named preset that locks benchmark invariants; or explicit `--trajectory-start-min/max-ratio` flags) each trace instead samples a per-trace snapshot instant `t*` inside `[min_ratio, max_ratio] × trace_duration` (the scenario applies `0.0..1.0`; it also pins `--synthesis-idle-gap-cap` to `10.0`). When the window is active (`--trajectory-start-max-ratio > 0`) and no explicit warmup phase is configured, AIPerf injects an automatic WARMUP phase ahead of PROFILING: + +- Warmup dispatches exactly **one priming credit per chain live at `t*`** — that chain's *boundary turn*, the last node of the per-session chain (root chain or subagent chain) recorded before `t*`. Trie prompts are cumulative, so priming the boundary turn's prompt (output capped by `AIPERF_GRAPH_WARMUP_MAX_OUTPUT_TOKENS`, default `1`) warms the chain's whole prefix in the server KV cache. +- A chain with no pre-`t*` node needs no priming; a chain entirely before `t*` is not live and is skipped. +- Priming credits burst at phase start: leading recorded offsets are dropped and recorded inter-turn gaps are never replayed during warmup. +- With `t* = 0` (window `[0, 0]`, or a zero-duration trace) the warmup graph is empty and the phase finalizes immediately. + +Profiling then replays only the at/after-`t*` portion of each trace at the full recorded output lengths, measuring against the warmed prefix. On multi-`--url` runs each trace instance keeps deterministic URL affinity across the warmup and profiling phases, so the priming and the measured replay hit the same backend. + +For a lower-level architecture walkthrough, see [Graph Async Dataflow Runtime Architecture](../reference/graph-async-dataflow-runtime.md). + +### Extended warmup (cache pressure + handoff) + +`--agentic-cache-warmup-duration ` extends the boundary-priming warmup with a cache-pressure stage. After every priming credit returns, the warmup phase continues replaying each lane's post-`t*` remainder with zero idle delay (all recorded inter-turn gaps collapsed) and 1-token outputs, recycling fresh templates onto freed lanes, for the configured duration — driving the server KV cache to steady-state pressure before any profiled request. + +When the duration elapses, in-flight requests drain and profiling resumes each lane at its **execution frontier** instead of re-firing from `t*`: nodes already executed during warmup/pressure are dropped from the profiling graph (the server holds their KV; the trie envelope keeps the full prompt prefix, so resume prompts are exact), and each chain's first pending node fires after its **residual delay** — the recorded gap to that turn minus the wall time already spent draining — so the phase boundary ramps instead of bursting. + +Notes: + +- Only the weka/dynamo graph-IR replay path honors this flag; it also implies a WARMUP phase even when the `t*` window is inactive (`t* = 0` runs pressure the full corpus compressed). +- Each resumed frontier's residual delay is clamped by `AIPERF_GRAPH_HANDOFF_RESIDUAL_CAP` (default 60s, matching the recorded idle-gap cap); `--burst-phase-starts` collapses the resumed leading offsets as usual (deliberately asymmetric: only true leading offsets are zeroed — a folded AND-join residual is mid-graph pacing and survives the burst collapse). +- Every pressure lane is honored at the profiling handoff: lanes still live at drain resume their execution frontier, while lanes that completed during pressure fresh-start on the next cursor template (a full `t*=0` replay in a dedicated `.f0` id namespace, not a re-run of the `t*` resume the pressure stage already executed; bounded runs only — a single-pass run keeps its cover-the-corpus-once pass and re-serves its pass-0 plan instead), and `--num-conversations` gates only recycles — never the drained lanes themselves — matching agentx's `_build_handoff_trajectories`. +- Warmup records (priming and pressure alike) are excluded from metrics as usual; only the resumed profiling turns are measured. +- Any **terminal request failure during warmup** — boundary priming or cache pressure — aborts the run before profiling (agentx parity): a warmup that could not prime the cache faithfully leaves a degraded pool, and profiling it would produce numbers that look valid but are not. Self-inflicted drain cancellations at the pressure deadline are *not* failures (the pressure stage cancels its own in-flight executors when the duration elapses) and never abort the run. +- The warmup phase is **mode-owned** when this flag is set: any user-configured warmup phase (`--warmup-request-count` / `--warmup-duration` / `--num-warmup-sessions`) is superseded by the auto boundary-priming + pressure shape (a notice is logged), so the stage always gets its full duration and no count cap can starve it. An explicit `--warmup-grace-period` is the one setting carried through verbatim; otherwise the drain waits `min(pressure duration, AIPERF_GRAPH_PRESSURE_DRAIN_GRACE_CAP)` (default cap 300s) for in-flight returns after sending completes, so a wedged or lost return cannot hang the run. On grace expiry the in-flight requests are cancelled: cancelled turns are excluded from the handoff ledger (not executed — profiling refires them), so a drain whose cancellations land yields a valid handoff, while a drain that force-completes with credits still unreturned trips the stash completeness gate and skips the re-cut (profiling then starts from the plain t\* plans). + +## Validation checklist + +Graph validation helpers are available for native graph files. The current profile ingest path parses agentic workloads and may surface parser/model errors before execution, but it does not run the semantic validator automatically — validator-only rules fire when you call `validate()` yourself (the full rule set is documented in [Graph IR Validation Reference](../reference/graph-ir-validation.md)). There is no CLI lint command yet; to run the semantic validator on a native file: + +```bash +python -c " +from pathlib import Path +from aiperf.dataset.graph.parser import parse_native +from aiperf.dataset.graph.validator import validate +for issue in validate(parse_native(Path('hello.graph.yaml'))): + print(issue) +" +``` + +| Check | Enforced | What to fix | +|---|---|---| +| Unknown record / node kind | parse time | Use `kind: graph` and `kind: trace` records with `node_type: llm` nodes; any other record or node kind (including the former `mix`, `subgraph`, rich node types, and conditional edges) fails as unknown. | +| More than one graph record | parse time | Keep one top-level graph record per file. | +| Graph record after trace records | parse time | Put the graph record first in JSONL. | +| Cycles | validator only | Native Graph IR expects an acyclic execution graph. Pre-unroll any loops in the trace topology (cycles are a future graph feature). | +| Unreachable nodes | validator only | Add edges from `START` or predecessor nodes, or remove the unused node. | +| Duplicate channel writers | validator only | Overwrite-reducer channels must have exactly one writer node. | +| Timing anchor shape | validator only | An edge carries at most one of the completion/dispatch anchors, and first-token anchors require their dispatch fallback. | + +Constructs outside this list (e.g. channel splice wiring) are not statically validated; errors there surface at parse decode, at the native lowering's gates, or at runtime. + +Parse and validation errors are reported with the file location or graph location when available, so fix parse-time structural errors before debugging runtime behavior. + +## Choosing the right format + +Use **native Graph IR** when: + +- The workload is easier to describe as nodes, edges, and channel state than as chat sessions. +- You need per-node request overrides via `extra_body`. +- You want explicit fan-out/fan-in inside one trace. +- You are converting another graph-like trace format into AIPerf's canonical graph representation. + +Use the **legacy `dag_jsonl` plane** only when you depend on the `BranchStats` export, the legacy DAG child stop-condition rules, or the reactive fork/spawn orchestration documented in [DAG Benchmarks](./dag.md) — the graph runtime does not replicate these. Otherwise, run existing `dag_jsonl` files on the graph plane (`--graph-format dag_jsonl`); the profiling wire bodies are payload-identical (canonical order-insensitive comparison; [run-level accounting differs](./dag.md#known-documented-divergences)). + +Use regular `single_turn`, `multi_turn`, or raw-payload replay when the workload is linear and does not need graph scheduling. + +## Related references + +- [DAG Benchmarks](./dag.md) — the legacy branching conversation mode Graph IR supersedes. +- [Native Graph IR schema reference](#native-graph-ir-schema-reference) — record, graph, node, prompt, and trace fields on this page. +- [Validation checklist](#validation-checklist) — common authoring errors and fixes. +- [Runtime behavior](#runtime-behavior) — user-facing execution model on this page. +- [Graph Async Dataflow Runtime Architecture](../reference/graph-async-dataflow-runtime.md) — lower-level runtime reference. +- [`profile_export_aiperf.json` Schema](../reference/json-export-schema.md) — metrics export schema after a run. diff --git a/docs/benchmark-modes/dag.md b/docs/benchmark-modes/dag.md index 4c2c2752ab..e04e369a62 100644 --- a/docs/benchmark-modes/dag.md +++ b/docs/benchmark-modes/dag.md @@ -5,6 +5,8 @@ SPDX-License-Identifier: Apache-2.0 # DAG Benchmarks: Branching Conversations +> **Status: legacy mode.** DAG mode predates AIPerf's Graph IR. For new agentic workloads use [Agentic Workloads (Graph IR)](./agentic.md). Existing `dag_jsonl` files run on the graph runtime by replacing `--custom-dataset-type dag_jsonl` with `--graph-format dag_jsonl` — profiling wire bodies are payload-identical (canonical order-insensitive comparison; see [Running a `dag_jsonl` file on the graph runtime](#running-a-dag_jsonl-file-on-the-graph-runtime) and the [documented divergences](#known-documented-divergences)). DAG mode remains the only source of `BranchStats` export, the legacy child stop-condition rules, and reactive fork/spawn orchestration. + Most benchmark conversations are a straight line: turn 1, then turn 2, then turn 3. DAG mode lets a single turn branch into **multiple follow-up conversations that run in parallel**. Picture a planner turn whose answer is then picked up by two different specialist turns at the same time, each continuing on its own from there. This guide walks through the feature from zero: what it is, when to reach for it, and how to author a file. No prior AIPerf knowledge is assumed beyond the basics in the README. @@ -13,8 +15,8 @@ This guide walks through the feature from zero: what it is, when to reach for it Reach for DAG when your workload looks like one of these: -- **Prefix-cache or KV-aware routing tests.** You want several follow-up requests to share the same long preamble so the server's cache is exercised. DAG's **FORK** mode makes the children look like continuations of the parent and routes them all to the same worker. -- **Agentic sub-agent trees.** A parent turn completes, then independent sub-agents kick off. Each sub-agent should start fresh, not inherit the parent's history. DAG's **SPAWN** mode handles this. +- **Prefix-cache or KV-aware routing tests.** You want several follow-up requests to share the same long preamble so the server's cache is exercised. DAG's **FORK** mode makes the children look like continuations of the parent and carries parent metadata for locality-aware routing experiments. +- **Agentic sub-agent trees.** A parent turn completes, then independent sub-agents kick off. Each sub-agent should start fresh, not inherit the parent's history. DAG's **SPAWN** mode handles this — for new agentic workloads prefer [Agentic Workloads (Graph IR)](./agentic.md); reach for SPAWN when you need its reactive orchestration or `BranchStats`. If your workload is a plain sequence of turns with no branching, you do **not** need DAG — stick with `multi_turn` or `raw_payload`. @@ -24,14 +26,14 @@ DAG mode exposes one primitive with two flavors, selected by a shorthand key on | Mode | Shorthand on parent turn | What the child sees | Routing | Parent fate | |---|---|---|---|---| -| **FORK** | `"forks": [...]` | Inherits the parent's full conversation history, including the captured model response. | Pinned to the same worker as the parent (locality). | Bare-string entries terminate; `{"child": ..., "background": true}` keeps the parent running. | +| **FORK** | `"forks": [...]` | Inherits the parent's full conversation history, including the captured model response. | Carries parent metadata for locality-aware routing; exact worker pinning depends on router wiring. | Bare-string entries terminate; `{"child": ..., "background": true}` keeps the parent running. | | **SPAWN** | `"spawns": [...]` | Starts from an empty history. Only the child's own messages go on the wire. | Free to land on any worker. | Continues; suspends only at an explicit `join_at` (or the next-turn auto-join). | Both keys can appear on the same turn — the scheduler treats them independently, so one turn can both fork continuations and spawn fresh sub-agents. ## A minimal example, walked through -Below is the shipped example at `examples/dag_jsonl/example.dag.jsonl`. Each line is one conversation; the three conversations together describe one tree. +Below is a minimal DAG JSONL example. Each line is one conversation; the three conversations together describe one tree. ```jsonl {"session_id":"root","turns":[{"model":"Qwen3-0.6B","messages":[{"role":"system","content":"You are a careful assistant."},{"role":"user","content":"Please summarize the attached document."}],"max_tokens":128,"forks":["branch-a","branch-b"]}]} @@ -53,7 +55,7 @@ flowchart TD **Line 2 — `branch-a`.** Two turns. Because it was reached via `forks`, it starts with `root`'s full accumulated history plus the real model response already in place. Its own messages get appended onto that, then dispatched. -**Line 3 — `branch-b`.** Also two turns, also forked from `root`. Runs in parallel with `branch-a` — both are sticky-routed to the same worker as `root`, so the server sees matching prefixes across the two siblings. +**Line 3 — `branch-b`.** Also two turns, also forked from `root`. Runs in parallel with `branch-a`; FORK children carry parent metadata so locality-aware routing can preserve sibling prefix locality when the router wiring supports it. Run it against any OpenAI-compatible chat endpoint: @@ -63,12 +65,12 @@ aiperf profile \ --endpoint-type chat \ --streaming \ --url localhost:8000 \ - --input-file examples/dag_jsonl/example.dag.jsonl \ + --input-file branch-example.dag.jsonl \ --custom-dataset-type dag_jsonl \ --concurrency 1 ``` -The example file has exactly one root (`root`); `branch-a` and `branch-b` are FORK targets, not roots. The autodefault sets `--num-conversations` to the root count, so `--concurrency` may not exceed `1` here. To exercise concurrency, supply your own multi-root DAG file or pass `--num-conversations N` explicitly. (FORK fanout still produces multiple in-flight requests per session — see the "concurrency" reference section below.) +The example has exactly one root (`root`); `branch-a` and `branch-b` are FORK targets, not roots. The autodefault sets `--num-conversations` to the root count when no request, duration, or session count is supplied. Higher `--concurrency` values are accepted, but with one root there may be no additional root sessions to keep every slot busy. FORK fanout can still produce multiple in-flight requests from that one root — see the "concurrency" reference section below. That is enough to get started. The rest of this document is reference material you can skim on demand. @@ -134,7 +136,7 @@ Each turn is a flat object validated against a strict schema (`DagTurn` in `src/ **Native vs. extra.** The top-level whitelist matches AIPerf's native `Turn` concepts (`messages`, `model`, `max_tokens`, `tools`) — the same fields AIPerf already tracks per-turn for any dataset. Anything else — sampling knobs (`temperature`, `top_p`, `seed`, `stop`, `logprobs`), response shaping (`response_format`), vendor tunables (`ignore_eos`, `min_tokens`, `top_k`) — lives in `extra`. At dispatch time the `extra` keys are merged into the top level of the wire body, so name them exactly as the server expects. -**What gets sent on the wire.** Structural keys (`forks`, `spawns`, `delay`) are consumed by the scheduler; every native field and everything under `extra` is forwarded to the chat-completions request body. +**What gets sent on the wire.** Structural keys (`forks`, `spawns`, `delay`) are consumed by the scheduler; native fields and everything under `extra` are translated by the endpoint layer into the request body. For example, `max_tokens` may become either `max_tokens` or `max_completion_tokens` depending on endpoint configuration. **Message shape.** Each entry in `messages` is a free-form dict — the only structural requirement is a `role` key, matching `MooncakeTrace`. `content` may be a string, a list of OpenAI multimodal parts (e.g. `[{"type": "text", "text": "..."}, {"type": "image_url", "image_url": {"url": "..."}}]`), or omitted for assistant messages that are purely `tool_calls`. Paste whatever the server expects; AIPerf forwards it verbatim onto the wire. @@ -143,7 +145,7 @@ Each turn is a flat object validated against a strict schema (`DagTurn` in `src/ `forks: [session_id, ...]` desugars into FORK-mode branches. When the parent turn completes, each listed child session: - Inherits the parent's accumulated message history (including the captured real assistant response), merged under the system-prompt rule below. -- Sticky-routes to the parent's worker so the server sees sibling requests with a common prefix and can exercise its prefix cache. +- Carries parent correlation metadata for locality-aware routing. Whether this pins to the parent's worker depends on the active router implementation. Each listed `session_id` must be declared as its own top-level conversation in the same file. A conversation can be the FORK target of **at most one** parent (ambiguous seed messages otherwise). See [Join Semantics](#join-semantics) below for how a parent can gate a later turn on its FORK/SPAWN children completing. @@ -170,7 +172,7 @@ A `DagFork` entry with `background: true` is the inherit-context-AND-parent-cont | Property | bare `forks: ["c"]` | `forks: [{"child": "c", "background": true}]` | |---|---|---| | Child inherits parent context | yes | yes | -| Sticky-routing to parent worker | yes | yes | +| Parent metadata for locality-aware routing | yes | yes | | Parent's remaining turns | not allowed (must be terminal) | run normally | | Join semantics | n/a (parent terminates) | none (fire-and-forget) | | Allowed on non-final turns | no | yes | @@ -196,7 +198,7 @@ DAG-style conversations can declare that a turn dispatches only after children f For v1, the orchestrator honors these gate shapes: -- **FORK**: no gate; child inherits parent context and sticky-routes. +- **FORK**: no gate; child inherits parent context and carries parent correlation metadata. - **SPAWN, immediate join (legacy bare-string form)**: parent suspends on the turn immediately after the spawning turn (`join_at = spawn_turn + 1`). - **SPAWN, delayed join (`DagSpawn.join_at = K`)**: busy-parent semantics. The parent runs turns `[spawn_turn+1 .. K-1]` concurrently with the spawned children and only suspends when it is about to dispatch turn `K`. - **SPAWN, fan-in (multiple branches gating one turn)**: a single gated turn may carry SPAWN_JOIN prereqs referencing multiple branches (across one or more spawning turns); the orchestrator pre-seeds an `outstanding` set and only fires when every referenced branch drains. Multi-consumer is also supported — one branch_id may be gated by prereqs on more than one downstream turn. @@ -249,12 +251,9 @@ If you need each phase to wrap the previous response with a new "system-like" fr ## Reference: routing and `agent_depth` -Every AIPerf session has its own `x_correlation_id` that pins it to a specific worker via sticky routing. In a DAG, FORK children inherit their parent's routing key: the router keys on the root session's correlation id, not each child's. That means: - -- All siblings in a fork hit the **same worker** as the parent. -- Siblings send the same root prefix, so the worker (and its server) see a clean prefix-cache hit pattern across sibling pairs. +Every AIPerf session has its own `x_correlation_id`, and DAG credits also carry parent correlation metadata. FORK children inherit the parent's accumulated context, so they send the same root prefix. Locality-aware routing can use the parent metadata to keep sibling requests near the parent, but the current sticky router still routes by the credit's own correlation id unless additional DAG-aware wiring is active. -This is what makes FORK mode useful for exercising prefix-cache and KV-aware routing — without sticky routing across the fork, siblings would scatter across workers and the prefix-share benefit would be invisible on the server. +This is why FORK mode is useful for prefix-cache and KV-aware routing experiments: it creates sibling requests with shared prefixes and carries the metadata needed for locality-aware routing, while the exact worker pinning behavior must be verified for the router configuration under test. Every credit and request record is tagged with two DAG-aware fields: @@ -278,13 +277,7 @@ Children are dispatched reactively by `BranchOrchestrator` at credit-return time ### `--num-conversations` autodefault for `dag_jsonl` -When neither `--request-count` nor `--num-conversations` is supplied for a `dag_jsonl` run, AIPerf auto-defaults `--num-conversations` to the **root count** of the file (sessions not referenced by any other conversation's `forks` list) rather than auto-defaulting `--request-count`. Auto-defaulting `--request-count` for a forking dataset would silently truncate the DAG mid-tree because the cap counts fork-spawned children. The CLI logs: - -``` -No request count or conversation count provided for forking dataset; -defaulting --num-conversations to N (run each root in the file once). -Use --request-count for a literal wire-request cap instead. -``` +When neither `--request-count`, `--benchmark-duration`, nor `--num-conversations` is supplied for a `dag_jsonl` run, AIPerf auto-defaults `--num-conversations` to the **root count** of the file. Roots are sessions not referenced by any other conversation's `forks`, `spawns`, or `pre_session_spawns`. Auto-defaulting `--request-count` for a forking dataset would silently truncate the DAG mid-tree because the cap counts fork-spawned children. If you do want a wire-request cap, pass `--request-count` explicitly — but be aware of the cap-applies-to-children behavior described above. @@ -294,15 +287,15 @@ Using the example file above, here is what happens on the wire: 1. `root`'s turn 0 dispatches as-is (accumulator is empty, so walking `turn_list` yields just the authored system + user). 2. When its response arrives, the worker appends a captured `{role: assistant, content: }` Turn onto `root.turn_list`. -3. The orchestrator sees `forks=["branch-a","branch-b"]` and sticky-routes both children to `root`'s worker; at the worker, `UserSessionManager.create_and_store` seeds each child's `turn_list` from the parent session's accumulator. Both children's turn 0 then dispatch concurrently. +3. The orchestrator sees `forks=["branch-a","branch-b"]`, creates children with parent metadata, and seeds each child's `turn_list` from the parent session's accumulator. Both children's turn 0 then dispatch concurrently. 4. `branch-a`'s turn 0 has its authored `raw_messages` appended into the child's `turn_list`; the chat endpoint walks the list and concatenates every turn's messages, producing `[root sys, root user, root assistant_response, child user_a, child user_b]`. No system-prompt rewriting happens — accumulation is pure concatenation. 5. `branch-a`'s turn 1 follows the same rule, now on top of the captured response from turn 0. 6. `branch-b` runs concurrently with `branch-a`, independently. -7. `root` has no further turns, so it terminates at the fork point. Its session is pinned in the worker cache (declared DAG branches) so late-arriving siblings can still seed their `turn_list` from it. +7. `root` has no further turns, so it terminates at the fork point after its fork children have been launched. ## Reference: validation and error messages -The loader performs strict structural checks at load time. Every error message includes the offending `file:line`. +The loader performs strict structural checks at load time. JSON parse and Pydantic schema errors include the offending line; cross-record semantic errors identify the offending session, turn, branch, or cycle path. | Failure | Example message | |---|---| @@ -390,7 +383,7 @@ aiperf profile \ --concurrency 3 ``` -With neither `--num-conversations` nor `--request-count` supplied, AIPerf logs `defaulting --num-conversations to 3` (one per root). The wire sees exactly nine requests: three roots and six children. `BranchStats.children_spawned` will be `6`, `children_completed` will be `6`, and all other counters will be `0`. +With neither `--num-conversations`, `--request-count`, nor `--benchmark-duration` supplied, AIPerf sets `--num-conversations` to `3` from the root count. The DAG JSONL loader's default sampler is random, so this config samples three root sessions and dispatches each sampled root's children; it does not guarantee one pass over each distinct root unless the sampling strategy is changed accordingly. If each of the three roots is sampled once, the wire sees nine requests: three roots and six children. In that case, `BranchStats.children_spawned` is `6`, `children_completed` is `6`, and the other counters are `0`. ## When NOT to use DAG mode @@ -399,6 +392,169 @@ With neither `--num-conversations` nor `--request-count` supplied, AIPerf logs ` - **Synthetic prompt generation** — DAG mode takes authored turn objects as given (messages are appended to the accumulator as-is). There is no synthetic input generator in v1. - **Diamond topologies** — a session with two FORK parents is explicitly rejected. DAG mode ships tree topology only. +## Successor mode: agentic workloads (Graph IR) + +DAG JSONL is the earlier plane: branching chat conversations authored as +sessions connected by `forks` and `spawns`, built on the multi-turn session +system. Its context inheritance is tree-shaped — a session cannot have two +FORK parents (no diamonds); SPAWN join gates provide control-flow fan-in. +[Agentic Workload Benchmarks](./agentic.md) is the successor: explicit +dataflow nodes, channels, per-node request overrides, imported Weka/Dynamo +agent traces, and fan-out/fan-in including context fan-in. + +Prefer Graph IR for new workloads; reach for the legacy plane for the +`BranchStats` export, the legacy child stop-condition rules, and the reactive +fork/spawn orchestration the graph runtime does not replicate. + +### Running a `dag_jsonl` file on the graph runtime + +The **same** DAG JSONL file can run on either of two execution planes: + +- The **legacy DAG plane** (`--custom-dataset-type dag_jsonl`) — the sub-agent + orchestrator and credit plumbing this whole guide describes. +- The **graph plane** (`--graph-format dag_jsonl`) — the file is lowered onto + AIPerf's graph IR and replayed by the [graph async dataflow + runtime](./agentic.md#runtime-behavior). + +Both accept the identical JSONL file format and run the identical load-time +validation (the graph adapter reuses the legacy loader, so every error message +in the [validation table](#reference-validation-and-error-messages) above still +applies). You choose the plane with the selector flag; nothing in the file +changes. + +Selection is **explicit-only**. `dag_jsonl` is deliberately excluded from graph +workload auto-detection, so an existing `--custom-dataset-type dag_jsonl` run is +never silently re-routed onto the graph plane — you get the graph plane only +when you ask for it with `--graph-format dag_jsonl`. + +#### What the graph plane does with the file + +On the graph plane, the DAG lowers to a dataflow graph rather than to reactive +orchestrator branches: + +- `forks`, `spawns`, `join_at`, and `pre_session_spawns` become static dataflow + edges plus AND fan-in gates (a joining node only fires once every gated + producer has landed). +- Accumulated conversation history is interned **verbatim** into the unified + segment store, so shared prefixes dedup into a real KV-cache prefix; each + ancestor's **live** captured reply is spliced in worker-side at run time. +- Per-turn `model`, `max_tokens`, `tools`, and `extra` ride dispatch overrides + in the legacy wire key order; per-turn `delay` maps to edge delays (the root + turn-0 delay is ignored, exactly as on the legacy plane). + +#### When to prefer each plane + +| You want… | Use | +|---|---| +| The graph async dataflow runtime, replay lanes, corpus-pass semantics, edge-delay replay timing, and unified-segment-store KV-prefix dedup | **Graph plane** (`--graph-format dag_jsonl`) | +| `BranchStats` export in `profile_export_aiperf.json` and the legacy DAG child stop-condition rules (`--request-count` honored per child, `--num-conversations` bypassed for children) | **Legacy plane** (`--custom-dataset-type dag_jsonl`) | + +If you need the `BranchStats` export or the legacy child stop-condition rules, +stay on the legacy plane — it is the plane every other section of this document +describes. Otherwise prefer the graph plane (`--graph-format dag_jsonl`); the +profiling wire bodies are payload-identical (canonical order-insensitive comparison), with the run-level differences listed +in the [divergences section](#known-documented-divergences). + +#### Wire-parity guarantee + +For a given run configuration, the graph plane emits **payload-identical wire +request bodies** to the legacy plane: same keys, same values, compared as +canonical (sorted-keys) serializations — byte-strict on values and types, but +key ORDER may differ (the graph plane authors `model` / `stream` / the token +cap through native node fields, whose wire positions differ from the legacy +`format_payload` order). This is proven for both streaming and non-streaming +chat, by two complementary gates: the **non-streaming** case by the in-process +transport-capture gate (`test_dag_jsonl_byte_parity.py`), and the **streaming** +case by a live run against the in-repo mock server +(`test_dag_jsonl_fidelity_real.py`). The guarantee is a **wire-body property**: +it is about the JSON that goes on the wire per request, not about run-level +accounting (see the divergences below). + +To reproduce the proof on your own file, run both planes with +`--export-level raw` against the **same** endpoint, then diff the raw exports +with the bundled fidelity tool: + +```bash +# Same server for both runs: dag_jsonl fork/spawn children splice the parent's +# LIVE reply into their prompt, so both planes must see identical replies. +aiperf-mock-server --port 8000 --ttft 0 --itl 0 & + +aiperf profile --input-file conv.dag.jsonl --custom-dataset-type dag_jsonl \ + --url http://127.0.0.1:8000 --endpoint-type chat --model test-model \ + --streaming --export-level raw --artifact-dir ./legacy --random-seed 42 + +aiperf profile --input-file conv.dag.jsonl --graph-format dag_jsonl \ + --url http://127.0.0.1:8000 --endpoint-type chat --model test-model \ + --streaming --export-level raw --artifact-dir ./graph --random-seed 42 + +python tools/dag_jsonl_fidelity.py \ + ./legacy/**/profile_export_raw.jsonl \ + ./graph/**/profile_export_raw.jsonl +``` + +The tool matches requests by `(session_id, turn_index)` and asserts the exported +`payload.messages`, the full parsed payload, and the re-serialized full body are +all equal on every matching PROFILING request. A SPAWN template instantiated more +than once per tree contributes its `#n` instances to the same +`(session_id, turn_index)` key, so the tool compares each key as a **multiset**: +it sorts both planes' bodies for that key and checks them element-wise, and +additionally requires the per-key multiplicity to match (a template fired N times +on one plane and M on the other fails loudly). The in-repo proofs are +`tests/component_integration/graph/test_dag_jsonl_byte_parity.py` (in-process +golden gate) and `tests/integration/graph/test_dag_jsonl_fidelity_real.py` +(real `aiperf profile` subprocesses against a live mock server). + +#### Known, documented divergences + +Byte parity is a per-request wire-body property. These behaviors intentionally +differ between the two planes: + +1. **Stop conditions and branch accounting.** The graph plane runs on lanes and + trace accounting, so it exports **no** `BranchStats` and treats + `--request-count` / `--num-conversations` per graph-replay semantics rather + than the legacy DAG child rules described under [stop conditions for DAG + children](#reference-stop-conditions-for-dag-children). Byte parity is a + wire-body property only. +2. **Warmup phase.** The two planes warm differently — the legacy plane prepends + a system prefix to warmup requests, while the graph plane applies a 1-token + output cap. The parity statements above apply to the **PROFILING** phase. + +`--extra-inputs` is **no longer** a divergence. Per-turn `extra` beats +`--extra-inputs` on an overlapping key on both planes, matching the legacy +contract (the run-level value is inserted first — fixing the key's position in +the wire body — and the turn value then wins). The graph adapter owns the fold +at parse: it merges `endpoint.extra` into every node's `extra_body` and +stamps the node `endpoint_extra_applied`, so the worker skips its own +`endpoint.extra` re-merge. The proof is the `mixed-full-extra-inputs` +parametrization in +`tests/component_integration/graph/test_dag_jsonl_byte_parity.py`, which runs +both planes with an `--extra-inputs` key overlapping a turn `extra` key (plus a +fresh vendor key) under the full-body byte-equality comparison. + +`agent_depth` / `parent_correlation_id` record stamping is **no longer** a +divergence. The graph plane stamps both fields on every dag record with the +legacy semantics: roots carry `agent_depth=0` with no parent, +FORK/SPAWN child instances carry the owner's depth + 1, and pre-session SPAWN +children carry `agent_depth=1` with no parent — every turn of a child instance +shares the instance's identity, so root-vs-child post-hoc filtering works +identically on both planes. One shape note: the graph plane mints a correlation +id per **node**, so a child's `parent_correlation_id` is the spawning parent +*turn's* correlation id (the graph analogue of the legacy parent *session* id). +The proofs are `test_fork_children_records_carry_dag_identity` and +`test_prespawn_child_record_carries_depth_one_no_parent` in +`tests/component_integration/graph/test_dag_jsonl_e2e.py`. + +Responses carrying `tool_calls` are **no longer** a divergence. Both planes +capture the reply through the same `build_assistant_turn`, which reassembles the +full assistant message (`content` plus `tool_calls`); the legacy plane splices +that message into the FORK child's inherited seed verbatim and the graph plane +serializes the identical dict into its dynamic pool and splices it back, so a +mixed text+`tool_calls` reply and a `content: null` `tool_calls`-only reply both +round-trip into child seeds byte-identically. The proof is the `fork_toolcalls` +parametrization in +`tests/component_integration/graph/test_dag_jsonl_byte_parity.py`, which exercises +both variants under the same byte-exact multiset comparison. + ## Related docs - [Raw Payload Replay](../tutorials/raw-payload-replay.md) — the non-forking analogue. diff --git a/docs/benchmark-modes/timing-modes-reference.md b/docs/benchmark-modes/timing-modes-reference.md index 507c3e6e9e..ccaf1fc6ca 100644 --- a/docs/benchmark-modes/timing-modes-reference.md +++ b/docs/benchmark-modes/timing-modes-reference.md @@ -342,7 +342,7 @@ With `--num-users 15` and `--user-centric-rate 1.0`, each user has 15 seconds be |--------|------|---------|-------------| | `--session-turns-mean` | float | 1.0 | Mean turns per session (`--user-centric-rate` requires ≥ 2) | | `--session-turns-stddev` | float | 0.0 | Standard deviation of turns | -| `--dataset-sampling-strategy` | enum | shuffle | Dataset sampling: `sequential`, `shuffle` (not with `--fixed-schedule`) | +| `--dataset-sampling-strategy` | enum | dataset-type-dependent | Dataset sampling: `sequential`, `random`, `shuffle`. Default depends on dataset type (`sequential` for traces, `shuffle` for synthetic). Not compatible with `--fixed-schedule` | ### Multi-URL Load Balancing diff --git a/docs/cli-options.md b/docs/cli-options.md index 6b57d2d488..bc2c3b9e2f 100644 --- a/docs/cli-options.md +++ b/docs/cli-options.md @@ -32,11 +32,15 @@ Expand a sweep config and print the resulting variations. Validate an AIPerf config file. +### [`dynamo`](#aiperf-dynamo) + +Dynamo agent-trace tooling. Use `aiperf dynamo trace-report` to aggregate metrics from a captured trace. + ### [`profile`](#aiperf-profile) Run the Profile subcommand. -[Endpoint](#endpoint) • [Tokenizer](#tokenizer) • [Input](#input) • [Fixed Schedule](#fixed-schedule) • [Goodput](#goodput) • [Conversation Input](#conversation-input) • [Prompt](#prompt) • [Prefix Prompt](#prefix-prompt) • [Input Sequence Length (ISL)](#input-sequence-length-isl) • [Output Sequence Length (OSL)](#output-sequence-length-osl) • [Audio Input](#audio-input) • [Image Input](#image-input) • [Video Input](#video-input) • [Rankings](#rankings) • [Synthesis](#synthesis) • [Load Generator](#load-generator) • [Warmup](#warmup) • [User-Centric Rate](#user-centric-rate) • [Request Cancellation](#request-cancellation) • [Output](#output) • [HTTP Trace](#http-trace) • [Server Metrics](#server-metrics) • [Network Latency](#network-latency) • [GPU Telemetry](#gpu-telemetry) • [UI](#ui) • [Multi-Run](#multi-run) • [Accuracy](#accuracy) • [Service](#service) • [Workers](#workers) • [ZMQ Communication](#zmq-communication) +[Endpoint](#endpoint) • [Tokenizer](#tokenizer) • [Input](#input) • [Fixed Schedule](#fixed-schedule) • [Goodput](#goodput) • [Conversation Input](#conversation-input) • [Prompt](#prompt) • [Prefix Prompt](#prefix-prompt) • [Input Sequence Length (ISL)](#input-sequence-length-isl) • [Output Sequence Length (OSL)](#output-sequence-length-osl) • [Audio Input](#audio-input) • [Image Input](#image-input) • [Video Input](#video-input) • [Rankings](#rankings) • [Synthesis](#synthesis) • [Load Generator](#load-generator) • [Warmup](#warmup) • [User-Centric Rate](#user-centric-rate) • [Request Cancellation](#request-cancellation) • [Output](#output) • [HTTP Trace](#http-trace) • [Server Metrics](#server-metrics) • [Network Latency](#network-latency) • [GPU Telemetry](#gpu-telemetry) • [UI](#ui) • [Multi-Run](#multi-run) • [Accuracy](#accuracy) • [Service](#service) • [Workers](#workers) • [ZMQ Communication](#zmq-communication) • [Scenario](#scenario) ### [`plot`](#aiperf-plot) @@ -50,7 +54,7 @@ Explore AIPerf plugins: aiperf plugins [category] [type] Run an AIPerf service in a single process. -[Parameters](#parameters) • [Endpoint](#endpoint) • [Tokenizer](#tokenizer) • [Input](#input) • [Fixed Schedule](#fixed-schedule) • [Goodput](#goodput) • [Conversation Input](#conversation-input) • [Prompt](#prompt) • [Prefix Prompt](#prefix-prompt) • [Input Sequence Length (ISL)](#input-sequence-length-isl) • [Output Sequence Length (OSL)](#output-sequence-length-osl) • [Audio Input](#audio-input) • [Image Input](#image-input) • [Video Input](#video-input) • [Rankings](#rankings) • [Synthesis](#synthesis) • [Load Generator](#load-generator) • [Warmup](#warmup) • [User-Centric Rate](#user-centric-rate) • [Request Cancellation](#request-cancellation) • [Output](#output) • [HTTP Trace](#http-trace) • [Server Metrics](#server-metrics) • [Network Latency](#network-latency) • [GPU Telemetry](#gpu-telemetry) • [UI](#ui) • [Multi-Run](#multi-run) • [Accuracy](#accuracy) • [Service](#service) • [Workers](#workers) • [ZMQ Communication](#zmq-communication) +[Parameters](#parameters) • [Endpoint](#endpoint) • [Tokenizer](#tokenizer) • [Input](#input) • [Fixed Schedule](#fixed-schedule) • [Goodput](#goodput) • [Conversation Input](#conversation-input) • [Prompt](#prompt) • [Prefix Prompt](#prefix-prompt) • [Input Sequence Length (ISL)](#input-sequence-length-isl) • [Output Sequence Length (OSL)](#output-sequence-length-osl) • [Audio Input](#audio-input) • [Image Input](#image-input) • [Video Input](#video-input) • [Rankings](#rankings) • [Synthesis](#synthesis) • [Load Generator](#load-generator) • [Warmup](#warmup) • [User-Centric Rate](#user-centric-rate) • [Request Cancellation](#request-cancellation) • [Output](#output) • [HTTP Trace](#http-trace) • [Server Metrics](#server-metrics) • [Network Latency](#network-latency) • [GPU Telemetry](#gpu-telemetry) • [UI](#ui) • [Multi-Run](#multi-run) • [Accuracy](#accuracy) • [Service](#service) • [Workers](#workers) • [ZMQ Communication](#zmq-communication) • [Scenario](#scenario) ### [`speed-bench-report`](#aiperf-speed-bench-report) @@ -229,6 +233,12 @@ Path to an AIPerf YAML config to validate.
+## `aiperf dynamo` + +Dynamo agent-trace tooling. Use `aiperf dynamo trace-report` to aggregate metrics from a captured trace. + +
+ ## `aiperf profile` Run the Profile subcommand. @@ -347,6 +357,17 @@ Use the legacy 'max_tokens' field instead of 'max_completion_tokens' in request Use server-reported token counts from API usage fields instead of client-side tokenization. When enabled, tokenizers are still loaded (needed for dataset generation) but tokenizer.encode() is not called for computing metrics. Token count fields will be None if the server does not provide usage information. For OpenAI-compatible streaming endpoints (chat/completions), stream_options.include_usage is automatically configured when this flag is enabled.
_Flag (no value required)_ +#### `--cache-bust` `` + +Where to inject a per-trace-instance cache-bust marker on the graph-IR replay path. 'none' (default): recorded bytes are sent verbatim. 'first_turn_prefix': prepend a '[rid:<12hex>]' marker to the first user turn. The marker is shared across all turns of one trace instance (its own prefix stays cacheable) but differs between distinct instances and resets on recycle, so the inference server's KV cache sees a distinct prefix per trace instance, defeating cross-instance prefix-cache hits that would inflate cache-hit metrics. + +**Choices:** + +| | | | +|-------|:-------:|-------------| +| `none` | _default_ | Cache-bust disabled: sessions send their recorded bytes verbatim. | +| `first_turn_prefix` | | Prepend the marker to the first turn's prompt prefix (token 0), so the whole prompt's KV-cache prefix diverges per session. Later turns inherit the changed prefix naturally; only the first turn is independently busted. | + #### `--connection-reuse-strategy` `` Transport connection reuse strategy. 'pooled' (default): connections are pooled and reused across all requests. 'never': new connection for each request, closed after response. 'sticky-user-sessions': connection persists across turns of a multi-turn conversation, closed on final turn (enables sticky load balancing). @@ -379,6 +400,15 @@ Content type for request body serialization. By default, requests are sent as 'a HTTP header name used to carry the per-session affinity identifier. When set, replaces the default `X-Correlation-ID` header with the provided name (e.g., `--session-header X-Session-ID`). +#### `--session-routing` `` + +Session-aware routing mode: stamps per-session identity on every request for router affinity. Built-ins: dynamo_headers (X-Dynamo-Session-ID + parent header), dynamo_nvext (nvext.session_control bind/close request-body metadata), smg_routing_key (X-SMG-Routing-Key for the SGLang Model Gateway manual policy), session_id_header (custom additive header). Parameterize with --session-routing-opt. +
_Choices: [`dynamo_headers`, `dynamo_nvext`, `smg_routing_key`, `session_id_header`]_ + +#### `--session-routing-opt` `` + +Repeatable key=value option for the selected --session-routing mode (e.g. --session-routing-opt timeout_seconds=600), validated against the plugin's Options model. + ### Tokenizer #### `--tokenizer` `` @@ -432,11 +462,25 @@ Dataset-specific filter in key=value form. Repeat for multiple filters. Only sup Format specification for custom dataset provided via `--input-file`. Determines parsing logic and expected file structure. Options: `single_turn` (JSONL with single exchanges), `multi_turn` (JSONL with conversation history), `mooncake_trace`/`bailian_trace`/`baseten_trace` (timestamped trace files), `random_pool` (directory of reusable prompts; when using `random_pool`, `--conversation-num` defaults to 100 if not specified; batch sizes > 1 sample each modality independently from a flat pool and do not preserve per-entry associations - use `single_turn` if paired modalities must stay together). Requires `--input-file`. Mutually exclusive with `--public-dataset`.
_Choices: [`burst_gpt_trace`, `bailian_trace`, `baseten_trace`, `mooncake_trace`, `raw_payload`, `inputs_json`, `dag_jsonl`, `sagemaker_data_capture`, `multi_turn`, `random_pool`, `single_turn`, `speed_bench_qualitative`, `speed_bench_coding`, `speed_bench_humanities`, `speed_bench_math`, `speed_bench_multilingual`, `speed_bench_qa`, `speed_bench_rag`, `speed_bench_reasoning`, `speed_bench_roleplay`, `speed_bench_stem`, `speed_bench_summarization`, `speed_bench_writing`, `speed_bench_throughput_1k`, `speed_bench_throughput_2k`, `speed_bench_throughput_8k`, `speed_bench_throughput_16k`, `speed_bench_throughput_32k`, `speed_bench_throughput_1k_low_entropy`, `speed_bench_throughput_1k_mixed`, `speed_bench_throughput_1k_high_entropy`, `speed_bench_throughput_2k_low_entropy`, `speed_bench_throughput_2k_mixed`, `speed_bench_throughput_2k_high_entropy`, `speed_bench_throughput_8k_low_entropy`, `speed_bench_throughput_8k_mixed`, `speed_bench_throughput_8k_high_entropy`, `speed_bench_throughput_16k_low_entropy`, `speed_bench_throughput_16k_mixed`, `speed_bench_throughput_16k_high_entropy`, `speed_bench_throughput_32k_low_entropy`, `speed_bench_throughput_32k_mixed`, `speed_bench_throughput_32k_high_entropy`]_ +#### `--graph-format` `` + +Select the agentic-workload format for `--input-file`, overriding graph-adapter auto-detection. Registered names: `dynamo_trace`, `weka_trace`, `dag_jsonl`, `native`. Requires `--input-file`. Independent of `--custom-dataset-type` (that selects a custom dataset loader; this selects a graph adapter). Use `native` to load a hand-authored Graph IR workload file, or `dag_jsonl` to run a DAG conversation file on the graph runtime. +
_Choices: [`dynamo_trace`, `weka_trace`, `dag_jsonl`, `native`]_ + #### `--dataset-sampling-strategy` `` Strategy for selecting entries from dataset during benchmarking. `sequential`: Iterate through dataset in order, wrapping to start after end. `random`: Randomly sample with replacement (entries may repeat before all are used). `shuffle`: Shuffle dataset and iterate without replacement, re-shuffling after exhaustion. Default behavior depends on dataset type (e.g., `sequential` for traces, `shuffle` for synthetic).
_Choices: [`random`, `sequential`, `shuffle`]_ +#### `--max-context-length` `` + +Maximum per-trace context length (tokens) for graph-plane dataset selection. Graph traces whose input+output context would exceed this cap are excluded from the eligible set before selection. Applies to graph-adapter workloads (`--input-file` graph traces); ignored by synthetic/public datasets. Unset (the default) applies no context-length filter. +
_Constraints: ≥ 1_ + +#### `--allow-dataset-wrap`, `--no-allow-dataset-wrap` + +Allow graph-plane dataset selection to wrap (reuse the finite trace pool) when `--request-count` exceeds the number of eligible traces. `--allow-dataset-wrap` forces wrapping on; `--no-allow-dataset-wrap` forces it off. Unset (the default, None) defers to a derived default computed at resolution time and surfaced on `run.resolved.allow_dataset_wrap`. Applies to graph-adapter workloads; ignored by synthetic/public datasets. + #### `--random-seed` `` Random seed for deterministic data generation. When set, makes synthetic prompts, sampling, delays, and other random operations reproducible across runs. Essential for A/B testing and debugging. Uses system entropy if not specified. Initialized globally at config creation. @@ -491,7 +535,7 @@ The total number of unique conversations to generate. Each conversation represen #### `--num-dataset-entries`, `--num-prompts` `` -Total number of unique entries to generate for the dataset. Each entry represents one user message that can be used as a turn in conversations. Entries are reused across conversations and turns according to `--dataset-sampling-strategy`. Higher values provide more diversity. +Total number of unique entries to generate for the dataset. Each entry represents one user message that can be used as a turn in conversations. Entries are reused across conversations and turns according to `--dataset-sampling-strategy`. Higher values provide more diversity. On the graph plane (graph-adapter workloads), this bounds the distinct traces selected: unset means all eligible traces (after `--max-context-length` and other filters) are used; N means up to N distinct traces are selected after filters. It does not count requests -- see `--request-count`.
_Constraints: ≥ 1_
_Default: `100`_ @@ -900,6 +944,15 @@ Maximum input sequence length for filtering. Traces with input_length > max_isl Maximum output sequence length cap. Traces with output_length > max_osl are capped to max_osl.
_Constraints: ≥ 1_ +#### `--prompt-corpus` `` + +Corpus backing recorded graph (`weka_trace`, `dynamo_trace`) real-content synthesis. `coding` (default when unset) uses the procedural CodingContentGenerator pool the recorded weka workloads were captured against; `sonnet` uses the Shakespeare pool, which yields matching token counts but different bytes (useful only to reproduce golden fixtures built from that pool). Ignored by non-graph datasets. + +#### `--synthesis-idle-gap-cap` `` + +Per-trace idle-gap cap (seconds) for recorded graph (`weka_trace`, `dynamo_trace`) replay. Recorded inter-request idle gaps longer than this are compressed to the cap so a recorded multi-hour gap does not park the replay. Defaults to 60s when unset. Ignored by non-graph datasets. +
_Constraints: ≥ 0.0_ + ### Load Generator #### `--benchmark-duration` `` @@ -937,11 +990,11 @@ Smoothness parameter for gamma distribution arrivals (--arrival-pattern gamma). #### `--request-count`, `--num-requests` `` -The maximum number of requests to send. If not set, will be automatically determined based on the timing mode and dataset size. For synthetic datasets, this will be `max(10, concurrency * 2)`. Pass a comma-separated list (e.g. `--request-count 100,500,1000`) to sweep over multiple request counts; the converter promotes the list to a sweep on phases.profiling.requests before AIPerfConfig validation. +The maximum number of requests to send. If not set, will be automatically determined based on the timing mode and dataset size. For synthetic datasets, this will be `max(10, concurrency * 2)`. Pass a comma-separated list (e.g. `--request-count 100,500,1000`) to sweep over multiple request counts; the converter promotes the list to a sweep on phases.profiling.requests before AIPerfConfig validation. On the graph plane (graph-adapter workloads), this counts individual requests, not traces -- a single trace can emit many requests, so `--request-count` does not bound the number of distinct traces selected (use `--num-dataset-entries` for that). When it exceeds the eligible trace pool, `--allow-dataset-wrap` controls whether selection wraps to reuse traces. #### `--concurrency-ramp-duration` `` -Duration in seconds to ramp session concurrency from 1 to target. Useful for gradual warm-up of the target system. +Duration in seconds to ramp session concurrency from 1 to target. On graph-IR replay (weka/dynamo/dag traces) this ramps LANE admission: concurrent trace-instance lanes are admitted 1 -> --concurrency over the duration. Useful for gradual warm-up of the target system.
_Constraints: > 0_ #### `--prefill-concurrency-ramp-duration` `` @@ -991,6 +1044,16 @@ Maximum adaptive scale control value. Inferred from the phase target when omitte SLA filter for adaptive scale. Format: 'metric_tag:stat:op:threshold'. Latency-family metrics (request_latency, time_to_first_token/ttft, inter_token_latency/itl/tpot) support percentile stats; window scalar/rate metrics (request_throughput, output_token_throughput, goodput, goodput_ratio, success_rate, error_rate, cancellation_rate) support {avg, min, max}. Full metric/stat table: [Adaptive SLA metric support](tutorials/yaml-config.md#adaptive-sla-metric-support). op in {lt, le, gt, ge}; threshold is a float. Repeatable. Example: --adaptive-scale-sla 'request_latency:p95:le:30000'. +#### `--trajectory-start-min-ratio` `` + +Lower bound (fraction of each trace's recorded wall-clock duration) of the per-trace t* snapshot-instant sampling window for recorded graph replay. Unset resolves to 0.0. With both bounds at 0.0 the window is OFF (full native replay). Per trace, t* is drawn uniformly from [min, max] x trace_duration with a trace-salted RNG: nodes firing before t* become cache-priming WARMUP history, nodes at/after t* are PROFILING. Set min == max for a deterministic t*. Must be <= --trajectory-start-max-ratio. The inferencex-agentx-mvp scenario auto-applies 0.0 when unset and locks an explicit conflicting value. +
_Constraints: ≥ 0.0, ≤ 1.0_ + +#### `--trajectory-start-max-ratio` `` + +Upper bound (fraction of each trace's recorded wall-clock duration) of the per-trace t* snapshot-instant sampling window for recorded graph replay. Unset resolves to 0.0 (window OFF, full native replay); any value > 0 engages the t* snapshot machinery (per-trace t* sampling plus the auto boundary-priming warmup phase). See --trajectory-start-min-ratio for the sampling contract. The inferencex-agentx-mvp scenario auto-applies 1.0 when unset and locks an explicit conflicting value. +
_Constraints: ≥ 0.0, ≤ 1.0_ + ### Warmup #### `--warmup-request-count`, `--num-warmup-requests` `` @@ -1003,6 +1066,16 @@ The maximum number of warmup requests to send before benchmarking. If not set an The maximum duration in seconds for the warmup phase. If not set, it will use the `--warmup-request-count` value. If neither are set, no warmup phase will be used.
_Constraints: > 0_ +#### `--agentic-cache-warmup-duration` `` + +Extended (cache-pressure) warmup duration in seconds for the recorded-trace graph-IR replay path. After the boundary-priming warmup drains, AIPerf continues the live post-t* replay with zero idle delay and one-token outputs for this long, recycling fresh templates onto freed lanes, then drains and hands each lane's execution frontier to profiling (with residual delays), driving the server KV cache to steady-state pressure before benchmarking. If not set, no extended warmup runs. +
_Constraints: > 0_ + +#### `--burst-phase-starts` + +Burst-vs-spread control for phase starts on the recorded-trace graph-IR replay path. Default False = SPREAD: each WARMUP / PROFILING trace's first firing is offset by its recorded (idle-gap-warped, <=60s) lead from the t* snapshot instant. True = BURST: both phase starts collapse into a synchronized burst -- every warmup priming credit and every trace's earliest profiling firing fire at once at phase-time 0 (the leading per-firing dispatch offset is dropped); relative inter-turn delays after the first are still honored. Governs ONLY the two phase-start dispatch patterns. +
_Flag (no value required)_ + #### `--num-warmup-sessions` `` The number of sessions to use for the warmup phase. If not set, it will use the `--warmup-request-count` value. @@ -1588,6 +1661,17 @@ Directory path for ZMQ IPC (Inter-Process Communication) socket files. When usin Select the ZMQ dual-bind communication backend (IPC + TCP). All dual-bind knobs are cluster-managed; this flag only selects the discriminator and the converter routes downstream to the default.
_Flag (no value required)_ +### Scenario + +#### `--scenario` `` + +Lock all benchmark invariants for a named scenario (e.g. 'inferencex-agentx-mvp'). Conflicts with the locked invariants raise ScenarioLockError at startup unless --unsafe-override is also passed. Distinct from the sweep ``scenarios`` strategy (hand-picked named runs). + +#### `--unsafe-override` + +Convert scenario lock errors to warnings; stamps submission_valid=false in the aggregate output. No-op without --scenario. +
_Flag (no value required)_ +
## `aiperf plot` @@ -1675,7 +1759,7 @@ Explore AIPerf plugins: aiperf plugins [category] [type] #### `--category` `` Category to explore. -
_Choices: [`accumulator`, `accuracy_benchmark`, `accuracy_grader`, `analyzer`, `api_router`, `arrival_pattern`, `communication`, `communication_client`, `console_exporter`, `convergence_criterion`, `custom_dataset_loader`, `data_exporter`, `dataset_backing_store`, `dataset_client_store`, `dataset_composer`, `dataset_sampler`, `endpoint`, `gpu_telemetry_collector`, `plot`, `public_dataset_loader`, `ramp`, `record_observer`, `record_processor`, `search_planner`, `search_recipe`, `search_recipe_post_process`, `service`, `service_manager`, `stream_exporter`, `timing_strategy`, `transport`, `ui`, `url_selection_strategy`, `zmq_proxy`]_ +
_Choices: [`accumulator`, `accuracy_benchmark`, `accuracy_grader`, `analyzer`, `api_router`, `arrival_pattern`, `communication`, `communication_client`, `console_exporter`, `convergence_criterion`, `custom_dataset_loader`, `data_exporter`, `dataset_backing_store`, `dataset_client_store`, `dataset_composer`, `dataset_sampler`, `endpoint`, `gpu_telemetry_collector`, `graph_adapter`, `plot`, `public_dataset_loader`, `ramp`, `record_observer`, `record_processor`, `search_planner`, `search_recipe`, `search_recipe_post_process`, `service`, `service_manager`, `session_routing`, `stream_exporter`, `timing_strategy`, `transport`, `ui`, `url_selection_strategy`, `zmq_proxy`]_ #### `--name` `` @@ -1805,6 +1889,17 @@ Use the legacy 'max_tokens' field instead of 'max_completion_tokens' in request Use server-reported token counts from API usage fields instead of client-side tokenization. When enabled, tokenizers are still loaded (needed for dataset generation) but tokenizer.encode() is not called for computing metrics. Token count fields will be None if the server does not provide usage information. For OpenAI-compatible streaming endpoints (chat/completions), stream_options.include_usage is automatically configured when this flag is enabled.
_Flag (no value required)_ +#### `--cache-bust` `` + +Where to inject a per-trace-instance cache-bust marker on the graph-IR replay path. 'none' (default): recorded bytes are sent verbatim. 'first_turn_prefix': prepend a '[rid:<12hex>]' marker to the first user turn. The marker is shared across all turns of one trace instance (its own prefix stays cacheable) but differs between distinct instances and resets on recycle, so the inference server's KV cache sees a distinct prefix per trace instance, defeating cross-instance prefix-cache hits that would inflate cache-hit metrics. + +**Choices:** + +| | | | +|-------|:-------:|-------------| +| `none` | _default_ | Cache-bust disabled: sessions send their recorded bytes verbatim. | +| `first_turn_prefix` | | Prepend the marker to the first turn's prompt prefix (token 0), so the whole prompt's KV-cache prefix diverges per session. Later turns inherit the changed prefix naturally; only the first turn is independently busted. | + #### `--connection-reuse-strategy` `` Transport connection reuse strategy. 'pooled' (default): connections are pooled and reused across all requests. 'never': new connection for each request, closed after response. 'sticky-user-sessions': connection persists across turns of a multi-turn conversation, closed on final turn (enables sticky load balancing). @@ -1837,6 +1932,15 @@ Content type for request body serialization. By default, requests are sent as 'a HTTP header name used to carry the per-session affinity identifier. When set, replaces the default `X-Correlation-ID` header with the provided name (e.g., `--session-header X-Session-ID`). +#### `--session-routing` `` + +Session-aware routing mode: stamps per-session identity on every request for router affinity. Built-ins: dynamo_headers (X-Dynamo-Session-ID + parent header), dynamo_nvext (nvext.session_control bind/close request-body metadata), smg_routing_key (X-SMG-Routing-Key for the SGLang Model Gateway manual policy), session_id_header (custom additive header). Parameterize with --session-routing-opt. +
_Choices: [`dynamo_headers`, `dynamo_nvext`, `smg_routing_key`, `session_id_header`]_ + +#### `--session-routing-opt` `` + +Repeatable key=value option for the selected --session-routing mode (e.g. --session-routing-opt timeout_seconds=600), validated against the plugin's Options model. + ### Tokenizer #### `--tokenizer` `` @@ -1890,11 +1994,25 @@ Dataset-specific filter in key=value form. Repeat for multiple filters. Only sup Format specification for custom dataset provided via `--input-file`. Determines parsing logic and expected file structure. Options: `single_turn` (JSONL with single exchanges), `multi_turn` (JSONL with conversation history), `mooncake_trace`/`bailian_trace`/`baseten_trace` (timestamped trace files), `random_pool` (directory of reusable prompts; when using `random_pool`, `--conversation-num` defaults to 100 if not specified; batch sizes > 1 sample each modality independently from a flat pool and do not preserve per-entry associations - use `single_turn` if paired modalities must stay together). Requires `--input-file`. Mutually exclusive with `--public-dataset`.
_Choices: [`burst_gpt_trace`, `bailian_trace`, `baseten_trace`, `mooncake_trace`, `raw_payload`, `inputs_json`, `dag_jsonl`, `sagemaker_data_capture`, `multi_turn`, `random_pool`, `single_turn`, `speed_bench_qualitative`, `speed_bench_coding`, `speed_bench_humanities`, `speed_bench_math`, `speed_bench_multilingual`, `speed_bench_qa`, `speed_bench_rag`, `speed_bench_reasoning`, `speed_bench_roleplay`, `speed_bench_stem`, `speed_bench_summarization`, `speed_bench_writing`, `speed_bench_throughput_1k`, `speed_bench_throughput_2k`, `speed_bench_throughput_8k`, `speed_bench_throughput_16k`, `speed_bench_throughput_32k`, `speed_bench_throughput_1k_low_entropy`, `speed_bench_throughput_1k_mixed`, `speed_bench_throughput_1k_high_entropy`, `speed_bench_throughput_2k_low_entropy`, `speed_bench_throughput_2k_mixed`, `speed_bench_throughput_2k_high_entropy`, `speed_bench_throughput_8k_low_entropy`, `speed_bench_throughput_8k_mixed`, `speed_bench_throughput_8k_high_entropy`, `speed_bench_throughput_16k_low_entropy`, `speed_bench_throughput_16k_mixed`, `speed_bench_throughput_16k_high_entropy`, `speed_bench_throughput_32k_low_entropy`, `speed_bench_throughput_32k_mixed`, `speed_bench_throughput_32k_high_entropy`]_ +#### `--graph-format` `` + +Select the agentic-workload format for `--input-file`, overriding graph-adapter auto-detection. Registered names: `dynamo_trace`, `weka_trace`, `dag_jsonl`, `native`. Requires `--input-file`. Independent of `--custom-dataset-type` (that selects a custom dataset loader; this selects a graph adapter). Use `native` to load a hand-authored Graph IR workload file, or `dag_jsonl` to run a DAG conversation file on the graph runtime. +
_Choices: [`dynamo_trace`, `weka_trace`, `dag_jsonl`, `native`]_ + #### `--dataset-sampling-strategy` `` Strategy for selecting entries from dataset during benchmarking. `sequential`: Iterate through dataset in order, wrapping to start after end. `random`: Randomly sample with replacement (entries may repeat before all are used). `shuffle`: Shuffle dataset and iterate without replacement, re-shuffling after exhaustion. Default behavior depends on dataset type (e.g., `sequential` for traces, `shuffle` for synthetic).
_Choices: [`random`, `sequential`, `shuffle`]_ +#### `--max-context-length` `` + +Maximum per-trace context length (tokens) for graph-plane dataset selection. Graph traces whose input+output context would exceed this cap are excluded from the eligible set before selection. Applies to graph-adapter workloads (`--input-file` graph traces); ignored by synthetic/public datasets. Unset (the default) applies no context-length filter. +
_Constraints: ≥ 1_ + +#### `--allow-dataset-wrap`, `--no-allow-dataset-wrap` + +Allow graph-plane dataset selection to wrap (reuse the finite trace pool) when `--request-count` exceeds the number of eligible traces. `--allow-dataset-wrap` forces wrapping on; `--no-allow-dataset-wrap` forces it off. Unset (the default, None) defers to a derived default computed at resolution time and surfaced on `run.resolved.allow_dataset_wrap`. Applies to graph-adapter workloads; ignored by synthetic/public datasets. + #### `--random-seed` `` Random seed for deterministic data generation. When set, makes synthetic prompts, sampling, delays, and other random operations reproducible across runs. Essential for A/B testing and debugging. Uses system entropy if not specified. Initialized globally at config creation. @@ -1949,7 +2067,7 @@ The total number of unique conversations to generate. Each conversation represen #### `--num-dataset-entries`, `--num-prompts` `` -Total number of unique entries to generate for the dataset. Each entry represents one user message that can be used as a turn in conversations. Entries are reused across conversations and turns according to `--dataset-sampling-strategy`. Higher values provide more diversity. +Total number of unique entries to generate for the dataset. Each entry represents one user message that can be used as a turn in conversations. Entries are reused across conversations and turns according to `--dataset-sampling-strategy`. Higher values provide more diversity. On the graph plane (graph-adapter workloads), this bounds the distinct traces selected: unset means all eligible traces (after `--max-context-length` and other filters) are used; N means up to N distinct traces are selected after filters. It does not count requests -- see `--request-count`.
_Constraints: ≥ 1_
_Default: `100`_ @@ -2358,6 +2476,15 @@ Maximum input sequence length for filtering. Traces with input_length > max_isl Maximum output sequence length cap. Traces with output_length > max_osl are capped to max_osl.
_Constraints: ≥ 1_ +#### `--prompt-corpus` `` + +Corpus backing recorded graph (`weka_trace`, `dynamo_trace`) real-content synthesis. `coding` (default when unset) uses the procedural CodingContentGenerator pool the recorded weka workloads were captured against; `sonnet` uses the Shakespeare pool, which yields matching token counts but different bytes (useful only to reproduce golden fixtures built from that pool). Ignored by non-graph datasets. + +#### `--synthesis-idle-gap-cap` `` + +Per-trace idle-gap cap (seconds) for recorded graph (`weka_trace`, `dynamo_trace`) replay. Recorded inter-request idle gaps longer than this are compressed to the cap so a recorded multi-hour gap does not park the replay. Defaults to 60s when unset. Ignored by non-graph datasets. +
_Constraints: ≥ 0.0_ + ### Load Generator #### `--benchmark-duration` `` @@ -2395,11 +2522,11 @@ Smoothness parameter for gamma distribution arrivals (--arrival-pattern gamma). #### `--request-count`, `--num-requests` `` -The maximum number of requests to send. If not set, will be automatically determined based on the timing mode and dataset size. For synthetic datasets, this will be `max(10, concurrency * 2)`. Pass a comma-separated list (e.g. `--request-count 100,500,1000`) to sweep over multiple request counts; the converter promotes the list to a sweep on phases.profiling.requests before AIPerfConfig validation. +The maximum number of requests to send. If not set, will be automatically determined based on the timing mode and dataset size. For synthetic datasets, this will be `max(10, concurrency * 2)`. Pass a comma-separated list (e.g. `--request-count 100,500,1000`) to sweep over multiple request counts; the converter promotes the list to a sweep on phases.profiling.requests before AIPerfConfig validation. On the graph plane (graph-adapter workloads), this counts individual requests, not traces -- a single trace can emit many requests, so `--request-count` does not bound the number of distinct traces selected (use `--num-dataset-entries` for that). When it exceeds the eligible trace pool, `--allow-dataset-wrap` controls whether selection wraps to reuse traces. #### `--concurrency-ramp-duration` `` -Duration in seconds to ramp session concurrency from 1 to target. Useful for gradual warm-up of the target system. +Duration in seconds to ramp session concurrency from 1 to target. On graph-IR replay (weka/dynamo/dag traces) this ramps LANE admission: concurrent trace-instance lanes are admitted 1 -> --concurrency over the duration. Useful for gradual warm-up of the target system.
_Constraints: > 0_ #### `--prefill-concurrency-ramp-duration` `` @@ -2449,6 +2576,16 @@ Maximum adaptive scale control value. Inferred from the phase target when omitte SLA filter for adaptive scale. Format: 'metric_tag:stat:op:threshold'. Latency-family metrics (request_latency, time_to_first_token/ttft, inter_token_latency/itl/tpot) support percentile stats; window scalar/rate metrics (request_throughput, output_token_throughput, goodput, goodput_ratio, success_rate, error_rate, cancellation_rate) support {avg, min, max}. Full metric/stat table: [Adaptive SLA metric support](tutorials/yaml-config.md#adaptive-sla-metric-support). op in {lt, le, gt, ge}; threshold is a float. Repeatable. Example: --adaptive-scale-sla 'request_latency:p95:le:30000'. +#### `--trajectory-start-min-ratio` `` + +Lower bound (fraction of each trace's recorded wall-clock duration) of the per-trace t* snapshot-instant sampling window for recorded graph replay. Unset resolves to 0.0. With both bounds at 0.0 the window is OFF (full native replay). Per trace, t* is drawn uniformly from [min, max] x trace_duration with a trace-salted RNG: nodes firing before t* become cache-priming WARMUP history, nodes at/after t* are PROFILING. Set min == max for a deterministic t*. Must be <= --trajectory-start-max-ratio. The inferencex-agentx-mvp scenario auto-applies 0.0 when unset and locks an explicit conflicting value. +
_Constraints: ≥ 0.0, ≤ 1.0_ + +#### `--trajectory-start-max-ratio` `` + +Upper bound (fraction of each trace's recorded wall-clock duration) of the per-trace t* snapshot-instant sampling window for recorded graph replay. Unset resolves to 0.0 (window OFF, full native replay); any value > 0 engages the t* snapshot machinery (per-trace t* sampling plus the auto boundary-priming warmup phase). See --trajectory-start-min-ratio for the sampling contract. The inferencex-agentx-mvp scenario auto-applies 1.0 when unset and locks an explicit conflicting value. +
_Constraints: ≥ 0.0, ≤ 1.0_ + ### Warmup #### `--warmup-request-count`, `--num-warmup-requests` `` @@ -2461,6 +2598,16 @@ The maximum number of warmup requests to send before benchmarking. If not set an The maximum duration in seconds for the warmup phase. If not set, it will use the `--warmup-request-count` value. If neither are set, no warmup phase will be used.
_Constraints: > 0_ +#### `--agentic-cache-warmup-duration` `` + +Extended (cache-pressure) warmup duration in seconds for the recorded-trace graph-IR replay path. After the boundary-priming warmup drains, AIPerf continues the live post-t* replay with zero idle delay and one-token outputs for this long, recycling fresh templates onto freed lanes, then drains and hands each lane's execution frontier to profiling (with residual delays), driving the server KV cache to steady-state pressure before benchmarking. If not set, no extended warmup runs. +
_Constraints: > 0_ + +#### `--burst-phase-starts` + +Burst-vs-spread control for phase starts on the recorded-trace graph-IR replay path. Default False = SPREAD: each WARMUP / PROFILING trace's first firing is offset by its recorded (idle-gap-warped, <=60s) lead from the t* snapshot instant. True = BURST: both phase starts collapse into a synchronized burst -- every warmup priming credit and every trace's earliest profiling firing fire at once at phase-time 0 (the leading per-firing dispatch offset is dropped); relative inter-turn delays after the first are still honored. Governs ONLY the two phase-start dispatch patterns. +
_Flag (no value required)_ + #### `--num-warmup-sessions` `` The number of sessions to use for the warmup phase. If not set, it will use the `--warmup-request-count` value. @@ -3046,6 +3193,17 @@ Directory path for ZMQ IPC (Inter-Process Communication) socket files. When usin Select the ZMQ dual-bind communication backend (IPC + TCP). All dual-bind knobs are cluster-managed; this flag only selects the discriminator and the converter routes downstream to the default.
_Flag (no value required)_ +### Scenario + +#### `--scenario` `` + +Lock all benchmark invariants for a named scenario (e.g. 'inferencex-agentx-mvp'). Conflicts with the locked invariants raise ScenarioLockError at startup unless --unsafe-override is also passed. Distinct from the sweep ``scenarios`` strategy (hand-picked named runs). + +#### `--unsafe-override` + +Convert scenario lock errors to warnings; stamps submission_valid=false in the aggregate output. No-op without --scenario. +
_Flag (no value required)_ +
## `aiperf speed-bench-report` diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 2a58f8f33e..0de8dd3904 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -39,6 +39,15 @@ Accuracy benchmark settings. Tunables for accuracy benchmarking: the cancel-path | `AIPERF_ACCURACY_CANCEL_RESULT_WAIT_SEC` | `5.0` | ≥ 0.0 | Bounded time (seconds) the SystemController waits on the cancel (Ctrl+C) path for the RecordsManager's ProcessAccuracyResultMessage before stopping. The normal completion path blocks on the accuracy shutdown gate indefinitely, but the cancel path must not hang forever, so it waits at most this long for the graded accuracy summary to arrive over pub/sub before proceeding to export. Set to 0 to skip the wait entirely. | | `AIPERF_ACCURACY_LCB_RELEASE_TAG` | `'v4_v5'` | — | LiveCodeBench dataset subset (HF config name) passed as the positional ``name`` arg to ``load_dataset("livecodebench/code_generation_lite", name, split="test", trust_remote_code=True)``. Pins which monthly snapshot LCB serves so accuracy numbers are reproducible across runs and branches. Default ``v4_v5`` matches lighteval's base subset; bump (e.g. to ``v6``) when the team rebaselines against a newer snapshot. The loader always passes ``trust_remote_code=True`` so LCB's dataset-loading script can execute on ``datasets`` v4+ (mirrors lighteval's reference opt-in). Consumed by ``aiperf.accuracy.benchmarks.lcb_codegeneration``. | +## AGENTX + +InferenceX AgentX scenario runtime settings. Runtime knobs for the locked AgentX scenario submission contract. The context-overflow substring allowlist and rate threshold flip a scenario run to ``submission_valid=false`` when the runtime overflow rate exceeds the limit (RFC §7). Mirrors ``Environment.AGENTX`` on ``ajc/agentx``. + +| Environment Variable | Default | Constraints | Description | +|----------------------|---------|-------------|-------------| +| `AIPERF_AGENTX_CONTEXT_OVERFLOW_SUBSTRINGS` | `['context length', 'maximum context', 'context_length_exceeded', 'prompt is too long']` | — | Case-insensitive substring allowlist used to classify a server error response as a context-overflow event. Matched against the raw response body and the OpenAI-style nested 'error.message' field. Extend via AIPERF_AGENTX_CONTEXT_OVERFLOW_SUBSTRINGS to support additional inference-server vocabularies (vLLM, TGI, TensorRT-LLM, ...). Empty list disables runtime detection. | +| `AIPERF_AGENTX_CONTEXT_OVERFLOW_RATE_LIMIT` | `0.01` | ≥ 0.0, ≤ 1.0 | Strict upper bound on the per-run context-overflow rate (context_overflow_count / total_responses) before a scenario submission is flipped to submission_valid=false with reason 'context_overflow_rate_exceeded'. Default 0.01 (1%) matches the scenario spec RFC §7. Comparison is strictly greater-than: rate exactly equal to the limit is accepted. Has no effect on non-scenario runs (no --scenario flag) or runs with zero responses. | + ## APISERVER API server settings. Controls the host and port of the API server. @@ -90,6 +99,26 @@ Dataset loading and configuration. Controls timeouts and behavior for dataset lo | `AIPERF_DATASET_MEDIA_DOWNLOAD_TIMEOUT` | `60.0` | ≥ 1.0, ≤ 100000.0 | Timeout in seconds per media URL download when inline encoding is required | | `AIPERF_DATASET_MEDIA_DOWNLOAD_MAX_CONCURRENCY` | `10` | ≥ 1, ≤ 100 | Maximum number of concurrent media URL downloads | | `AIPERF_DATASET_INLINE_RECORDS_WARN_THRESHOLD` | `500` | ≥ 1 | Soft warning threshold for the number of inline `records:` entries on a `FileDataset`. When total inline records exceed this value, the config loader logs a warning suggesting the user move the dataset to a JSONL file. No hard cap. | +| `AIPERF_DATASET_WEKA_GRAPH_PARALLEL_THRESHOLD` | `8` | ≥ 0 | Minimum number of Weka graph rows/files required before the graph adapter switches from serial parsing to the multi-process forkserver parse path. Set to 0 to force the pool path for any non-empty corpus. | +| `AIPERF_DATASET_WEKA_GRAPH_PARALLEL_WORKERS` | `0` | ≥ 0, ≤ 256 | Number of worker processes for Weka graph parsing. 0 = auto (min(cpu_count - 1, WEKA_GRAPH_PARALLEL_AUTO_MAX_WORKERS, item_count when known)), matching AgentX's full-corpus reconstruction default. Set to 1 to force a single worker. | +| `AIPERF_DATASET_WEKA_GRAPH_PARALLEL_AUTO_MAX_WORKERS` | `16` | ≥ 1, ≤ 256 | Upper bound used when WEKA_GRAPH_PARALLEL_WORKERS=0 selects auto worker sizing for Weka graph parsing. | +| `AIPERF_DATASET_WEKA_GRAPH_PARALLEL_PREFETCH_MULTIPLIER` | `16` | ≥ 1, ≤ 64 | Ordered pool submit-window multiplier for Weka graph parsing. The adapter keeps at most workers * multiplier rows in flight so heavy traces do not idle fast workers while still bounding parent-side memory. Window-sizing rule: the window (workers * multiplier) must cover the rows remaining behind the single heaviest trace, or fast workers stall head-of-line while that one trace drains; a multiplier of 16 gives a window of 256 at the auto 16 workers, which covers this 393-row corpus whose heaviest row sits at index 140. Memory tradeoff: the parent buffers up to workers * multiplier completed payloads plus the raw rows queued behind them; measured on the 062126 full corpus at 16 workers this costs ~7.3 GiB parent VmHWM (~17.5 GiB process tree), vs ~2.8 GiB parent at multiplier 4. | +| `AIPERF_DATASET_WEKA_GRAPH_PARALLEL_ITEM_TIMEOUT_SECONDS` | `600.0` | > 0.0 | Per-item timeout in seconds for one Weka graph parse pool result (a single trace file/row). A raw multiprocessing Pool never completes the in-flight task of a worker killed mid-parse (OOM kill / external SIGKILL), so an unbounded wait presents as a silent indefinite hang; this bound turns it into a clear RuntimeError instead. Raise it for corpora whose single heaviest trace legitimately parses slower than the default. | +| `AIPERF_DATASET_DYNAMO_GRAPH_PARALLEL_THRESHOLD` | `8` | ≥ 0 | Minimum number of Dynamo session-trees in ONE `from_dynamo_trace` call required before the graph adapter switches from the serial tree-scoped build to the multi-process forkserver build path. The parallel unit is TREES (a root session plus its subagent descendants), so even one large file with many trees parallelizes. At or below this count the build stays serial (no pool spawn). Set to 0 to force the pool path for any multi-tree capture. | +| `AIPERF_DATASET_DYNAMO_GRAPH_PARALLEL_WORKERS` | `0` | ≥ 0, ≤ 256 | Number of worker processes for Dynamo session-tree graph building. 0 = auto (min(cpu_count - 1, DYNAMO_GRAPH_PARALLEL_AUTO_MAX_WORKERS, tree_count)). Set to 1 to force the serial path (single-tree captures always build serially). | +| `AIPERF_DATASET_DYNAMO_GRAPH_PARALLEL_AUTO_MAX_WORKERS` | `16` | ≥ 1, ≤ 256 | Upper bound used when DYNAMO_GRAPH_PARALLEL_WORKERS=0 selects auto worker sizing for Dynamo session-tree graph building. | +| `AIPERF_DATASET_DYNAMO_GRAPH_PARALLEL_PREFETCH_MULTIPLIER` | `16` | ≥ 1, ≤ 64 | Ordered pool submit-window multiplier for Dynamo session-tree graph building, doubling as the batches-per-worker factor: the trees are shuffled into at most workers * multiplier CONTIGUOUS weight-balanced per-batch temp files (balanced by a cumulative record-line byte-length proxy resolved WITHOUT parsing the recorded `input_sequence_hashes`) and at most that many batches are kept in flight. Contiguous batching preserves the global sorted-by-root tree order so the parent's in-order union is byte-identical to the serial build; a larger multiplier gives finer batches so a few heavy trees do not idle fast workers, at the cost of more parent-buffered batch results. | +| `AIPERF_DATASET_DYNAMO_GRAPH_PARALLEL_ITEM_TIMEOUT_SECONDS` | `600.0` | > 0.0 | Per-item timeout in seconds for one Dynamo session-tree batch build pool result. A raw multiprocessing Pool never completes the in-flight task of a worker killed mid-build (OOM kill / external SIGKILL), so an unbounded wait presents as a silent indefinite hang; this bound turns it into a clear RuntimeError instead. Raise it for captures whose single heaviest batch legitimately builds slower than the default. | +| `AIPERF_DATASET_WEKA_HF_SPLIT` | `'train'` | — | HuggingFace `datasets` split loaded when a weka graph workload is given as a HuggingFace `org/name` repo id (e.g. `semianalysisai/cc-traces-weka-062126`) instead of a local path. The published weka corpora ship a single `traces.jsonl` exposed as the `train` split, so the default is `train`. Datasets slice syntax is accepted (e.g. `train[:100]`) to bound a streamed slice without forcing a full-split download. | +| `AIPERF_DATASET_WEKA_HF_REVISION` | `None` | — | Optional HuggingFace dataset revision (commit SHA, branch, or tag) pinned when loading a weka corpus by repo id. Pinning a commit SHA makes a streamed ingest reproducible across re-runs even if the dataset's `main` branch is updated. None loads the default revision. | + +## DYNAMO + +Dynamo agent-trace adapter tunables. Bounds the `parent_link` chain depth (cycle guard) walked while flattening the Dynamo trace into the shared segment-trie IR (`_guard_chain_forest` in `src/aiperf/dataset/graph/adapters/dynamo/trace.py`). + +| Environment Variable | Default | Constraints | Description | +|----------------------|---------|-------------|-------------| +| `AIPERF_DYNAMO_MAX_SUBAGENT_DEPTH` | `16` | ≥ 1 | Maximum `parent_link` chain depth walked by the Dynamo trie-IR cycle guard. Chains deeper than this raise DynamoTraceAdapterError. | ## GPU @@ -105,6 +134,22 @@ GPU telemetry collection configuration. Controls GPU metrics collection frequenc | `AIPERF_GPU_SHUTDOWN_DELAY` | `5.0` | ≥ 1.0, ≤ 300.0 | Delay in seconds before shutting down GPU telemetry service to allow command response transmission | | `AIPERF_GPU_THREAD_JOIN_TIMEOUT` | `5.0` | ≥ 1.0, ≤ 300.0 | Timeout in seconds for joining GPU telemetry collection threads during shutdown | +## GRAPH + +Graph-mode runtime tunables. Controls async-dataflow graph runtime guardrails and adapter behavior. Every live producer lowers to flat LlmNode graphs. + +| Environment Variable | Default | Constraints | Description | +|----------------------|---------|-------------|-------------| +| `AIPERF_GRAPH_DYNAMIC_POOL_MAX_BYTES` | `64 * 1024 * 1024` | ≥ 1 | Per-worker byte cap for the graph dynamic-content pool (captured assistant responses keyed by trace/variant/ordinal, spliced into successor prompts under per-trace sticky routing). LRU-evicts whole trace entries when exceeded; an evicted live trace surfaces as a loud pool-missing error at its consumer, never a silent truncation. Load-bearing backstop for lost GraphTraceEnd notifications. | +| `AIPERF_GRAPH_MERGE_CONSECUTIVE_USER` | `False` | — | When a dynamic-content producer fails or returns no replayable content, its assistant turn is OMITTED from the successor's spliced prompt (dynamic-content-pool spec 6.1), which can leave consecutive user-role messages (e.g. an init user turn followed by a static user turn). Some chat APIs reject consecutive same-role messages. When True, the worker merges consecutive user messages in a spliced prompt into one (contents joined by a newline). Default False sends the messages as-is. | +| `AIPERF_GRAPH_IGNORE_EDGE_DELAYS` | `False` | — | When true, the executor skips honoring StaticEdge.delay_after_predecessor_us, StaticEdge.delay_after_predecessor_start_us and StaticEdge.min_start_delay_us — captured inter-node idle gaps are collapsed and successor nodes fire as soon as their channel-readiness inputs are satisfied. Use for compressed-time replays where the user wants to stress the server independent of the original trace's wall-clock pacing. Default (false): honor edge delays exactly as captured. | +| `AIPERF_GRAPH_WARMUP_MAX_OUTPUT_TOKENS` | `1` | ≥ 1 | Per-request output-token cap applied UNCONDITIONALLY to WARMUP boundary turns (via a `max_tokens` dispatch override in `rewrite_for_warmup` / the warmup delta envelope), so the warmup wire payload always carries `max_tokens == 1`. Warmup is cache-priming history: it primes each live chain's input prefix, and the (1-token) generated output is discarded (warmup records are excluded from metrics). Does NOT affect the PROFILING phase, which replays each turn's full recorded output length. | +| `AIPERF_GRAPH_HANDOFF_RESIDUAL_CAP` | `60.0` | > 0 | Upper bound in seconds on any single extended-warmup handoff residual delay (the recorded inter-turn gap minus drain time a resumed profiling frontier waits before firing; see `chop_trie_at_frontier`). Graph edge delays are NOT bounded by the build-plane idle-gap warp when another stream's active interval covers the gap, so this explicit clamp keeps a handoff lane from parking for minutes at the profiling start; the 60s default matches the recorded idle-gap cap. None disables the clamp (faithful-but-unbounded residuals). | +| `AIPERF_GRAPH_PRESSURE_DRAIN_GRACE_CAP` | `300.0` | > 0 | Upper bound in seconds on the extended (cache-pressure) warmup phase's drain grace period. When a pressure duration is set and no explicit --warmup-grace-period overrides it, the warmup phase waits min(duration, this cap) for in-flight returns after sending completes -- instead of the boundary-priming warmup's unbounded wait, so a wedged or lost return cannot hang the run. On grace expiry the in-flight requests are cancelled; cancelled turns are excluded from the handoff ledger (not-executed, profiling refires them), so a drain whose cancellations land yields a VALID handoff, and a drain that force-completes with unreturned credits trips the completeness gate and skips the re-cut (profiling then starts from the plain t* plans). | +| `AIPERF_GRAPH_DISPATCH_TIMEOUT` | `300.0` | > 0 | Deadlock guard for the CreditDispatchAdapter Future bridge. The executor's `dispatch(...)` issues a graph credit and awaits the correlated CreditReturn via an asyncio.Future; this caps that wait. If a return is dropped (worker crash, lost message, unrouted return) the Future is rejected with TimeoutError instead of hanging the trace's TaskGroup forever. Sized to comfortably exceed a slow single dispatch round-trip (large prefill + long output); raise it for very long generations, lower it to fail faster in CI. | +| `AIPERF_GRAPH_EXECUTOR_WATCHDOG_TIMEOUT` | `None` | > 0 | Optional wall-clock deadlock guard on TraceExecutor.run's firing loop. DISPATCH_TIMEOUT only bounds a node that ISSUED a credit; a node wedged on an unsatisfiable channel input (a producer-accounting bug in mark_producer_done / fan-in gating) never dispatched and so has no per-dispatch guard. When set, each executor run is wrapped in asyncio.wait_for(timeout=this) so such a wedge raises TimeoutError instead of hanging. Default None (disabled): bare/count/session runs replay recorded idle gaps faithfully with no wall-clock bound; a --benchmark-duration already bounds the phase. Set a generous value in CI/dev to convert an executor deadlock into a diagnosable failure without truncating a legitimate long gap. | +| `AIPERF_GRAPH_IDLE_GAP_NO_DURATION_WARN_SECONDS` | `30.0` | ≥ 0.0, ≤ 86400.0 | Threshold (seconds) for the once-per-run PROFILING-phase advisory GraphIRReplayStrategy logs when an idle-gap recorded corpus is run WITHOUT `--benchmark-duration`. A faithful recorded-trace replay honors every recorded inter-turn idle gap verbatim (each per-gap-warped to the `--synthesis-idle-gap-cap`), so a count/session/bare run spans the slowest admitted trace's full recorded wall time -- sending completes only after every scheduled idle gap elapses. When any inter-turn edge delay exceeds this threshold and no duration budget is set, the strategy emits a NOTICE advising `--benchmark-duration` (the bound) or Ctrl+C (graceful partial export). Set 0.0 to always advise; raise to silence the advisory for corpora with shorter gaps. | + ## HTTP HTTP client socket and connection configuration. Controls low-level socket options, keepalive settings, DNS caching, and connection pooling for HTTP clients. These settings optimize performance for high-throughput streaming workloads. Video Generation Polling: For async video generation APIs that use job polling (e.g., SGLang /v1/videos), the poll interval is controlled by AIPERF_HTTP_VIDEO_POLL_INTERVAL. The max poll time uses the --request-timeout-seconds CLI argument. diff --git a/docs/index.yml b/docs/index.yml index 192b70a0d9..98edd23442 100644 --- a/docs/index.yml +++ b/docs/index.yml @@ -187,10 +187,13 @@ navigation: path: benchmark-modes/timing-modes-reference.md - page: Trace Replay with Mooncake Traces path: benchmark-modes/trace-replay.md + - page: Agentic Workload Benchmarks + path: benchmark-modes/agentic.md - page: Baseten Trace Replay path: tutorials/baseten-trace.md - - page: Conversation DAG Benchmarks + - page: Conversation DAG Benchmarks (Legacy) path: benchmark-modes/dag.md + slug: conversation-dag-benchmarks - section: Accuracy collapsed: true @@ -223,6 +226,27 @@ navigation: path: reference/api-endpoints.md - page: YAML Config Roadmap path: reference/yaml-config-roadmap.md + - section: Graph IR & Replay Internals + collapsed: true + contents: + - page: Graph IR Schema + path: reference/graph-ir-schema.md + - page: Graph IR Validation + path: reference/graph-ir-validation.md + - page: Graph Ingest & Build Pipeline + path: reference/graph-ingest-build-pipeline.md + - page: Graph Structural Sidecar Handoff + path: reference/graph-structural-handoff.md + - page: Graph Segment Unified Store + path: reference/graph-segment-unified-store.md + - page: AIPerf Graph Async Dataflow Runtime Architecture + path: reference/graph-async-dataflow-runtime.md + - page: Graph Worker Materialization + path: reference/graph-worker-materialization.md + - page: Graph Trie-Route Prompt Convention + path: reference/graph-trie-prompt-convention.md + - page: Graph Runtime Troubleshooting + path: reference/graph-runtime-troubleshooting.md - section: Server Metrics collapsed: true diff --git a/docs/metrics-reference.md b/docs/metrics-reference.md index 36f377af4f..d425c670a6 100644 --- a/docs/metrics-reference.md +++ b/docs/metrics-reference.md @@ -97,6 +97,7 @@ This document provides a comprehensive reference of all metrics available in AIP - [Request Latency](#request-latency) - [Request Throughput](#request-throughput) - [Request Count](#request-count) + - [Streamed Request Count](#streamed-request-count) - [Error Request Count](#error-request-count) - [Minimum Request Timestamp](#minimum-request-timestamp) - [Maximum Response Timestamp](#maximum-response-timestamp) @@ -195,7 +196,7 @@ any knowledge of the individual request/response data. ## Streaming Metrics > [!NOTE] -> All metrics in this section require the `--streaming` flag with a token-producing endpoint and at least one non-empty response chunk. +> All metrics in this section are computed **over streamed requests only** — the per-record subset counted by [Streamed Request Count](#streamed-request-count) — and require a token-producing endpoint with at least one non-empty response chunk. Applicability is per-record: a request contributes iff it streamed on the wire (`RequestRecord.streamed`). Standard runs stream when launched with the global `--streaming` flag; graph-IR replays (weka / dynamo / native graph) stream individual recorded nodes from their own recorded mode even when the global flag is off, so a non-streamed request is simply excluded rather than dragging its full request latency in as, e.g., a first-token time. ### Time to First Token (TTFT) @@ -216,6 +217,7 @@ ttft_seconds = ttft_ns / 1e9 ``` **Notes:** +- Computed over streamed requests only ([Streamed Request Count](#streamed-request-count) is the denominator); non-streamed requests are excluded so their full request latency is never miscounted as a first-token time. - Includes network latency, queuing time, prompt processing, and generation of the first token (or chunk of tokens). - Raw timestamps are in nanoseconds; converted to milliseconds for display and seconds for rate calculations. - Response chunks refer to individual messages with non-empty content received during streaming. @@ -238,6 +240,7 @@ ttst_ms = ttst_ns / 1e6 ``` **Notes:** +- Computed over streamed requests only ([Streamed Request Count](#streamed-request-count) is the denominator); non-streamed requests are excluded. - Requires at least 2 non-empty response chunks to compute the time between first and second tokens. - Raw timestamps are in nanoseconds; converted to milliseconds for display. @@ -261,6 +264,7 @@ ttfo_ms = ttfo_ns / 1e6 ``` **Notes:** +- Computed over streamed requests only ([Streamed Request Count](#streamed-request-count) is the denominator); non-streamed requests are excluded. - TTFO vs TTFT: Time to First Output (TTFO) measures time to the first non-reasoning token, while Time to First Token (TTFT) measures time to any first token including reasoning tokens. For models without reasoning, TTFO and TTFT are equivalent. - Non-reasoning tokens include TextResponseData with non-empty text, or ReasoningResponseData with non-empty content field (regardless of reasoning field). - Requires at least one non-empty non-reasoning response chunk. @@ -286,6 +290,7 @@ inter_token_latency_ms = inter_token_latency_ns / 1e6 ``` **Notes:** +- Computed over streamed requests only ([Streamed Request Count](#streamed-request-count) is the denominator); non-streamed requests are excluded. - Requires at least 2 non-empty response chunks and valid `time_to_first_token`, `request_latency`, and `output_sequence_length` metrics. - Result is in seconds when used for throughput calculations (Output Token Throughput Per User). @@ -303,6 +308,7 @@ inter_chunk_latency = [request.content_responses[i].perf_ns - request.content_re ``` **Notes:** +- Computed over streamed requests only ([Streamed Request Count](#streamed-request-count) is the denominator); non-streamed requests are excluded. - Requires at least 2 response chunks. - Unlike ITL (which produces a single average), ICL provides the full distribution of inter-chunk times. - Useful for detecting variability, jitter, or issues in streaming delivery. @@ -1502,6 +1508,23 @@ request_count = sum(1 for r in records if r.valid) --- +### Streamed Request Count + +**Type:** [Aggregate Metric](#aggregate-metrics) + +The total number of requests that were sent in streaming mode (i.e. actually streamed on the wire). This is the streaming denominator displayed beside [Request Count](#request-count): the streaming record metrics ([TTFT](#time-to-first-token-ttft), [TTST](#time-to-second-token-ttst), [TTFO](#time-to-first-output-token-ttfo), [ITL](#inter-token-latency-itl), [ICL](#inter-chunk-latency-icl)) are computed over exactly this subset, not over all requests. + +**Formula:** +```python +streamed_request_count = sum(1 for r in records if r.valid and r.request.streamed) +``` + +**Notes:** +- Membership is per-record: a request counts iff it streamed on the wire (`RequestRecord.streamed`), independent of the global `--streaming` flag. In a graph-IR replay (weka / dynamo / native graph) individual recorded nodes stream from their own recorded mode even when the run's global `--streaming` flag is off, so this count can be non-zero without `--streaming`. +- Omitted from the export entirely when no request streamed (a run with zero streamed requests reports no `streamed_request_count`). + +--- + ### Error Request Count **Type:** [Aggregate Metric](#aggregate-metrics) @@ -1875,6 +1898,25 @@ http_req_chunks_received = trace.response_chunks_count --- +## Trace Replay Metrics + +> [!NOTE] +> Metrics in this section are produced by a standalone results accumulator (not the standard `MetricRegistry` derivation walk) and are emitted only when a trace loader supplies the underlying per-turn metadata. They are absent on workloads that do not. Each tag is still registered as a display-only metric class (following the [Externally-Injected Derived Metric pattern](dev/patterns.md#externally-injected-derived-metric-pattern)) so the realtime dashboard and console exporter can resolve its display metadata. + +### Theoretical Prefix Cache Hit + +**Type:** Standalone results accumulator (`theoretical_prefix_cache`); display-only class `TheoreticalPrefixCacheHitMetric` + +**Tag:** `theoretical_prefix_cache_hit` · **Unit:** `%` · **Console group:** [`CACHE`](#group-cache) (rendered as `LLM Metrics: Cache`) + +The cumulative infinite-cache prefix-hit rate over the trace's KV-block hash ids (WEKA: recorded hash ids; Dynamo: recorded `input_sequence_hashes` whenever a record carries replay metadata, per-session virtual block ids otherwise). The shared trie build (`segment_ir.prefix_cache.stamp_theoretical_prefix_cache`) walks every request's `hash_ids` at parse time in recorded global-time order under one shared per-trace seen-set and stamps each node with two integers — leading prefix-cache `hit_blocks` and `total_blocks` — under `LlmNode.metadata["dispatch"]["theoretical_prefix_cache"]`. The seen-set is per trace file even for `hash_id_scope: "global"` corpora (recorded `t` is conversation-relative, so no cross-file ordering exists); under global scope the stamped counts are therefore a lower bound — a block another trace already sent still counts as a miss. The accumulator sums them over valid profiling records: + +```text +theoretical_prefix_cache_hit = 100 * sum(hit_blocks) / sum(total_blocks) +``` + +This is the THEORETICAL ceiling an infinite per-trace prefix cache would achieve for the replayed request mix; it carries no hash ids through the request path at runtime (only a per-node metadata lookup plus two integer adds). Emitted for trie graph workloads (WEKA and Dynamo traces) on every ingest route — local file, directory, and HuggingFace-streamed corpora alike: the streaming store build recovers the per-node counts from the merged content-free structural graph. Absent when that structural merge fails (fail-soft; the run then also skips the `graph_meta` sidecar), for graph workloads with no recorded KV-block hash ids to walk (native and dag_jsonl never stamp prefix-cache counts), and for all non-graph datasets. Matches the agentx v0.4-sync metric of the same name, formula, and unit. + ## GPU Power Efficiency Metrics > [!NOTE] diff --git a/docs/plugins/plugin-system.md b/docs/plugins/plugin-system.md index 838ec32eb2..30df610fb4 100644 --- a/docs/plugins/plugin-system.md +++ b/docs/plugins/plugin-system.md @@ -120,6 +120,19 @@ AIPerf supports 34 plugin categories organized by function, including `api_route | `dataset_composer` | `ComposerType` | Dataset generation (synthetic, custom, rankings) | | `custom_dataset_loader` | `CustomDatasetType` | JSONL format loaders | | `public_dataset_loader` | `PublicDatasetType` | Shared benchmark dataset fetchers (HTTP, HuggingFace) | +| `graph_adapter` | `GraphAdapterType` | Trace/log-to-`ParsedGraph` converters for graph benchmark mode (weka, dynamo, native, dag_jsonl) | + +Graph adapters implement `GraphAdapterProtocol` +(`src/aiperf/dataset/graph/adapters/protocols.py`): `can_load(path)` is a cheap +extension + content sniff for auto-detection (ties broken by the metadata +`detection_priority`, higher wins), and `parse(path, ctx)` does the full +conversion. `ctx` is an optional `GraphParseContext` carrying run-derived parse +knobs (content seed, tokenizer, corpus, idle-gap cap, dispatch defaults, ...); +each adapter maps only the fields it consumes and forwards a field ONLY when it +is set, so a ctx-less `parse(path)` stays byte-equal to the protocol-default +entry and a partial ctx never clobbers a non-`None` entry default. Adapters may +raise adapter-specific `ValueError` subclasses; the parser layer wraps them +into `GraphParseError`. ### Endpoint and Transport Categories @@ -129,6 +142,50 @@ AIPerf supports 34 plugin categories organized by function, including `api_route | `endpoint` | `EndpointType` | API endpoint implementations (chat, completions, embeddings, etc.) | | `transport` | `TransportType` | Network transport (HTTP via aiohttp) | +### Session Routing Category + +| Category | Enum | Description | +|----------|------|-------------| +| `session_routing` | `SessionRoutingType` | Stamps per-session identity (headers or body metadata) onto outbound requests so an external router pins every turn of a session to one worker; selected via `--session-routing` | + +**Purpose.** A session-routing plugin gives an external router (SGLang Model Gateway, Dynamo, a +generic session-affinity load balancer) the identity of the session a request belongs to, so all of +a conversation's turns re-land on the replica holding its KV prefix. Exactly one mode runs per +invocation; the selected plugin is instantiated once per worker by `InferenceClient` and invoked at +the request-serialization chokepoint. The base class is `SessionRoutingBase` +(`src/aiperf/workers/session_routing.py`); passthrough is the default (no headers, body unchanged). + +**Built-ins:** + +| Name | Class | Description | +|------|-------|-------------| +| `dynamo_headers` | `DynamoHeadersRouting` | Dynamo session affinity via `X-Dynamo-Session-ID`, plus `X-Dynamo-Parent-Session-ID` on subagent children. No options. | +| `dynamo_nvext` | `DynamoNvextRouting` | Dynamo session affinity via `nvext.session_control` request-body metadata (bind on non-final turns, close on the final turn). Only for Dynamo builds that implement `session_control`. Option: `timeout_seconds` (default 300). | +| `smg_routing_key` | `SmgRoutingKeyRouting` | SGLang Model Gateway `manual`-policy stickiness via `X-SMG-Routing-Key`. No options. | +| `session_id_header` | `SessionIdHeaderRouting` | Generic additive session-affinity header carrying the session's correlation ID. Option: `header_name` (default `X-Session-ID`). | + +**Options (`--session-routing-opt key=value`).** Each plugin exposes an `Options` Pydantic model +(`extra="forbid"`, so unknown keys are rejected at startup). Repeated `--session-routing-opt` +pairs populate it; values are coerced to the model's field types and canonicalized at config +resolution, so downstream code always sees typed values. `--session-routing-opt` without +`--session-routing` is an error, and parameterless modes (`dynamo_headers`, `smg_routing_key`) +reject every opt key. + +**`mutates_body`.** A plugin sets the `mutates_body` class var to `True` when `transform_body` +changes the payload (only `dynamo_nvext` does); header-only modes leave the body untouched. It is +protocol metadata that marks a mode as incompatible with any verbatim-bytes request path. +`transform_body` must never mutate its input — the request path shares cached `Turn.raw_payload` +dicts with the dataset — so it returns a new dict. + +**`on_session_end` contract.** Fires strictly after the session's last worker-side activity, on +every terminal path (final turn, cancellation, cancel-before-start). It is post-session cleanup +only and must be idempotent (default no-op). + +**Stateful-plugin rule.** A stateful plugin must key its instance state on `ctx.x_correlation_id` +only. A session tree deliberately spans workers, so tree-keyed worker state fragments across +processes. For tree-scoped behavior, use the stateless per-request context facts +`root_correlation_id` and `is_tree_final` instead of accumulating state. + ### Processing Categories | Category | Enum | Description | @@ -316,6 +373,7 @@ Category-specific metadata is validated against Pydantic models in `aiperf.plugi | `PlotMetadata` | `display_name`, `category` | | `ServiceMetadata` | `required`, `auto_start`, `disable_gc`, `replicable` | | `GPUTelemetryCollectorMetadata` | `is_local` | +| `GraphAdapterMetadata` | `detection_priority` | ## CLI Commands diff --git a/docs/reference/graph-async-dataflow-runtime.md b/docs/reference/graph-async-dataflow-runtime.md new file mode 100644 index 0000000000..2f527f7901 --- /dev/null +++ b/docs/reference/graph-async-dataflow-runtime.md @@ -0,0 +1,532 @@ + + +# AIPerf Graph Async Dataflow Runtime Architecture + +Internal developer reference for the Graph v1 async dataflow runtime used by +Weka graph replay. This document describes runtime execution, cross-trace +admission, and credit backpressure. It intentionally +separates executable runtime behavior from static graph analysis helpers. + +Primary implementation files: + +- `src/aiperf/graph/executor.py` +- `src/aiperf/graph/channel_store.py` +- `src/aiperf/graph/scheduler.py` +- `src/aiperf/graph/credit_dispatch_adapter.py` +- `src/aiperf/graph/dispatch/` +- `src/aiperf/timing/strategies/graph_ir_replay.py` + +Related graph references: + +- [Agentic Workload Benchmarks](../benchmark-modes/agentic.md) for user-facing native Graph IR authoring and graph adapter selection. +- [Graph IR Schema](./graph-ir-schema.md) for authorable records, node types, channels, and trace fields. +- [Graph IR Validation](./graph-ir-validation.md) for parser and validator rule behavior. +- [Graph Ingest and Build Pipeline](./graph-ingest-build-pipeline.md), [Graph Structural Sidecar Handoff](./graph-structural-handoff.md), and [Graph Segment Unified Store](./graph-segment-unified-store.md) for build-plane storage contracts. +- [Graph Worker Materialization](./graph-worker-materialization.md) for worker-side request reconstruction. +- [Graph Runtime Troubleshooting](./graph-runtime-troubleshooting.md) for stall, timeout, and return-routing diagnostics. + +## Runtime object model + +The graph runtime is built around a per-graph `TraceExecutor`. The executor owns +immutable, graph-derived collaborators: the parsed graph, the scheduler, the +injected clock, static producer counts, timing flags (`compress_edge_delays`, +`absolute_start_offsets`), and the credit issuer. The per-trace +`VersionedChannelStore` is built per run via a module-level store factory and +owns its own reducer registry. Mutable per-run state lives in `_TraceContext`, +which carries the trace, channel store, the task group, scheduled-node +bookkeeping (`scheduled_node_ids` / `tasks_by_node_id`), per-node finish / +dispatch / first-token wall timestamps, per-node first-token latches, and an +`overflow_terminated` flag for early termination. + +`GraphIRReplayStrategy` is the phase-level driver for Weka replay. It owns trace +admission, lane fan-out and recycle, per-instance `CreditDispatchAdapter` +construction, graph return routing, and phase completion signaling. A graph +credit carries an instance id in `credit.trace_id`; the strategy routes each +return to the adapter registered for that instance. + +```mermaid +flowchart TD + parsed["ParsedGraph and TraceRecord"] --> strategy["GraphIRReplayStrategy"] + strategy --> adapter["CreditDispatchAdapter per instance"] + strategy --> executor["TraceExecutor per run"] + executor --> ctx["_TraceContext"] + ctx --> store["VersionedChannelStore"] + executor --> scheduler["Scheduler adjacency"] + executor --> dispatch["dispatch/* singledispatch handlers"] + dispatch --> adapter + adapter --> issuer["CreditIssuer.issue_graph_credit"] + issuer --> router["StickyCreditRouter"] + router --> worker["worker request materialization"] + worker --> returns["CreditCallbackHandler graph observer"] + returns --> strategy + strategy --> adapter +``` + +## End-to-end execution lifecycle + +`TraceExecutor.run()` creates one `VersionedChannelStore` for a trace and drives +the frontier. If +`Environment.GRAPH.EXECUTOR_WATCHDOG_TIMEOUT` is set, the frontier driver is +wrapped in `asyncio.wait_for`; otherwise pre-dispatch deadlocks are not bounded +inside the executor. + +The frontier driver opens a per-trace `asyncio.TaskGroup`, schedules every entry +node returned by `Scheduler`, and then lets node tasks schedule successors as they +complete. + +```mermaid +sequenceDiagram + participant S as GraphIRReplayStrategy + participant E as TraceExecutor + participant C as VersionedChannelStore + participant D as dispatch handler + participant A as CreditDispatchAdapter + participant I as CreditIssuer + participant R as graph return observer + + S->>A: create and register adapter by instance id + S->>E: run(trace) + E->>C: create per-trace store + E->>E: schedule entry _fire tasks + E->>C: await_inputs requirements + E->>E: apply timing gate + E->>D: execute node body + D->>A: dispatch graph credit for LlmNode + A->>I: issue_graph_credit(turn) + R->>A: resolve credit return + A-->>D: complete parked Future + D-->>E: NodeExecutionResult writes + E->>C: write outputs and mark producers done + E->>E: schedule successors + E-->>S: TraceResult channel snapshot +``` + +## Scheduling model + +There is no central ready queue in the executor. Readiness is expressed through +channel waiters, futures, and `TaskGroup` task creation. + +`Scheduler` provides adjacency only: + +- `START` static successors become entry nodes. +- `END` is suppressed and never scheduled. +- Static successors fire after a node returns a result. Start-anchored successors + (`delay_after_predecessor_start_us`) are scheduled at the predecessor's DISPATCH + instead, and are excluded from the completion frontier (see below). + +Within a trace, concurrency is every scheduled node task whose channel inputs and +timing gate have cleared. `_schedule()` deduplicates fan-in by tracking scheduled +node ids. Rescheduling an already completed node is treated as a cycle bug. + +A node firing path is: + +1. Wait for declared `ChannelRequirement` values. +2. Capture the channel store sequence and compute a causal input snapshot. +3. Apply static and node-level timing gates. When the gate clears, stamp the + node's dispatch wall and schedule any start-anchored successors — they are + paced off this dispatch, not off the node's completion. +4. Execute the registered dispatch body. +5. Publish result writes to the channel store. +6. Mark all declared output producers done. +7. Schedule static successors. + +An `LlmNode` receives a causal global snapshot at the gate sequence +(`snapshot_at_seq`), so its prompt resolution sees every channel — including ones +it did not declare in `inputs` — as of the captured sequence. + +## Channel and producer semantics + +`VersionedChannelStore` is a per-trace versioned dataflow store. Node writes +append log entries with monotonically increasing `write_seq` values; the log is +never replaced, and an overwrite-typed channel rejects a second writer for the +trace lifetime with `OverwriteConflictError`. Initial +trace state is seeded at sequence zero and participates as the reducer seed, but +it does not count as a producer arrival for `count=N` requirements. + +A firing captures the first `N` non-initial writes needed for each requirement and +reduces them deterministically by `write_seq` and `writer_node_id`. + +Static producer counts are computed from every node's `write_channels` and seed +`count="all"` resolution. A node marks its declared output producers done when it +finishes (`mark_producer_done` in the `_fire` finally), whether or not it wrote; +waiters whose requested count is no longer reachable wake with an +`insufficient_producers_remaining` (producer completed without writing) or +`all_producers_cancelled` (producer failed/cancelled) orphan instead of sleeping +forever. + +## Dispatch by node kind + +Dispatch bodies are registered through `TraceExecutor._execute` using sibling +modules under `src/aiperf/graph/dispatch/`. `dispatch/__init__.py` imports the +modules for registration; the fallback handler in `executor.py` only runs if +registration failed, and the `singledispatch` default raises for any node type +that is not an `LlmNode`. + +`LlmNode` is the only dispatched node kind — every live producer lowers to +`LlmNode` + `StaticEdge`. Its body (`dispatch/llm.py`) builds a `DispatchRequest` +(node id) and awaits the injected credit issuer — for graph +replay the per-instance `CreditDispatchAdapter` — passing a per-dispatch +first-token stamp callback, then writes the result to `node.output`: a +type-correct empty list for messages-typed channels, the placeholder string +otherwise. It is also the only kind that issues a graph credit. + +## Per-request wire streaming mode + +Each graph credit resolves its wire streaming mode PER REQUEST from the recorded +per-node value, not from the global `--streaming` flag. When the worker +materializes a graph credit's body, it stamps `payload["stream"]` from the node +envelope's TOP-LEVEL `stream` key (`envelope["stream"]` in +`worker_materialize.py` — not `dispatch_overrides`), which the build plane +stamped from the recorded per-node mode (weka `"s"` → `True`, `"n"` → `False`; +dynamo derived from recorded `ttft_ms`). Only a payload whose envelope carries +NO recorded stream value (`stream` absent) falls back to the global +`endpoint.streaming`. That same +resolved value is carried onto `RequestInfo.stream_override`, and every non-graph +request leaves `stream_override=None` so it follows the global flag as before. +`stream_options.include_usage` keys on the FINAL stamped `stream`, so +server-token-count usage still rides whichever mode won. + +Resolution rule, in order of precedence: + +1. recorded per-node `stream` override (graph credits) — **wins**; +2. global `endpoint.streaming` (`--streaming`) — fallback for mode-less graph + payloads and the sole control for every non-graph run. + +On the wire, the transport's REQUEST side reads `RequestInfo.stream_override` +through `effective_streaming` (`base_transports.py`) to pick the `Accept` +header (`text/event-stream` vs `application/json`) and the streaming URL path per +request. The RESPONSE read side is already per-request and independent: it parses +SSE vs JSON by the SERVER's `text/event-stream` content type +(`aiohttp_client.py`), so the parse follows the actual response body regardless +of the request flag. The net effect is that a recorded-streaming source streams +(and emits its mid-flight `FirstToken`) even with the global flag OFF, while a +recorded `"n"` turn stays non-streaming — a single JSON body — inside an +otherwise-streaming run. + +> **Result metrics are gated PER RECORD, not by the global flag.** +> `STREAMING_ONLY` streaming metrics (TTFT / TTST / TTFO / ITL / ICL) are +> computed over the per-record subset of requests that actually streamed on the +> wire (`RequestRecord.streamed`), counted by the visible `streamed_request_count` +> aggregate (its denominator). The run-level gate in +> `BaseMetricsProcessor.get_filters` (inherited by `MetricRecordProcessor`; +> `src/aiperf/post_processors/base_metrics_processor.py:50-57`) only drops the +> whole streaming family when NOTHING in the run can stream — i.e. global +> `--streaming` off AND the input is not a graph workload; for a graph-IR replay +> the family stays enabled and the hidden `streamed_request` predicate excludes +> each non-streamed record individually. So a graph run replayed WITHOUT +> `--streaming` still reports TTFT / ITL / ICL for its recorded-streaming nodes, +> computed over those nodes alone, while its recorded `"n"` records are simply +> excluded (never dragging their full request latency in as a first-token time). +> `--streaming` is no longer required to see these metrics for a graph replay. + +## Timing and t-star behavior + +A node's firing gate is the maximum of: + +- incoming static edge `delay_after_predecessor_us` from the predecessor finish + time, +- incoming static edge `delay_after_predecessor_start_us` from the predecessor + DISPATCH time (see below), +- incoming static edge `delay_after_predecessor_first_token_us` from the + predecessor's OBSERVED FIRST-TOKEN time (post-TTFT anchoring; see below), +- incoming static edge `min_start_delay_us` from input readiness, +- node-level `min_start_delay_us`. + +| Edge field | Anchor point | Gate | Fallback | +| --- | --- | --- | --- | +| `delay_after_predecessor_us` | predecessor finish | `node_finish_wall_us[source] + delay` | — | +| `delay_after_predecessor_start_us` | predecessor dispatch | `node_dispatch_wall_us[source] + delay` | — | +| `delay_after_predecessor_first_token_us` | predecessor observed first token | `node_first_token_wall_us[source] + delay` | falls back to the start anchor when no first token is observed | + +`AIPERF_GRAPH_IGNORE_EDGE_DELAYS` bypasses all edge and node timing gates, +including the start-anchored gate. The per-executor `compress_edge_delays` flag +does the same for selected runtime instances, such as accelerated warmup. + +A start-anchored edge (`delay_after_predecessor_start_us`) paces a successor off +when its predecessor DISPATCHES rather than when it completes — the successor +does not wait for the predecessor's result. The runtime stamps +`_TraceContext.node_dispatch_wall_us[source]` the moment the predecessor's own +firing gate clears, then schedules the start-anchored successor right there, at +predecessor dispatch. The successor's gate is +`node_dispatch_wall_us[source] + delay_after_predecessor_start_us`. Because such +a successor is scheduled at dispatch time, it is deliberately excluded from the +completion-time `Scheduler.successors_after` frontier: a child that finishes +before its still-running parent would otherwise be re-scheduled into the cycle +guard. Start-anchored edges are likewise excluded from a node's fan-in channel +requirements (`with_fan_in_inputs`), so the successor gates only on the +dispatch-relative delay and never blocks on the predecessor's `_out` channel. A +single edge is either completion-anchored (`delay_after_predecessor_us`) or +start-anchored (`delay_after_predecessor_start_us`), never both; validator rule +54 (`_rule_54_edge_delay_exclusivity`) rejects an edge that sets both. + +Anchor kinds must also be uniform PER TARGET: a node with one start-anchored +in-edge plus one completion in-edge (a START edge counts as completion-kind) is +rejected at `Scheduler` construction with a `NotImplementedError` naming the +node and both edges. The runtime would otherwise fire the node at its start +anchor — silently ignoring the completion predecessor's recorded ordering — and +then re-schedule the completed node into the cycle guard when that predecessor +finishes (a spurious "cycle detected"). No shipped lowering emits the shape +(`apply_start_anchors` replaces an anchored node's whole in-edge set), so the +gate only affects hand-authored graphs and future adapters. + +A first-token-anchored edge REFINES a start-anchored one for a successor whose +recording began after the predecessor's first token arrived. Its +`delay_after_predecessor_first_token_us` (call it `D'`) paces the successor off +the predecessor's OBSERVED first token rather than off its dispatch. `D'` is only +valid alongside `delay_after_predecessor_start_us` (validator rule 55): the start +anchor is the mandatory fallback. `_apply_firing_delay` first awaits the source's +first-token latch (`ctx.first_token_event(source)`) so the observed wall — or its +settled ABSENCE — is known before the gate is computed. When the predecessor emits +a first token the runtime stamps `ctx.node_first_token_wall_us[source]` and gates +the successor at `first_token_wall + D'`, which SUPERSEDES the dispatch fallback +for that edge. When the predecessor TERMINATES WITHOUT a first token, +`_finalize_node` still sets the latch (with no wall entry), so the gate cleanly +falls back to the start anchor `node_dispatch_wall_us[source] + +delay_after_predecessor_start_us`. `AIPERF_GRAPH_IGNORE_EDGE_DELAYS` and +`compress_edge_delays` skip the first-token wait and the gate uniformly, exactly +as they skip the other anchors. + +> **First-token anchoring needs the SOURCE node to stream.** The observed first +> token exists only when the worker parses the source node's SSE stream. Each +> graph node now dispatches per its own recorded `streaming` mode (a per-request +> override on the wire), so a recorded-streaming source streams regardless of the +> global `--streaming` flag — the flag is only the fallback for mode-less +> payloads and non-graph runs. The residual failure is a first-token edge whose +> source `LlmNode` carries `streaming=False`: it emits no first-token event, so +> every one of its first-token-anchored successors SILENTLY degrades to its start +> anchor (completion-relative dispatch delay) — a fidelity loss with no other +> signal. This is possible only in hand-authored/degenerate graphs; recorded +> corpora are consistent by construction (the same recorded ttft drives both the +> edge refinement and the source node's streaming mode). The TimingManager emits +> a one-shot configure-time warning when a first-token source has +> `streaming=False` (`_advise_non_streaming_first_token_sources`); set +> `streaming=True` on the source node to restore post-TTFT anchoring. + +For Weka replay, `GraphIRReplayStrategy` can slice a trace at a sampled `t*`. +Arrivals before `t*` are warmup; arrivals at or after `t*` are profiled and +rebased to offsets relative to `t*` (zero only at the boundary). The strategy +uses absolute start offsets so surviving frontier turns are anchored to the +instance run start rather than to when their inputs happen to become ready. + +`--trajectory-start-min-ratio`, `--trajectory-start-max-ratio` (per-run config), +the run's `--random-seed`, and lane index determine t-star sampling. The +ratio ordering constraint is enforced by the graph conversation source rather +than by the environment field definitions. + +## Concurrency and backpressure layers + +Graph replay has several independent concurrency controls. They should not be +collapsed into one concept. + +| Layer | Owner | What it bounds | +| --- | --- | --- | +| Node tasks | `TraceExecutor` | In-trace dataflow tasks whose inputs are ready | +| Trace lanes | `GraphIRReplayStrategy` | Concurrent trace instances admitted for replay | +| Graph credit issue | `CreditIssuer` and stop checker | Whether another graph request may be sent | +| Prefill slots | `CreditIssuer` and callback handler | In-flight prefill pressure per sent request | +| Adapter waiters | `CreditDispatchAdapter` | Graph requests awaiting correlated worker returns | +| Router load | `StickyCreditRouter` | Worker choice and in-flight credit load | + +Cross-trace concurrency is resolved by `GraphIRReplayStrategy`: explicit override +wins, then positive phase `concurrency`, then `1` (the plain aiperf default; raise +cross-trace parallelism with `--concurrency`). Lanes recycle templates while a stop +condition exists and still allows new work. With no stop condition, a bare graph +run performs a single corpus pass. + +Graph credits bypass the normal linear session-slot lifecycle, but they still +check the DAG-child stop gate, acquire one prefill slot before sending, increment +normal sent counters, and route through the normal credit router. Prefill slots +release on first token, or on return when no first-token release was recorded. +Graph credits do not acquire or release session slots. + +## Credit bridge and return path + +`CreditDispatchAdapter.dispatch()` mints a graph correlation id and turn index, +parks one `Future` in `_waiters`, creates a `TurnToSend`, and issues it +immediately. The awaiting node +is bounded by `Environment.GRAPH.DISPATCH_TIMEOUT` once it reaches the adapter. +That timeout does not bound nodes waiting earlier on unsatisfied channel inputs. + +`CreditIssuer.issue_graph_credit()` places graph work onto the normal credit +path. A refused issue sets a `GraphDispatchError` on the parked future so the +executor can unwind the node rather than wait for a return that will never +arrive. + +The graph return observer is registered before any graph credit is issued. It +runs before gated phase-handler logic and routes returns by `credit.trace_id` to +the live adapter. Adapter resolution treats success, error, cancellation, and +context overflow as terminal outcomes for the parked future. Context overflow is +converted into expected early termination for the trajectory. + +Unknown return keys are dropped after a debug log. Unknown adapter instance ids +are also dropped; a parked dispatch can then only resolve by timeout or external +cancellation. + +## First-token fan-out (post-TTFT anchoring) + +A first-token-anchored edge needs the runtime to learn a predecessor's observed +first token WHILE that predecessor is still streaming. This rides a dedicated +`FirstToken` fan-out that parallels the return path but fires earlier. + +Only the nodes that SOURCE a first-token-anchored edge are opted in. The strategy +computes `first_token_sources(graph)` per trace (the set of such source node ids) +and the adapter stamps `first_token_event=True` on exactly those credits' +`TurnToSend`. The worker builds its first-token SSE callback only when prefill +concurrency limiting is active OR the credit carries `first_token_event`, so a +non-graph run without prefill-concurrency limiting — and a graph run with no +first-token edges — pays no per-chunk SSE parsing overhead. + +```mermaid +sequenceDiagram + participant W as Worker + participant R as Credit router + participant H as CreditCallbackHandler + participant S as GraphIRReplayStrategy + participant A as CreditDispatchAdapter + participant E as TraceExecutor / _TraceContext + W->>W: parse SSE, first meaningful chunk (streaming only) + W->>R: FirstToken(credit_id, phase, ttft_ns,
trace_id, x_correlation_id, turn_index) + R->>H: on_first_token(first_token) + H->>S: graph first-token observer (set_graph_first_token_observer) + S->>S: de-mux by trace_id -> owning adapter + S->>A: on_first_token(x_correlation_id, turn_index) + A->>E: stamp closure: node_first_token_wall_us[source] set,
first_token latch released + E->>E: successor gate = first_token_wall + D' +``` + +The worker emits `FirstToken` per credit carrying `credit_id`, `phase`, +`ttft_ns`, plus the graph routing keys `trace_id` / `x_correlation_id` / +`turn_index`. The router hands it to `CreditCallbackHandler.on_first_token`, +which also releases a prefill slot when one was held. The single graph +first-token observer (installed via `set_graph_first_token_observer` before any +credit issues) de-multiplexes by `trace_id` to the owning per-trace adapter, just +as the return observer does. The adapter's `on_first_token` looks up the parked +per-dispatch first-token callback by `(x_correlation_id, turn_index)` and invokes +the executor's stamp closure, which records `node_first_token_wall_us[source]` and +releases the source's first-token latch. A late token for an unknown trace id — +after the trace unwound, or a non-graph fast-path token — is a graceful no-op, and +the successor simply falls back to its start anchor. + +Because the observation depends on SSE parsing, the fan-out is inert only for a +successor whose SOURCE node is non-streaming. Each graph node dispatches per its +own recorded `streaming` mode (a per-request override), so a recorded-streaming +source emits its `FirstToken` — and anchors its successors — regardless of the +global `--streaming` flag; that flag is only the fallback for mode-less payloads +and non-graph runs (see [Per-request wire streaming +mode](#per-request-wire-streaming-mode)). Only when a source carries +`streaming=False` is no `FirstToken` emitted for its edges: the latch is then +released solely by `_finalize_node` at that predecessor's completion, and its +anchored successors degrade to their start-anchor fallback. The TimingManager +warns once at configure time per that source-node condition. + +## Failure, cancellation, and containment + +Failure handling is intentionally not uniform across node kinds: + +- `GraphDispatchError` is contained rather than re-raised. Segment-pool + graphs receive type-correct sentinel writes to normal output channels (`[]` + for messages-typed channels, `None` otherwise) so gate-only downstream + readers can continue. Other IRs may receive no writes, allowing downstream + channel waiters to orphan. +- Context overflow is treated as expected early termination. The node suppresses + successors, and downstream orphan cascades from missing overflowed outputs are + swallowed as clean exits. +- A graph dispatch timeout or cancellation drops the adapter waiter, so a late + return cannot resolve a dead dispatch. + +Without `AIPERF_GRAPH_EXECUTOR_WATCHDOG_TIMEOUT`, pre-dispatch dataflow +liveness bugs such as unsatisfied channel counts can hang until an external +cancellation or test timeout. `AIPERF_GRAPH_DISPATCH_TIMEOUT` only applies +after a node reaches the credit adapter. + +## Environment knobs + +The graph runtime reads the process-wide `Environment` singleton. Runtime code +reads live `Environment` attributes rather than reparsing `os.environ` on every +use, so tests and tools that mutate environment variables after import must also +reset or mutate the singleton state used by the code under test. + +Key graph runtime knobs: + +| Field | Default | Runtime effect | +| --- | --- | --- | +| `AIPERF_GRAPH_DISPATCH_TIMEOUT` | `300.0` | Bounds adapter waits after a graph credit is issued or deferred | +| `AIPERF_GRAPH_EXECUTOR_WATCHDOG_TIMEOUT` | unset | Optional wall-clock timeout for the executor frontier driver | +| `AIPERF_GRAPH_IGNORE_EDGE_DELAYS` | `False` | Collapses recorded edge and node delay gates globally | +| `AIPERF_GRAPH_WARMUP_MAX_OUTPUT_TOKENS` | `1` | Caps warmup materialization output tokens | +| `--agentic-cache-warmup-duration` (config, not env) | unset | Enables compressed-delay cache-pressure warmup | +| `AIPERF_GRAPH_IDLE_GAP_NO_DURATION_WARN_SECONDS` | `30.0` | Advisory threshold for faithful idle-gap replay without duration | +| `--trajectory-start-min-ratio` (config, not env) | `0.0` | Lower t-star sampling ratio; window off by default (full replay), `--scenario inferencex-agentx-mvp` auto-applies `0.0`/`1.0` | +| `--trajectory-start-max-ratio` (config, not env) | `0.0` | Upper t-star sampling ratio; `0.0` keeps the window off (full replay) | +| `--random-seed` (config, not env) | unset | Base seed for t-star sampling (shared with content synthesis) | +| `--burst-phase-starts` (config, not env) | `False` | Controls phase start burst behavior | +| `AIPERF_DATASET_WEKA_GRAPH_PARALLEL_THRESHOLD` | `8` | Dataset ingest parallelism threshold | +| `AIPERF_DATASET_WEKA_GRAPH_PARALLEL_WORKERS` | `0` | Explicit Weka graph ingest worker count | +| `AIPERF_DATASET_WEKA_GRAPH_PARALLEL_AUTO_MAX_WORKERS` | `16` | Auto worker upper bound | +| `AIPERF_DATASET_WEKA_GRAPH_PARALLEL_PREFETCH_MULTIPLIER` | `16` | Dataset prefetch sizing multiplier | + +Window-sizing rule: the submit window is `workers * multiplier` and must cover +the rows remaining behind the single heaviest trace, or fast workers stall +head-of-line while that one trace drains; the default `16` yields a window of +`256` at the auto 16 workers, which covers the 393-row weka corpus whose +heaviest row sits at index 140. Measured cost on that corpus: ~7.3 GiB parent +VmHWM (~17.5 GiB process tree) at multiplier 16 vs ~2.8 GiB parent at +multiplier 4. + +The former dataset flags selecting alternate trie store shapes (the +segment-trie IR gate, the mmap segment-store toggle, and the cross-run delta +cache toggle) were retired: the segment-trie IR with the interned unified store +is the sole trie build path, and there is no cross-run cache. + +## Static analysis helpers + +The `src/aiperf/graph/analysis/` helpers compute timeline, cohort, snapshot, +and trace-duration views over the same graph primitives. These outputs are used +for planning, t-star slicing, warmup/profiling partitioning, and diagnostics. +They are not a second runtime scheduler, and LLM cohorts should not be described +as synchronization barriers. The static elaboration follows BOTH anchor kinds — +completion successors and start-anchored (dispatch-time) successors — and keeps +a scheduled-but-unsatisfied node pending until later arrivals satisfy its +fan-in, so the dry-run's fired-node set matches the executor for the DAGs the +live adapters emit (start-anchored subtrees count toward `trace_duration_us` +and partition normally in `compute_snapshot`). + +Worker-side graph request materialization is stateless by node ordinal and trace +id for slot-less nodes: any worker can rebuild such a request from the shared +unified segment store (`GraphSegmentUnifiedClient`), with warmup token caps +applied during materialization. Slot-carrying nodes (dynamic content) +additionally read the per-worker dynamic pool and require per-trace sticky +routing. + +## Test coverage map + +Relevant unit tests exercise the runtime contract from multiple angles: + +- `tests/unit/graph/test_executor_runs_weka.py` covers executor scheduling, + fan-in, virtual time, and overflow behavior. +- `tests/unit/graph/test_executor_watchdog.py` covers watchdog-bounded + frontier deadlocks. +- `tests/unit/graph/test_credit_dispatch_adapter.py` covers adapter dispatch, + waiter, and return behavior. +- `tests/unit/graph/test_graph_return_bridge.py` covers graph return routing. +- `tests/unit/graph/test_lane_fanout_recycle.py` covers lane fan-out and + recycle behavior. +- `tests/unit/graph/test_tstar_activation.py` and warmup tests cover t-star, + profiling, and warmup variants. +- `tests/unit/graph/test_analysis_core.py` covers static analysis helpers. + +## Pitfalls for maintainers + +- Do not infer a central executor queue from the scheduling code. The executor + uses task creation, channel waiters, and futures. +- Do not document `AIPERF_GRAPH_DISPATCH_TIMEOUT` as a whole-executor + liveness timeout. It starts after adapter dispatch is reached. +- Do not describe graph credits as normal session-slot turns. They bypass session + slots while still using prefill slots and the credit router. +- Do not present analysis cohorts as runtime synchronization barriers. +- Be careful when documenting environment defaults: the docs and generated tables + describe field defaults, while runtime behavior depends on the singleton and + the strategy or executor code that consumes those fields. diff --git a/docs/reference/graph-ingest-build-pipeline.md b/docs/reference/graph-ingest-build-pipeline.md new file mode 100644 index 0000000000..a5452ecdb4 --- /dev/null +++ b/docs/reference/graph-ingest-build-pipeline.md @@ -0,0 +1,1093 @@ + + +# Graph Ingest and Build Pipeline (Internal Reference) + +Internal reference for the graph workload path (weka, dynamo, native, and +dag_jsonl): how a local JSON file, a local directory of JSON files, or a +HuggingFace corpus id becomes a `ParsedGraph`, how the segment-trie IR is +persisted into runtime stores, and how the DatasetManager, TimingManager, and +workers agree on node addressing. + +This page is build/runtime plumbing, not user-facing CLI guidance. Related +references: + +- [Graph Structural Sidecar Handoff](./graph-structural-handoff.md) +- [Graph Segment Unified Store](./graph-segment-unified-store.md) + +## Pipeline summary + +```mermaid +flowchart TD + input["--input-file (+ optional --graph-format)\nlocal .json / dir / HF org/name"] --> detect["resolve_graph_workload(run)\nmemoized GraphWorkloadRef (path + format)"] + detect --> build["GraphStoreBuilder._build_graph_store_streaming\ndispatch by format: two drains"] + build -->|"weka_trace"| items["work items (_WorkItem)\nfile paths / HF row dicts"] + build -->|"dynamo_trace"| storefirst["construct store BEFORE parse\nparse_graph_workload(run, path, direct_store=store)\nStoreBackedSegmentPool write-through (pool stays empty)"] + build -->|"native / dag_jsonl"| wholeparse["parse_graph_workload(run, path)\nregistry-dispatched parse(path, ctx), in-process\neager SegmentPool"] + items --> dispatch["_map_items (trace_parallel.py)\nserial at/below threshold\nbounded ordered pool above"] + dispatch --> payloads["iter_item_segment_payloads\nstreamed TraceSegmentPayloads"] + storefirst --> eagerdrain + wholeparse --> eagerdrain["build_unified_trie_store_interned\nin-process interned drain (same parse,\nslot-free and slot-carrying alike;\ndynamo pool already drained -> put loop no-ops)"] + eagerdrain --> unified + eagerdrain --> wsidecar["_write_graph_sidecar(parsed)\nsidecar direct from stripped parse\n(parse-order traces)"] + wsidecar --> sidecar + payloads --> drain["_build_graph_store_streaming_trie\nbuild_unified_trie_store_from_payloads"] + drain --> unified["GraphSegmentUnifiedBackingStore"] + drain --> structural["_merge_structural_graphs\nmerged content-free graphs (weka only)"] + structural --> sidecar["graph_meta sidecar\n+ per-node prefix-cache map"] + unified --> worker["worker materialize"] + sidecar -- "path via GraphSegmentClientMetadata on DatasetConfiguredNotification" --> tm["TimingManager schedule plane\nstructural ParsedGraph"] + tm --> strategy["GraphIRReplayStrategy\nnode ordinals"] + strategy --> worker +``` + +The Weka path has two coupled contracts: + +1. **Graph/topology contract:** every weka source — a local `.json` file, a + directory of `.json` files, or a weka-marked HuggingFace `org/name` id — + becomes an iterator of work items feeding ONE dispatch (`_map_items` in + `trace_parallel.py`). The build plane streams store payloads through + `stream_weka_trace_segment_payloads` and writes the mandatory structural + sidecar; the schedule plane loads that sidecar from the exact path the + graph-typed `DatasetConfiguredNotification` advertised and never re-parses + (a missing or unloadable sidecar is a hard configure-time failure). Every + build route resolves the same run-derived synthesis knobs into ONE + `GraphParseContext` (`resolve_graph_parse_context(run)`); the weka entry + functions thread them into each per-item parse via `_resolve_parse_kwargs`, + so the in-process parse and every spawn-started pool worker derive the same + `ParsedGraph` topology and node ordinals. +2. **Payload contract:** build persists each node's request body materialization + data into mmap-backed stores; workers later materialize by + `(trace_id, node_ordinal, phase_variant)`. + +## Workload detection and parser seam + +Entry points live in `src/aiperf/dataset/graph/workload_detect.py`. + +### Detection: one memoized resolution per process + +- `resolve_graph_workload(run)` is the single graph-ness accessor every + consumer calls (DatasetManager, TimingManager, scenario validator, + post-processors, ...). It returns a `GraphWorkloadRef` + (`src/aiperf/config/resolution/plan.py`) — the default dataset's input path + (kept verbatim for HF `org/name` ids, never `.resolve()`d) plus the adapter + format name — or `None` for a non-graph run. Detection runs AT MOST ONCE per + process: the result is memoized on `run.resolved.graph_workload` (with a + separate `graph_workload_resolved` marker, so "not a graph run" is + distinguishable from "never checked"), and the config resolver chain + (`DatasetResolver.resolve`) populates the memo eagerly in single-run mode so + child processes inherit it through the pickled run without re-walking the + registry. Any derivation failure degrades to `None`, exactly like the + per-consumer sniffs this accessor replaced. +- If `--graph-format` / `FileDataset.graph_format` is set, the input is forced + into graph mode for that adapter, including `native`. +- Without an override, detection walks registered `graph_adapter` plugins and + excludes `native` and `dag_jsonl` (`_AUTODETECT_EXCLUDED`), so plain JSONL/YAML + conversation datasets — and legacy `--custom-dataset-type dag_jsonl` runs — are + not hijacked into graph mode (the `dag_jsonl` graph adapter is opt-in via + `--graph-format dag_jsonl` only). +- `is_graph_workload_path(path)` is the path-level companion (same registry + detection, same `native`/`dag_jsonl` exclusion) for callers that have no run. +- The schedule plane composes its own veto on top of the accessor: + `timing/config.py` routes the warmup and profiling phases to + `TimingMode.GRAPH_IR` iff `resolve_graph_workload(run)` is non-`None` AND + `_explicit_non_graph_format(run)` is `False` — a user who pinned a + graph-incompatible `--custom-dataset-type` is never silently rerouted to the + graph pipeline for timing purposes. + +### Parse dispatch: registry-driven, one context + +`parse_graph_workload(run, path)` is the shared ingest seam for in-process +graph parsing. It is a build-plane-only seam — the DatasetManager (via +`GraphStoreBuilder`) is the only production caller; the TimingManager never +parses. There is no per-format call ladder: every non-native format goes +through the adapter registry; native is decoded directly by the parser +(byte-identical to the registered `NativeGraphAdapter`). + +1. publish the run's tokenizer trust/revision to the loader-preload env + (`publish_graph_loader_tokenizer_env`); +2. read the memoized `GraphWorkloadRef` for the format name (asserting the + passed path matches the run's own resolved input); +3. resolve every run-derived parse knob into ONE `GraphParseContext` via + `resolve_graph_parse_context(run)`; +4. dispatch `parse_graph(path, format=fmt, ctx=ctx)` (`parser.py`), which + short-circuits native to a direct decode and routes every non-native + format to the selected adapter's uniform `parse(path, ctx)`; +5. run the t*/dynamic-slot gate (`_gate_dynamic_slots_vs_tstar`) uniformly on + the parsed result (see below). + +Adapter failures surface uniformly as `GraphParseError`: the parser re-wraps +adapter-specific `ValueError` subclasses with the message text preserved, so +callers need exactly one except class. A registry-dispatched +`adapter.parse(path, ctx)` is byte-identical to the run's own parse for every +format (pinned by parity tests). + +`src/aiperf/dataset/graph/parser.py` provides the generic parser: + +- native YAML/JSONL graph files are decoded directly, assembled, start/end edges + are auto-injected, and `auto_derive` runs; +- non-native inputs dispatch through the graph adapter registry to + `adapter_cls.parse(path, ctx)`; the ctx is passed opaquely and each adapter + maps only the fields it consumes. + +### `GraphParseContext` + +`src/aiperf/dataset/graph/parse_context.py`. A frozen dataclass carrying every +run-config-derived knob an adapter needs to parse byte-identically to the run. +`resolve_graph_parse_context(run)` populates every field with the run's +RESOLVED value — each resolver is a pure function of the run config, so +resolving dag knobs for a weka run (and vice versa) is harmless: + +| Field | Consumers | Source | +|-------|-----------|--------| +| `content_root_seed` | weka, dynamo | `resolve_graph_content_seed(run)` (`--random-seed`) | +| `content_tokenizer` | weka, dynamo | `resolve_graph_content_tokenizer(run)` | +| `tokenizer_trust_remote_code` / `tokenizer_revision` | weka, dynamo | run tokenizer config; published to the loader-preload env via `publish_ctx_tokenizer_env` (only when trust is set, so a ctx-less parse never clobbers the run path's publish) | +| `prompt_corpus` | weka, dynamo | `synthesis.corpus` / `--prompt-corpus` (default `coding`) | +| `max_osl` | weka | `synthesis.max_osl` — caps each top-level chain request's wire `max_tokens` to `min(recorded out, max_osl)`; subagent-body chains stay uncapped | +| `num_dataset_entries` | weka, dynamo | `--num-dataset-entries` — cap on distinct traces / session-trees selected (filter-then-cap) | +| `max_context_length` | weka, dynamo | `--max-context-length` — per-trace context ceiling (input+output tokens) for graph-plane selection | +| `idle_gap_cap_seconds` | weka, dynamo | `synthesis.idle_gap_cap_seconds` (tri-state, below) | +| `trajectory_start_max_ratio` | parse-time t\* gate | `--trajectory-start-max-ratio` (scenario-auto-applied when unset); `0.0` = window off | +| `default_model` | dag_jsonl | primary model name | +| `run_streaming` | dag_jsonl | resolved endpoint `stream` flag | +| `delay_cap_seconds` | dag_jsonl | legacy `inter_turn_delay_cap_seconds` | +| `endpoint_extra` | dag_jsonl | `--extra-inputs` pairs | + +**Forward-only-when-set.** Adapters forward a ctx field to their entry function +ONLY when it is set, so a ctx-less parse (`parse(path)` — CLI tooling, direct +callers with no run config) stays byte-equal to the protocol-default entry, and +a partial ctx never clobbers a non-`None` entry default (dynamo's +`prompt_corpus="coding"`, dag_jsonl's `run_streaming=True`). + +**Tri-state idle-gap cap.** `idle_gap_cap_seconds` distinguishes three states: +`UNSET` (module sentinel; adapter default, 60 s), a float (that cap), and an +explicit `None` (warping DISABLED — the user's +`synthesis.idle_gap_cap_seconds: null`). `resolve_graph_parse_context` never +yields `UNSET` (a run always has a resolved answer), and the weka and dynamo +adapters forward the tri-state verbatim (`is not UNSET`, not `is not None`). + +The dynamo adapter has no behavior env knob left on its `parse` entry: +generation is always pinned to the recorded `output_tokens` (weka parity) and +recorded delays always replay, subject to the tri-state idle-gap cap; the only +remaining `Environment.DYNAMO.*` knob is `MAX_SUBAGENT_DEPTH` (the parent-link +chain-depth cycle guard). Seed semantics are unchanged: both +weka and dynamo pin `content_root_seed` at entry via +`resolve_effective_root_seed` (explicit seed → ambient bootstrap root seed → +per-run OS entropy), so unseeded content is internally consistent within one +run and distinct unseeded runs deliberately differ. + +### t*/dynamic-slot gate + +`_gate_dynamic_slots_vs_tstar(parsed, ctx.trajectory_start_max_ratio)` runs +uniformly on every `parse_graph_workload` result: a graph carrying dynamic +content slots (`graph_carries_assembly_slots` — prompt refs to LlmNode-written +channels) is rejected while the t* snapshot window is engaged +(`--trajectory-start-max-ratio > 0` — off by default; applied by +`--scenario inferencex-agentx-mvp` or explicit +`--trajectory-start-min/max-ratio` flags), because a slot producer +chopped into warmup would leave its consumer's pool value undefined. + +**Explicit-zero carve-out.** The gate is skipped iff EVERY node's +`arrival_offset_us` is explicitly the int `0`; `None` (the un-stamped default +on natively authored nodes) does NOT qualify and keeps gating. All-zero offsets +make the recorded duration 0, so any sampled t* is 0 and the snapshot chop is a +structural no-op — rejecting could only ever false-positive. dag_jsonl lowering +stamps `arrival_offset_us=0` on every node, and `DagJsonlGraphAdapter.parse` +enforces that invariant at its own seam (`_assert_dag_zero_arrival_offsets`); +if dag ever emits recorded offsets, that guard raises AND the carve-out stops +matching, so the gate re-engages by construction. The carve-out is keyed on the +invariant, not on a format branch: a NATIVE graph authored with all-zero +offsets now passes too (an intentional contract delta — the t*-degeneracy +argument is identical there, so the old rejection was the same false +positive). + +## Weka input forms + +Implementation: `src/aiperf/dataset/graph/adapters/weka/trace.py`. + +All three forms reduce to the same thing: an iterator of `_WorkItem` values +(a trace file path OR an in-memory row dict) fed to the unified dispatch in +`trace_parallel.py` — see +[Unified parse dispatch](#unified-parse-dispatch-and-seed-determinism). + +### Local file + +A single Weka trace is a JSON object with top-level discriminator keys including +`id`, `models`, `block_size`, `hash_id_scope`, and `requests`. Detection is a +bounded sniff: it reads roughly the first 4 KiB and only falls back to full-file +parse if the signature keys are already present. The adapter rejects foreign +objects with extra top-level keys rather than silently accepting objects that only +look Weka-like. + +During parse, the adapter: + +1. loads JSON with `orjson`; +2. validates into `WekaTrace`; +3. rejects an unrecognized `hash_id_scope` (anything other than `local` / + `global`) before Pydantic turns it into a generic schema error; +4. rejects empty traces; +5. builds the segment-trie graph through `build_trie_graph`; +6. attaches one `TraceRecord(id=weka_trace.id, tags=["from-weka-trace"])`; +7. returns a `ParsedGraph` whose `segment_pool` carries the deduplicated prompt + segments used by the build plane. + +A nonexistent local path surfaces the file-read error (`FileNotFoundError`) +directly; the "did you mean a HuggingFace dataset id?" hint appears only for +weka-marked `org/name` ids (see below). + +### Local directory + +A directory is detected as Weka when its lexicographically first `*.json` child +matches the same local-file sniff. Parsing wraps each `*.json` child (sorted by +name) as a file work item and feeds the unified dispatch; each item parses +through the same `_parse_single_file` / `_parse_trace_dict` core used by the +single-file and HF paths. + +Worker count defaults to `AIPERF_DATASET_WEKA_GRAPH_PARALLEL_WORKERS=0`, meaning +auto-size. The auto cap is controlled by +`AIPERF_DATASET_WEKA_GRAPH_PARALLEL_AUTO_MAX_WORKERS`; the switch to the +multiprocess path is controlled by +`AIPERF_DATASET_WEKA_GRAPH_PARALLEL_THRESHOLD`. + +### HuggingFace corpus id + +A non-existing `org/name` argument is treated as a Weka HF dataset id only when: + +- it has exactly one slash and matches the repo-id shape; +- it has no graph/file suffix such as `.json`, `.jsonl`, `.yaml`, or `.yml`; +- it contains the case-insensitive marker `weka`; +- it does not already exist as a filesystem path; +- its leading path component does not exist as a local directory + (`_looks_like_hf_dataset_id`, `return not p.parent.exists()`), so a typo'd + relative path like `traces/weka-061526` under an existing `traces/` dir stays a + local-path error instead of being routed to HuggingFace. + +The HF path uses `datasets.load_dataset(..., streaming=True)` and never writes a +temporary JSON file. The row iterator is controlled by: + +- `AIPERF_DATASET_WEKA_HF_SPLIT` (default `train`; slice syntax such as + `train[:100]` bounds the streamed rows); +- `AIPERF_DATASET_WEKA_HF_REVISION` (optional branch/tag/SHA pin). + +Each streamed row is shallow-copied to a plain dict and wrapped as a row work +item (`org/name#index` source label) for the same unified dispatch the file and +directory forms use. When an HF load fails, the error presents BOTH +interpretations (typo'd local path vs. inaccessible HF repo), because the +weka-marker heuristic can fire on a misspelled local path shaped like +`org/name`. + +### Unified parse dispatch and seed determinism + +Implementation: `src/aiperf/dataset/graph/adapters/weka/trace_parallel.py`. + +`_map_items` is the ONE serial-or-pool dispatch every weka route funnels +through. It prefetches at most `threshold + 1` items (lazy sources are never +fully consumed up front): at or below +`AIPERF_DATASET_WEKA_GRAPH_PARALLEL_THRESHOLD` (default 8) items, every item +parses serially in-process (no pool, no codec round-trip); above it, items +stream through a bounded, ordered pool window (forkserver on Linux, spawn on +macOS; `_loader_pool_context`) via +`_run_pool_streaming` (shared-memory corpus, per-item timeout, graceful +shutdown). Pool workers parse via the shared core (`_parse_single_file` for +file items, `_parse_trace_dict` for row items) — they never re-enter the +public `from_weka_trace` dispatcher. What crosses the pool boundary depends on +the consumer: the merged consumer's workers (`_parse_item_to_msgpack`) return +msgpack (msgspec codec) `ParsedGraph` blobs, while the streaming consumer's +workers (`_parse_item_to_segment_payloads`) return plain pickled +`list[TraceSegmentPayload]`. Either way no `ParsedGraph` instance crosses the +pool boundary, side-stepping the historically broken cross-process pickling of +`ParsedGraph` instances. + +The dispatch has two consumers: + +- `parse_items` merges the per-item results into ONE multi-graph `ParsedGraph` + (traces sorted by trace id, byte-deterministic across worker counts). This + backs `from_weka_trace` — the whole-graph weka adapter API used by direct + callers and tests (the DatasetManager's weka store build uses the streaming + consumer below, not this one). +- `iter_item_segment_payloads` streams per-trace `TraceSegmentPayload`s. This + backs `stream_weka_trace_segment_payloads` — the build plane serializes each + trace's envelopes into the unified store and drops the payloads before the + next arrive, so resident memory stays at ~one trace regardless of corpus + size. + +Both consumers pass the same `_resolve_parse_kwargs` dict, which pins the +content seed to a concrete int via `resolve_effective_root_seed` +(`shared/content.py`): an explicit `--random-seed` wins; otherwise the ambient +bootstrap-seeded manager's root seed; otherwise a fresh OS-entropy seed +(`secrets.randbits(64)`) generated once per resolution. Within ONE run the +serial path and every pool worker therefore synthesize identical bytes at any +threshold, while distinct unseeded runs deliberately differ — there is no +hardcoded fallback seed. The schedule plane never re-parses — it ingests the +structural sidecar this build wrote and broadcast — so no content bytes ever +need to agree across the build/schedule split; the only determinism that +matters is between the in-process parse and its spawn-started pool workers. + +## `ParsedGraph` shape from the segment-trie builder + +Implementation: the Weka-specific flatten/walk lives in +`src/aiperf/dataset/graph/adapters/weka/trie_build.py` (`build_trie_graph`). +The shared, format-agnostic trie core it calls into lives under +`src/aiperf/dataset/graph/segment_ir/` (`trie_content.py`, +`interval_order.py`, `pool.py`, `envelope.py`, +`store_builder.py`). The block-geometry, idle-gap-warp, and message-chaining +helpers now live inside `trie_content.py`. The t* / frontier snapshot chop +(`chop_trie_at_tstar` / `chop_trie_at_frontier`) has moved next to its only +consumer at `src/aiperf/timing/snapshot_chop.py`. +`weka/trie_build.py` re-exports only `ReconCallbacks` and `build_trie_graph`; the +trie primitives are imported directly from `segment_ir/` (there are no +back-compat aliases). + +`build_trie_graph(trace, ...)` returns `(ParsedGraph, SegmentPool)`. The emitted +IR is intentionally small: + +- one `LlmNode` per recorded normal (`type: "n"`) or streaming (`type: "s"`) + request, including recursive subagent-inner requests; +- `StaticEdge` waits-for dependencies only; +- no `SubgraphNode`, `SpawnNode`, `AwaitNode`, reducer, channel topology, or + chain/aux classification on the trie path; +- one top-level graph, no subgraphs; +- a `SegmentPool` containing prompt segments and assistant response segments. + +### Content lineage + +Every leaf request is flattened in recorded order. Its content parent is chosen +from prior requests by Weka `hash_ids` (`resolve_content_parents` in +`segment_ir/trie_content.py`, an incremental prefix-trie pass, byte-for-byte equal +to the pairwise scan but without the O(n²·m) double loop): + +1. prefer the earlier request whose `hash_ids` are the longest full prefix of the + current request; +2. tie-break full-prefix matches toward the most recent prior request; +3. if there is no full prefix, use the prior request with the longest partial LCP + as the branch point; +4. with no overlap, the request is a fresh root. + +The content parent no longer drives an incremental prompt "advance"; it now +serves two narrower roles: **role-boundary inheritance** (below) and reconstruction +statistics (LCP coverage). Timing dependency is derived independently — see +[Timing and dependency edges](#timing-and-dependency-edges). + +**Per-block frozen tags.** A single global pass (`compute_asst_caps` + +`assign_block_tags`) assigns every covered block a `(role, starts_new_message)` +tag. A block's tag is **frozen at the first node that creates it** and inherited +verbatim by every later node that shares it — never relabeled or coalesced. This +is the cache-safety invariant: any two requests sharing a block-aligned prefix see +the identical role segmentation over that prefix, so they render an identical +leading message chain. Assistant caps come from the recorded `out` of the +content-parent chain; the trailing user turn is guaranteed at block-creation time +(`asst -= 1` when a node would otherwise end assistant-only). + +**Message-unit emission.** `assemble_messages` groups the frozen per-block tags +into messages (a new message begins at each `starts_new_message`) and emits one +content-addressed pool entry per message via the existing +`segment_id(parent_id, role, tokens)` (blake2b-16), chained root→tip. Because tags +are frozen per trie position, a shared block prefix yields an identical +`(parent_id, role, tokens)` chain → identical segment ids → a real prefix-cache +hit. Emission is block-aligned: no partial tails, no synthesis of missing whole +blocks. A hard build-time gate (`assert_covered_isl` → `TrieISLMismatchError`) asserts each +node's reconstructed prompt equals the block-aligned **covered-count** +`min(len(hash_ids), in // block_size) * block_size` (covered-count, not +`(in // bs) * bs`, so a request storing fewer hash blocks than `in // bs` does not +abort the build). + +Non-block-aligned inputs are format-normalized before the trie: dynamo records +one trailing partial-tail hash spanning the remainder +(`(n-1)*bs < input_length <= n*bs`), which its lowering DROPS — engines +cache/share full blocks only, so +a partial block is never a prefix-cache hit — while `input_length` keeps the +tail tokens. Both recorded adapters then sample the sub-block remainder from +the same trace-scoped node seed, and both pass `small_prompt_fallback=True` so +a covered-count-0 (sub-block) recorded prompt synthesizes a single sampled +user message instead of an unreplayable empty prompt. The same recording +therefore lowers identically from either format. + +Each `LlmNode` receives: + +- `prompt`, materialized from the message-unit segment path for in-memory graph use; +- `metadata["trie"]["prompt_segment_ids"]`, the persisted per-message id chain used + by the worker (the sole trie metadata key; the assistant response segment is + reachable as the assistant pool entry chained onto the prompt tip); +- the native dispatch fields: `model` and `max_tokens` from the recorded request, + `streaming` from the request type; +- `arrival_offset_us` on the idle-gap-warped timeline. + +The worker ultimately uses `prompt_segment_ids`, not predecessor channel values, +to build the request body. + +The recorded `stream` override now REACHES THE WIRE per credit: weka lowering +sets `dispatch_overrides["stream"]` from the request type (`type: "s"` → `True`, +`type: "n"` → `False`) and the dynamo adapter derives it from recorded `ttft_ms` +(present → streaming). When the worker materializes a graph credit's body, +`apply_run_level_payload_options` stamps `payload["stream"]` from that recorded +value and carries it onto `RequestInfo.stream_override`, so the recorded per-node +mode WINS over the global `--streaming` flag (which is now only the fallback for +mode-less payloads and the sole control for non-graph runs). A recorded `"s"` +parent therefore streams — and emits its mid-flight first token — even without +`--streaming`, while a recorded `"n"` turn stays a single non-streaming JSON body +inside an otherwise-streaming run. `STREAMING_ONLY` result metrics +(TTFT / TTST / TTFO / ITL / ICL) are gated PER RECORD, not by the global flag: +they are computed over exactly the requests that streamed on the wire +(`RequestRecord.streamed`), counted by the visible `streamed_request_count` +aggregate. The run-level gate (`base_metrics_processor.py:50-57`) drops the whole +streaming family only when nothing can stream (global `--streaming` off AND not a +graph workload); for a graph replay the family stays enabled and each non-streamed +record is excluded individually by the hidden `streamed_request` predicate. So a +weka replay reports TTFT / ITL / ICL for its recorded `"s"` nodes even without +`--streaming`, over those nodes alone — the recorded `"n"` records are excluded +rather than reporting their full request latency as a first-token time. See +[per-request wire streaming mode](./graph-async-dataflow-runtime.md#per-request-wire-streaming-mode) +for the transport-side resolution. + +### Timing and dependency edges + +Timing dependency is derived independently of content ancestry, by an +**interval-order** pass over the flattened nodes (`build_interval_edges` in +`segment_ir/interval_order.py`, consuming the time-consistent `rank` from +`compute_ranks`). It replaces the older +`chain_prev` / subagent `spawner` / joined-leaf candidate set — those fields are +gone. The pass runs over ONE node set: global for weka/native/dag_jsonl, but per +SESSION-TREE for dynamo (a root plus its `parent_trajectory_id` descendants is +lowered independently, so edges stay within a tree and independent trees never +gain a cross-parent edge). + +- **Rank** is a linear extension of finished-before: nodes sorted by + `(warped_start, warped_end, node_id)`. Idle-warp monotonicity makes this a valid + topological order of the causality DAG. +- **Edge rule:** `A → B` iff `A` finished before `B` started on the raw recorded + clock (`A.raw_end <= B.raw_start`, i.e. `B.request.t`) **and** `rank(A) < rank(B)`. A request that + fully finished before another fully started must precede it; anything mid-flight + at that instant is off the hook (concurrent), so genuine racers stay parallel. +- **Async exclusion** (`_excluded_async`) drops fire-and-forget (`async_launched`) + subtree children from the candidate set before the frontier filter, so a + detached background agent never becomes a timing prerequisite of later work. +- **Frontier (transitive reduction):** only the maximal finished-before candidates + are kept — a candidate `c` transitively covered by another candidate `d` + (`c → d → B`) is dropped. So a node's edges are its immediate causes, not its + whole ancestry. (A content parent is therefore not automatically a timing + predecessor; it is dropped when a later cause dominates it.) +- **Binding-cause delay:** the latest-ending frontier predecessor (`max by .end`) + carries the warped end-to-start delay; every other frontier predecessor is a + zero-delay AND-join wait. +- **Empty frontier:** the node roots at `START` with `min_start_delay_us` equal to + its own warped arrival offset, preserving recorded concurrency. + +**Causal-overlap carve-out (start-anchored edges).** The finished-before rule +cannot express a request that BEGINS while its causal predecessor is still in +flight — a subagent spawned mid-turn, or a chain request that overlaps the +previous one. Interval-order alone would bind such a node to whatever unrelated +request happened to finish just before its start, embedding the causal parent's +concurrent server time as fake idle "think time." To fix this, the walk stamps +each node's `TrieNode.causal_parent_id` — the spawner for a subagent's first +inner request, else the previous n/s request in the same list (the dynamo +adapter stamps the same field: the previous turn in a chain, else the latest +parent-session turn at or before the child's start). After `build_interval_edges`, +`apply_start_anchors` (in `segment_ir/interval_order.py`) tests each node against +that parent: when the node's recorded start falls inside the parent's recorded +interval (`parent.raw_start <= node.raw_start < parent.raw_end`), it REPLACES the +node's interval-order edges with ONE start-anchored `StaticEdge(parent → node, +delay_after_predecessor_start_us=D)`, where `D` is the warped start-to-start gap. +At replay the successor is scheduled at the parent's DISPATCH and gated `D` later, +so recorded mid-flight concurrency tracks the parent causally instead of freezing +to the wall clock — see +[edge-gate semantics](./graph-async-dataflow-runtime.md#timing-and-t-star-behavior). +Nodes whose causal parent had already finished keep their interval-order edges +(end-anchoring is correct there by construction). + +On the SemiAnalysis 062126 corpus (393 traces, 98,827 requests) the carve-out +fires for 6,981 requests: 33.7% of subagent spawns (572 of 1,697) begin while +their spawner is still in flight, and 6.5% of all requests (6,409) overlap their +chain predecessor. Before it, interval-order bound each overlapped spawn to an +unrelated earlier request whose binding end-to-start delay was p50 97% concurrent +server time (the parent's own in-flight processing) rather than real idle time. +After it, all 6,981 overlapped nodes anchor to their stamped causal parent +(verified zero mismatches). + +**Post-TTFT refinement (first-token-anchored edges).** A start anchor still +overcounts when the child began after the parent's FIRST TOKEN rather than at the +parent's dispatch: the child's true "think time" starts once the parent began +streaming, not when it was sent. When the parent is a streaming request that +carries a recorded time-to-first-token (`parent.request.ttft`, in seconds — weka +lowering stamps it only for `type: "s"` requests; the dynamo adapter stamps +`ttft = ttft_ms / 1000.0`) AND the child's start falls at or after that first +token (`node.raw_start - parent.raw_start >= ttft`), `apply_start_anchors` +additionally carries `delay_after_predecessor_first_token_us = D' = max(0, +D - ttft*1e6)` on the same edge (`D` is the warped start-to-start gap). At replay +the runtime gates the child at the parent's OBSERVED first token `+ D'`, +superseding the dispatch anchor; the start anchor remains the mandatory fallback +(validator rule 55) for when the parent terminates without a first token. A +non-streaming parent (`ttft is None`) and a child that began BEFORE the parent's +first token keep the pure dispatch/start anchor. + +Of the 6,981 anchored nodes on the 062126 corpus, 1,097 gain a first-token +refinement (streaming parent, child begins post-TTFT → get `D'`); 908 began +pre-TTFT (streaming parent, child starts before the first token → no refinement); +and 4,976 have a non-streaming parent (no recorded `ttft`) — the latter two groups +stay purely dispatch-anchored. + +> **Post-TTFT anchoring needs the SOURCE node to stream.** `D'` is only stamped +> when the parent is a streaming request (recorded `ttft`), so in a built corpus +> every first-token edge's source node is itself `streaming=True` — the runtime +> observes that parent's first token because each graph node now dispatches per +> its own recorded `streaming` mode (a per-request override), independent of the +> global `--streaming` flag (see the +> [runtime first-token fan-out](./graph-async-dataflow-runtime.md#first-token-fan-out-post-ttft-anchoring)). +> The only degradation path is a hand-authored/degenerate graph whose first-token +> source carries `streaming=False`; the TimingManager warns once at configure time +> for exactly that case. + +The per-node frontier filter is `O(candidates²)` (up to `Θ(n²)` per node for a +pathological wide fan-in); accepted deliberately (wide fan-in is rare; no synthetic +barrier/collapse node is inserted). A sweep-line optimization is a deferred +follow-up. + +The idle-gap warp caps only true inactive gaps: it builds active intervals from +all flattened requests, collapses idle stretches longer than the cap, and never +cuts inside a request's `api_time` or overlapping subagent activity. This keeps +request durations and overlap relationships intact while preventing multi-hour +recorded dead air from parking warmup indefinitely. + +### t* snapshot chop + +`chop_trie_at_tstar(graph, t_star_us)` (in `src/aiperf/timing/snapshot_chop.py`, +next to its only consumer, not the weka trie build module) trims pre-`t*` nodes +for resume. Surviving +nodes keep their full `prompt_segment_ids` path, because pre-`t*` turns were +warmed and the server should already hold their KV. Nodes whose predecessors were +chopped are re-rooted from `START` with a t*-relative absolute offset, and input +requirements for dropped predecessor output channels are removed. + +## Build plane: DatasetManager routing and `GraphStoreBuilder` + +Implementation: `src/aiperf/dataset/graph/store_build.py` (the build itself), +`src/aiperf/dataset/dataset_manager.py` (thin routing), and +`src/aiperf/dataset/graph/segment_ir/store_builder.py` (the drains' store +primitives). + +`DatasetManager._configure_dataset` routes a graph run +(`resolve_graph_workload(run)` non-`None`) to `_configure_graph_dataset`, whose +build seam is `_build_graph_store`: first the endpoint gate +(`workload_detect.validate_graph_endpoint_type` — the graph dispatch path emits +a chat-completions body verbatim, so a non-chat endpoint would 422 on every +request; rejected at configure time, before any store work), then +`GraphStoreBuilder(run).build(graph_path)`. `_configure_graph_workload` is the +facet-only wrapper over the same seam for direct callers. + +`GraphStoreBuilder.build` owns the whole build: + +1. publishes the run's tokenizer trust/revision to the loader-preload env and + prestarts the trace-loader forkserver on the loop at a known-quiet point + (a lazily started helper inside the offloaded parse would race its + process-wide stdio swap against live logging); +2. resolves the mmap base path from `AIPERF_DATASET_MMAP_BASE_PATH` or the + system temp directory; +3. reads the memoized `GraphWorkloadRef` for the format and calls + `_build_graph_store_streaming` — the store build for every graph workload, + which dispatches by format to one of two drains: `weka_trace` streams + worker-pool payloads and merges their structural graphs; every other format + (dynamo / native / dag_jsonl) parses ONCE in-process and drains that whole + parse through the interned builder (`build_unified_trie_store_interned`), + which persists dynamic-slot envelopes too — so slot-carrying graphs + (live-reply dag_jsonl lineage, `@channel` native trie assembly items/capture) + ride this route with no separate fallback; +4. both drains finalize the unified store and write their own mandatory + structural sidecar (a catalog-divergent sidecar raises `DatasetError`; an + unwritable path fails the build with the underlying I/O error — the run is + unschedulable without it; a drain that completes without recording a sidecar + path is a hard build failure); +5. derives the per-node prefix-cache map (`_build_graph_prefix_cache_by_trace`) + from the full in-process parse (dynamo / native / dag_jsonl) or, on the weka + payload-stream drain, from the merged structural graphs; +6. returns a `GraphStoreBuildResult` — the `GraphDatasetMetadata` facet (the + trace universe `trace_ids` plus the per-node prefix-cache map, no + conversations), the sidecar path, and the store base path. + +`DatasetManager._configure_graph_dataset` then broadcasts +`DatasetMetadata(conversations=[], graph=...)` and a +`GraphSegmentClientMetadata` (unified-store base path, benchmark id, and the +exact sidecar path from the build result) on the +`DatasetConfiguredNotification`. There are no stub conversations, no +conversation mmap store, and no `inputs.json`: workers read the real graph +request bodies from the unified store by `(trace_id, node_ordinal)`, and the +TimingManager plans from the advertised sidecar. + +### Build drains (dispatched by format) + +`GraphStoreBuilder._build_graph_store_streaming` builds the unified store for +EVERY graph workload, dispatching on format to one of two drains: + +- **weka** (HF corpus id, local file, or local directory): the builder + resolves the run's parse knobs ONCE via `resolve_graph_parse_context` and + spreads its fields (seed, tokenizer, corpus, `max_osl`, idle-gap cap — the + tri-state cap forwards AS-IS, an explicit `None` means warping disabled) + verbatim into `stream_weka_trace_segment_payloads`, then streams the + worker-built `TraceSegmentPayload` values (`iter_item_segment_payloads`) into + `_build_graph_store_streaming_trie`, so the per-item parse is parallelized and + the parent never holds the whole corpus's real content. The rest of this + subsection describes that payload contract. +- **dynamo / native / dag_jsonl** (dag_jsonl is opt-in via `--graph-format + dag_jsonl`; an undetected `fmt=None` fails inside `parse_graph`'s + own detection): the parse is a single in-process lowering (a dynamo capture + groups its sessions into independent SESSION-TREES — a root plus its + `parent_trajectory_id` descendants — each lowered on its own node set so + interval-order edges stay WITHIN a tree and cross-parent edges never form; + the live write-through store pins this route to the serial tree-by-tree + build — the fused-parallel tree build, tuned by + `AIPERF_DATASET_DYNAMO_GRAPH_PARALLEL_*`, engages only for direct + callers/tooling with no `direct_store` — and dag_jsonl expands whole + conversation trees, so the store build has no per-item decomposition to fan + out). + `parse_graph_workload` runs once in-process (off-loop), and that SAME parse is + drained directly through `build_unified_trie_store_interned` — no payload + round trip, because in-process there is no worker pool to fan out to. Every + adapter lowers through the shared segment-trie core at parse time, so `parsed` + carries a `segment_pool` plus per-node `prompt_segment_ids`; a pool-less parse + raises (it is a lowering bug). **`dynamo_trace` additionally takes the DIRECT + write-through route:** the store is constructed BEFORE the parse and threaded as + `direct_store`, so the parse's `pool.add` calls (via `StoreBackedSegmentPool`) + intern each segment straight into the store's `content.blob` at parse time — + no second in-RAM `SegmentPool` copy — and the interned drain then no-ops over + the empty returned pool. Byte-identical to the eager drain by construction. See + [the in-process interned drain](#in-process-interned-drain-non-weka) below. + +Each weka payload contains: + +- `trace_id`; +- `node_ordinals`; +- profiling `NodeEnvelope` envelopes (carrying `prompt_segment_ids`); +- `(segment_id, role, content, wire_json)` segment tuples, where `wire_json` is + the verbatim `orjson.dumps(message)` raw-JSON blob for a raw-authored segment + (persisted byte-for-byte) and `None` for a role/content segment (the store + derives the `{"role", "content"}` blob); +- `structural_graph`: the content-free structural `ParsedGraph` (msgpack + bytes; `replay_outputs` and the segment pool emptied) — the weka pool emits + one single-trace parse per payload and attaches its structural graph (and + pool segments) to that payload. The streaming consumer merges the collected + graphs into the corpus structural graph that feeds the sidecar and + prefix-cache map. + +`_build_graph_store_streaming_trie` drains these weka payloads into the ONE +interned unified store via `build_unified_trie_store_from_payloads` — the SAME +unified store the non-weka in-process interned drain builds. There is no legacy +per-node store lane; the unified store is the sole build output. Streaming +`put_segment` deduplicates by content-addressed segment id, keeping parent +memory bounded. + +Alongside the store drain, each streamed trace's content-free structural graph +is collected and merged ONCE (`_merge_structural_graphs`, which hard-fails on an +empty or unmergeable stream); the merged graph feeds both the mandatory +`graph_meta` sidecar (`_write_graph_sidecar`) and the per-node prefix-cache map, +so weka payload-stream builds report both identically to the non-weka interned +drain. + +The interned unified node envelope contains int handles, not hex segment ids: + +```json +{ + "handles": [0, 1, 2], + "dispatch_overrides": {"max_tokens": 123, "model": "...", "stream": true}, + "stream": true +} +``` + +See [Graph Segment Unified Store](./graph-segment-unified-store.md) for the +on-disk file layout. + +Node ordinals are assigned by `trie_node_ordinals`: sort trie `LlmNode`s by +`arrival_offset_us`, tie-break by node id, and assign dense ordinals. The schedule +plane uses the same ordinal function, so `(trace_id, node_ordinal, "profiling")` +resolves to the manifest written at build time. + +### In-process interned drain (non-weka) + +Every non-weka format — dynamo, native, dag_jsonl — parses once in-process and +drains that SAME whole-graph parse through the interned builder, with no payload +round trip: `_build_interned_unified_store` drains the content-addressed pool +and every node's `prompt_segment_ids` manifest into the interned +`GraphSegmentUnifiedBackingStore` via `build_unified_trie_store_interned` +(resolving hex segment ids to int handles at build time), then the sidecar is +written DIRECTLY from the stripped whole parse +(`_write_graph_sidecar(parsed, ...)`), so its traces stay in PARSE order (the +weka merge reorders by id). + +**Dynamo direct write-through.** For `dynamo_trace` the store is constructed +BEFORE the parse and threaded into `parse_graph_workload(..., direct_store=store)`: +the adapter hands `build_trie_ir` a `StoreBackedSegmentPool` +(`adapters/dynamo/store_backed_pool.py`) whose `add()` write-throughs each segment +into `content.blob` at parse time, so the content pool is never materialized a +second time in RAM. The returned `ParsedGraph.segment_pool` (the shim) has an +empty `_by_id`, so the subsequent `_build_interned_unified_store` put loop +no-ops over it. Byte-identical to the eager drain by construction (both intern in +`build_trie_ir`'s content-loop first-occurrence order; three-way parity +direct == eager == streaming plus the golden digest pin it). Because the store is +live before the parse, the dynamo branch's `try/except → abort() + rmtree` covers +the parse too, so a mid-parse failure leaves no partial store. Native and +dag_jsonl keep the eager `SegmentPool` → drain path (their parse can fail before +any store dir exists). See +[Dynamo direct write-through route](./graph-segment-unified-store.md#dynamo-direct-write-through-route-storebackedsegmentpool) +for the shim contract and the measured content-pool collapse. + +This is also the only drain that persists dynamic-slot envelopes, so +slot-carrying graphs ride it with no separate fallback. A slot-carrying graph — +live-reply dag_jsonl lineage, or trie assembly items/capture stamped by the +native `@channel` lowering — cannot ride the weka streaming envelope +(`_trace_trie_envelopes` rejects slot metadata loudly), but it never reaches +that envelope: weka corpora carry no slots, and every non-weka format takes this +interned drain regardless of slots. `graph_carries_assembly_slots` no longer +routes the store build — it survives only as the schedule-plane t\*-gate +predicate (`workload_detect._gate_dynamic_slots_vs_tstar`). In practice most real +dag workloads carry slots (fork children inherit the parent's history including +its assistant reply, a live-reply lineage slot, so 7 of the 8 in-repo dag +fixtures carry slots; only spawn-only dag graphs are slot-free), but all of them +take the same interned drain. + +`build_unified_trie_store_interned` also serves as the parity-test oracle for +the weka payload-stream drain +(`tests/unit/dataset/test_dynamo_streaming_store_parity.py` and +`tests/unit/dataset/test_dag_jsonl_streaming_store_parity.py` pin the +payload-stream store byte-for-byte against the interned build); the non-weka +route face is pinned by `tests/unit/dataset/test_nonweka_interned_route.py`. + +### Store shape summary + +There is no environment flag and no per-shape branch selecting the store: every +graph parse (weka, dynamo, native, dag_jsonl) lowers onto the one interned +unified store. + +| Build path | Build function | Store written | +|------|-------|-------------| +| Weka payload stream (local file / dir / HF corpus id) | `build_unified_trie_store_from_payloads` | `GraphSegmentUnifiedBackingStore` (interned, A2, drained from worker-streamed payloads). See [unified store](./graph-segment-unified-store.md). | +| In-process interned drain (dynamo, native, dag_jsonl — one whole-graph parse, slot-free and slot-carrying alike; dynamo additionally write-throughs via `StoreBackedSegmentPool`, no second RAM pool copy) | `build_unified_trie_store_interned` | `GraphSegmentUnifiedBackingStore` (same store, drained from the whole-graph parse — the only drain that persists dynamic-slot envelopes). | + +The worker opens that one unified store; there is no other store shape. + +## Worker materialization + +Implementation: `src/aiperf/workers/worker.py` (store selection / error handling); +the `materialize_graph_request_unified` / `..._bytes` helpers live in +`src/aiperf/graph/worker_materialize.py` and are called from the worker. + +Workers resolve the store from the same `(base_path, benchmark_id)` used by +DatasetManager. + +`_graph_unified_reader` tries to open one `GraphSegmentUnifiedClient` for the run's +`benchmark_id` on the first credit and reuses the result (success or failure) for +the rest of the run. + +- `_graph_store_reader` returns that `GraphSegmentUnifiedClient`, or `None` after + caching a fatal `GraphStoreUnavailable` on the credit when the store is + absent or rejected. There is no legacy fallback. +- The one `GraphSegmentUnifiedClient` carries both the addressing face + (`get_node_envelope`) and the content face (`materialize_handles` / + `build_request_body_handles`); there is no separate segment-content store. + +Materialization then branches (both paths read the unified store): + +1. **Unified interned bytes path:** when cache busting is disabled + (`endpoint.cache_bust == CacheBustTarget.NONE`) and the node carries no dynamic + `items`, workers build the pre-serialized request body once from the unified + store's mmap content-pool slices via `materialize_graph_request_unified_bytes`. +2. **Unified interned dict path:** otherwise workers materialize a `messages` dict + via `materialize_graph_request_unified` (composing dynamic slots from the worker + pool when present), apply run-level payload options and cache-bust markers, then + send it. + +A failed store open becomes an explicit `GraphStoreUnavailable` (folding in +the A2-strict rejection reason when the store existed but was rejected); a missing +node manifest becomes `GraphEnvelopeMissing` with trace, instance, ordinal, and phase +variant in the message. + +## Structural handoff and trie caveat + +The structural sidecar mechanism is documented in +[Graph Structural Sidecar Handoff](./graph-structural-handoff.md). + +The sidecar is **mandatory** and written on both routes via +`_write_graph_sidecar`. For the non-weka in-process interned drain (dynamo, +native, dag_jsonl), the SAME real-content parse that built the store is stripped +to a content-free structural graph (`_write_graph_sidecar` → `strip_replay_text`) +and written to `graph_meta.msgpack` DIRECTLY — so its traces stay in parse order. +Trie ordinals are rebuilt via `flat_trie_ordinals` (keyed on the graph +topology), so the sidecar catalog matches the store; a structural catalog that +diverges from the store's build catalog raises `DatasetError` and fails the run +(a divergent sidecar would describe a DIFFERENT topology than the envelopes the +worker reads). + +The weka payload-stream build writes the same mandatory sidecar from a +**merged** structural graph: each streamed trace's content-free structural graph +is collected and merged once (`_merge_structural_graphs`, which raises +`DatasetError` on an empty or unmergeable stream), then written after the same +catalog cross-check. The merged structural graph also supplies the per-node +prefix-cache map. + +The TimingManager never re-parses: it ingests the sidecar from the exact path +the graph-typed `DatasetConfiguredNotification` advertised. A broadcast that is +not graph-typed, or whose advertised sidecar is missing, undecodable, or +index-divergent, is a hard `InvalidStateError` at configure time — no re-parse, +no env-convention path re-derivation. + +## Environment knobs + +All variables are under the `AIPERF_DATASET_` prefix unless noted. + +| Variable | Default | Pipeline effect | +|----------|---------|-----------------| +| `MMAP_BASE_PATH` | `None` | Base directory for the unified store (`aiperf_graph_segments_`) and the graph-meta sidecar (`aiperf_graph_meta_`). Falls back to system temp. Must be visible to DatasetManager and workers. | +| `WEKA_GRAPH_PARALLEL_THRESHOLD` | `8` | Work-item count (local trace files or HF rows) above which every weka route switches from serial in-process parsing to the multiprocess pool. `0` forces pool for any non-empty corpus. | +| `WEKA_GRAPH_PARALLEL_WORKERS` | `0` | Parse worker count. `0` auto-sizes from CPU count, item count, and auto max. | +| `WEKA_GRAPH_PARALLEL_AUTO_MAX_WORKERS` | `16` | Upper bound for auto worker sizing. | +| `WEKA_GRAPH_PARALLEL_PREFETCH_MULTIPLIER` | `16` | Ordered pool submit-window multiplier; bounds in-flight items to `workers * multiplier`. The window must cover the rows remaining behind the single heaviest trace, or fast workers stall head-of-line while it drains; the default `16` yields a `256`-row window at the auto 16 workers, covering the 393-row corpus whose heaviest row sits at index 140. Measured on that corpus: ~7.3 GiB parent VmHWM (~17.5 GiB tree) at `16` vs ~2.8 GiB parent at `4`. | +| `WEKA_GRAPH_PARALLEL_ITEM_TIMEOUT_SECONDS` | `600.0` | Per-item bound on one pool result. A worker killed mid-parse (OOM kill / external SIGKILL) otherwise presents as a silent indefinite hang; on expiry the parse raises a `RuntimeError` naming that cause. | +| `WEKA_HF_SPLIT` | `train` | HF split or slice syntax for Weka repo-id input; a slice such as `train[:100]` bounds the streamed rows. | +| `WEKA_HF_REVISION` | `None` | Optional HF revision pin for reproducible streamed ingest. | + +The former dataset flags that selected alternate trie store shapes (the +segment-trie IR gate, the mmap segment-store toggle, the unified-store toggle, and +the cross-run cache) were retired: the interned unified store is the sole +build path for every graph shape, and there is no cross-run cache. + +Run config synthesis fields also affect Weka parsing even though they are not +`Environment` fields — they reach every parse route through the ONE +`GraphParseContext` resolved by `resolve_graph_parse_context(run)`: + +- `synthesis.idle_gap_cap_seconds` / `--synthesis-idle-gap-cap`: per-trace idle + gap cap; default `60.0`, `null` disables warping. +- `synthesis.max_osl` / `--synthesis-max-osl`: caps each top-level chain + request's `dispatch_overrides["max_tokens"]` to `min(recorded out, max_osl)`; + subagent-body chains stay uncapped. `None` (default) leaves the recorded + `out` uncapped. +- `synthesis.corpus` / `--prompt-corpus`: prompt corpus for content synthesis; + default `coding`, with `sonnet` selecting the legacy Shakespeare pool. +- run random seed: pinned to a concrete int in the parent by + `resolve_effective_root_seed` before any parse — explicit `--random-seed` + first, else the ambient bootstrap-seeded root seed, else a fresh OS-entropy + seed (`secrets.randbits(64)`) generated once per resolution. One run's serial + and pool-worker parses synthesize identical bytes at any threshold; distinct + unseeded runs deliberately differ (no hardcoded fallback seed). The schedule + plane never re-parses — it ingests the structural sidecar this build wrote and + broadcast — so content bytes never need to cross the build/schedule split; the + only synthesis determinism that matters is between the in-process parse and + its spawn-started pool workers. The dynamo adapter resolves the identical + ladder at `from_dynamo_trace` entry, so dynamo content synthesis follows the + same seed semantics. + +## Corpus-scale memory measurement + +The build-plane RAM cost at corpus scale (order 1M dynamo nodes) is sized by an +opt-in, `@pytest.mark.slow` measurement isolate rather than a committed +multi-GB fixture: + +- `tests/harness/dynamo_synth_corpus.py` (`write_synthetic_dynamo_capture`) + emits a deterministic, chain-heavy `dynamo.request.trace.v1` capture: each + turn re-lists its full prefix plus `new_blocks_per_turn` fresh globally-unique + 64-bit hashes, so `input_sequence_hashes` grows every turn (the hash-id-slot + amplification real captures show) while `input_length` stays block-consistent + with the covered-count ISL gate. +- `tests/unit/dataset/graph/adapters/test_dynamo_corpus_scale_memory.py` runs a + real `from_dynamo_trace` under `tracemalloc` (corpus pre-warmed outside the + traced region), attributes the peak-window snapshot by allocation-site + filename into four tiers, and linearly extrapolates each to 1M nodes: + **hash-id ints/lists** (`dynamo/trace_reader.py` + `dynamo/trie_lowering.py`), + the **decode cache** (`adapters/shared/content.py`), the **resolution-trie** + transient (`segment_ir/trie_content.py`, measured in its own phase isolate + because it is freed before emission), and the **content pool** + (`segment_ir/pool.py`). The budget is a calibrated RATIO (measured peak within + 1.5x an analytic per-tier model derived from the generator parameters), never + an absolute-bytes gate; peak RSS is logged, never asserted. + +Run it (deselected by default) and, for a manual corpus-scale run, override the +node count following the `AIPERF_TEST_WEKA_CORPUS_DIR` precedent: + +```bash +uv run pytest -m slow \ + tests/unit/dataset/graph/adapters/test_dynamo_corpus_scale_memory.py -s +AIPERF_TEST_DYNAMO_SCALE_NODES=200000 uv run pytest -m slow \ + tests/unit/dataset/graph/adapters/test_dynamo_corpus_scale_memory.py -s +``` + +The measured tier table (recorded in the test's module docstring) is the go/no-go +instrument for the trie-memory reduction work: the hash-id-int tier is the +dominant residue at ~1M-node scale, followed by the decode cache, the +resolution-trie transient, and the content pool. + +### Dynamo decode-cache mechanics + +`dynamo_recon_callbacks` (`adapters/dynamo/trie_lowering.py`) keeps a PRIVATE +per-parse decode cache (the shared synthesizer's `pg._cache` is keyed by bare +hash id, so mixing dynamo's 16-token blocks with weka's 64-token blocks in it +would return wrong-sized blocks across adapters). That cache stores one `int` +corpus **offset** per unique hash id — not the decoded token list — via +`CorpusContentSynthesizer._decode_block_tokens_offset_cached` +(`adapters/shared/content.py`): a cache miss issues the identical +`reseed_for_hash_id(h)` + `randrange(corpus_size)` pair the list-cache decoder +would (a full per-hash reseed, so the start offset is a pure function of +`(seed, scope, hash_id, corpus_size)`), and every decode re-slices the block +from the shared corpus window. This shrinks the decode-cache tier from +`~block_size` ints per entry to one int per entry (measured at 1M-node +extrapolation: isolate 9.0 to 4.4 GB, full-parse window 5.5 to 0.9 GB) at a +measured ~2% full-parse CPU cost, byte-identical +by construction and pinned by a differential test plus the golden store-digest. +The cached offsets assume the corpus is immutable for the synthesizer's +lifetime (the `attach_shared_corpus` byte-identical-rebind contract); a +size-changing rebind fails loud. The weka path is untouched — it calls +`_decode_block_tokens` (the list cache) directly. + +### Dynamo hash-id interning + +Dynamo records the FULL prompt block-hash list on every `request_end`, so a +chain re-lists each earlier block on every later turn. `_collect_records` +(`adapters/dynamo/trace.py`) is the single point where the streaming reader +is fully drained into memory before lowering — every recorded +`input_sequence_hashes` slot across the whole capture is simultaneously live (the +"record-window plateau") before a single `TrieRequest` exists. As each record is +kept, `_intern_replay_hashes` rewrites its hash list in place through a per-parse +`dict[int, int]` intern table, so every re-listed occurrence of a hash VALUE +across all turns and sessions shares ONE canonical `int` object instead of a +fresh ~36 B `orjson` allocation. The `list()` copy in `dynamo_trie_nodes` +preserves that shared identity into `TrieRequest.hash_ids`, so the canonical +objects persist through resolution and the content loop. This collapses the +recorded hash-id-int tier from one int per re-listed slot to one per unique value +(measured ~149.0 → ~42.7 MB at the corpus-scale isolate; ~24.8 → ~7.1 GB @1M), +leaving only the per-slot list pointers on the total-slot count. + +The change is values-only: the table is keyed and probed by `==`, no value ever +changes, and every downstream consumer of the hashes is value-semantic, so the +store `content`/`nodes` bytes, the envelope `prompt_segment_ids`, and the +`graph_meta` sidecar are all byte-identical (the golden store digest and the +three-way direct/eager/streaming parity gate this). The intern table is a +per-parse local — born before the read loop, dropped when collection returns, +never module state and never crossing parses. Negative ids never enter it: +recorded hashes are validated non-negative at read time and the virtual negative +ids are minted later at lowering, so `hash(-1) == hash(-2)` is at most an extra +bucket probe, never a value collision. The weka path is untouched — its records +never flow through `_collect_records`. + +### Dynamo sid-string interning + +Each node's `prompt_segment_ids` re-lists its whole message chain root→tip, so a +`segment_id` hexdigest (`segment_ir/pool.py`) born once at a segment's first +occurrence is re-minted as a fresh 32-char str on every later re-listing — ~246k +duplicate strings for only ~17,850 unique segments at the corpus-scale isolate, +retained for the whole `ParsedGraph` lifetime on BOTH dynamo routes. `SegmentPool` +already stores exactly one canonical `Segment` (hence one canonical `Segment.id` +string) per unique value, so the fix reuses that as the intern table: the eager +route builds an `InterningSegmentPool` whose `add*` returns `self._by_id[sid].id` +(the first-born canonical), and the direct write-through `StoreBackedSegmentPool` +keeps a handle-indexed `_sids` list (dense `put_segment` insertion indexes) and +returns `_sids[handle]` on a repeat. Both live in +`adapters/dynamo/store_backed_pool.py` (the dynamo pool shims). This collapses the +sid-string addressing tier to one str per unique segment (measured ~19.7 → ~3.0 MB +`pool.py` eager / ~18.0 → ~1.3 MB direct at the isolate; ~3.0 → ~0.5 GB @1M), the +direct `_sids` pointer list adding ~0.2 MB (~24 MB @1M). + +The change is values-only — interning shares string OBJECTS, never changes VALUES. +`put_segment` receives identical sid/role/content in identical first-occurrence +order (the intern happens on `add`'s RETURN, after the store/pool write), so the +`content`/`nodes` store bytes are byte-identical and the golden digest passes +without re-pinning; the envelope serializes `prompt_segment_ids` by value +(`orjson.dumps`), and `strip_replay_text` empties `metadata["trie"]` and swaps in a +fresh `SegmentPool` before the sidecar is encoded, so equal-valued strings encode +identically regardless of identity. The eager subclass depends on `Segment.id` +being the first-born string (stable while `pool.py` is frozen); the direct shim's +defensive fall-through degrades to a fresh value-correct sid if the store was +pre-populated (never in production). Weka is out of scope — its per-worker payload +stream can't share a per-parse intern table, and it constructs plain `SegmentPool`s. + +## Operational invariants + +- Weka `hash_id_scope` must be `local` (per-trace hash namespace: equal hash ids + across traces synthesize different bytes) or `global` (one hash namespace + shared across all traces: equal hash ids synthesize byte-identical blocks, + reproducing recorded cross-trace KV sharing); any other value is rejected + with `WekaHashScopeError`. +- Graph dispatch emits chat-completions request bodies verbatim; non-chat endpoint + types are rejected before store build. +- Build and schedule must derive node ordinals from the same parsed graph shape; + trie ordinals are dense per trace and sorted by `(arrival_offset_us, node_id)`. +- Trie prompt bytes are addressed through interned int handles, not through + predecessor channel contents. +- The interned unified store is the sole build output for every graph shape + (weka — file, directory, or HF id — via the worker-pool payload stream; + dynamo, native, and dag_jsonl — slot-free and slot-carrying alike — via the + in-process interned drain of one whole-graph parse); the legacy per-node store + was removed. There is no cross-run graph cache. +- Timing edges are interval-order (`A → B` iff `A` finished-before `B` and + `rank(A) < rank(B)`), reduced to the finished-before frontier, over the pass's + node set — global for weka/native/dag_jsonl, but per SESSION-TREE for dynamo + (each root + `parent_trajectory_id` descendants is lowered independently, so + independent trees never gain a cross-parent edge). The one carve-out + is `apply_start_anchors`: a node whose recorded start overlaps its stamped + `causal_parent_id` interval has that frontier replaced by a single start-anchored + edge, optionally REFINED with a `delay_after_predecessor_first_token_us` when the + parent streamed and the child began post-TTFT. There is still no `spawner` / + `chain_prev` / joined-leaf CANDIDATE SET feeding the base rule; `causal_parent_id` + is a single hint consumed only by the overlap carve-out. +- Post-TTFT anchoring is inert only when the SOURCE node is non-streaming: the + refinement is built into the IR unconditionally, and the runtime observes a first + token whenever the worker parses that parent's SSE stream. Each graph node + dispatches per its own recorded `streaming` mode (a per-request override), so a + recorded-streaming source is anchored regardless of the global `--streaming` + flag — the flag is only the fallback for mode-less payloads and non-graph runs. + Only a first-token edge whose source carries `streaming=False` silently degrades + its refined children to their start anchor; the TimingManager warns once at + configure time for exactly that case. +- A block's `(role, starts_new_message)` tag is frozen at its first creator and + inherited verbatim; two requests sharing a block-aligned prefix render an + identical leading message-id chain (never relabeled or coalesced). +- Reconstructed prompt length equals the block-aligned covered-count + `min(len(hash_ids), in // block_size) * block_size`; a mismatch is a hard build + abort (`TrieISLMismatchError`). + +## Validation boundary + +The offline unit and component tests +(`tests/unit/graph/test_weka_trie_interval_order.py`, +`tests/component_integration/graph/test_weka_trace_fidelity.py`) assert the +**structural** invariants: interval-order edge topology, frozen per-block tags, +shared-prefix identical message-id chains, boundary preservation, and the +block-aligned covered-count ISL gate. These are provable in-process because they +are properties of the reconstructed IR, not of any server. + +The **end-to-end cache-hit claim** — that a shared block prefix actually produces +a KV prefix-cache hit — is provable only against a real prefix-caching inference +engine (vLLM, SGLang, or a real Dynamo deployment). AIPerf's own mock server is +throughput-only: it has no KV cache or prefix-cache simulation, so it can neither +confirm nor refute the cache-hit claim. The weka trie path is **synthesize-only** +(it reconstructs faithful request bodies from the recorded blocks); it does not +emit a `hash_replay` dispatch. Block-hash replay against a KV simulator is a +Dynamo-path concern (validated against `dynamo-mocker`) and is out of scope for +the weka path regardless of its declared `hash_id_scope`. + +## Key symbols + +| Symbol | Location | +|--------|----------| +| `resolve_graph_workload` / `is_graph_workload_path` / `parse_graph_workload` / `resolve_graph_parse_context` / `validate_graph_endpoint_type` / `_gate_dynamic_slots_vs_tstar` | `src/aiperf/dataset/graph/workload_detect.py` | +| `GraphParseContext` / `UNSET` / `publish_ctx_tokenizer_env` | `src/aiperf/dataset/graph/parse_context.py` | +| `GraphWorkloadRef` | `src/aiperf/config/resolution/plan.py` | +| `GraphAdapterProtocol` | `src/aiperf/dataset/graph/adapters/protocols.py` | +| `parse_graph` / `GraphParseError` | `src/aiperf/dataset/graph/parser.py` | +| `WekaTraceAdapter` / `from_weka_trace` | `src/aiperf/dataset/graph/adapters/weka/trace.py` | +| `stream_weka_trace_segment_payloads` | `src/aiperf/dataset/graph/adapters/weka/trace.py` | +| `DynamoTraceAdapter` / `from_dynamo_trace` | `src/aiperf/dataset/graph/adapters/dynamo/trace.py` | +| `_WorkItem` / `file_work_items` / `row_work_items` | `src/aiperf/dataset/graph/adapters/weka/trace_parallel.py` | +| `_map_items` / `parse_items` / `iter_item_segment_payloads` | `src/aiperf/dataset/graph/adapters/weka/trace_parallel.py` | +| `resolve_effective_root_seed` | `src/aiperf/dataset/graph/adapters/shared/content.py` | +| `build_trie_graph` (Weka flatten/walk) | `src/aiperf/dataset/graph/adapters/weka/trie_build.py` | +| `build_trie_ir` / `resolve_content_parents` / `compute_asst_caps` / `assign_block_tags` / `assemble_messages` / `assert_covered_isl` / `TrieISLMismatchError` | `src/aiperf/dataset/graph/segment_ir/trie_content.py` | +| `compute_ranks` / `build_interval_edges` / `apply_start_anchors` / `_excluded_async` | `src/aiperf/dataset/graph/segment_ir/interval_order.py` | +| `apply_idle_gap_warp` / `ActiveIdleWarp` / `compute_turn_block_geometry` / `block_role_split` / `add_message_chain` | `src/aiperf/dataset/graph/segment_ir/trie_content.py` | +| `chop_trie_at_tstar` | `src/aiperf/timing/snapshot_chop.py` | +| `Segment` / `SegmentPool` / `segment_id` | `src/aiperf/dataset/graph/segment_ir/pool.py` | +| `DatasetManager._configure_graph_dataset` / `_configure_graph_workload` / `_build_graph_store` | `src/aiperf/dataset/dataset_manager.py` | +| `GraphStoreBuilder` / `GraphStoreBuildResult` | `src/aiperf/dataset/graph/store_build.py` | +| `GraphStoreBuilder._build_graph_store_streaming` / `_build_graph_store_streaming_trie` / `_build_interned_unified_store` / `_write_graph_sidecar` | `src/aiperf/dataset/graph/store_build.py` | +| `GraphStoreBuilder._merge_structural_graphs` / `_build_graph_prefix_cache_by_trace` | `src/aiperf/dataset/graph/store_build.py` | +| `trie_node_ordinals` / `flat_trie_ordinals` / `TraceSegmentPayload` / `iter_trace_segment_payloads` / `graph_carries_assembly_slots` | `src/aiperf/dataset/graph/segment_ir/store_builder.py` | +| `build_unified_trie_store_interned` / `build_unified_trie_store_from_payloads` | `src/aiperf/dataset/graph/segment_ir/store_builder.py` | +| `TimingManager._load_graph_sidecar` / `_sidecar_passes_index_check` | `src/aiperf/timing/manager.py` | +| `GraphSegmentClientMetadata` / `GraphDatasetMetadata` | `src/aiperf/common/models/dataset_models.py` | +| `Worker._graph_unified_reader` / `Worker._graph_store_reader` | `src/aiperf/workers/worker.py` | diff --git a/docs/reference/graph-ir-schema.md b/docs/reference/graph-ir-schema.md new file mode 100644 index 0000000000..d8fad987d3 --- /dev/null +++ b/docs/reference/graph-ir-schema.md @@ -0,0 +1,287 @@ + + + +# Graph IR schema reference + +This page documents the native AIPerf graph intermediate representation (IR): the authorable JSON/YAML/JSONL shape, the typed in-memory model it decodes to, and the boundary between stable workload fields and adapter/runtime internals. + +The canonical typed models live in `src/aiperf/dataset/graph/models.py`. Native files are parsed by `parser.py`, loose/legacy shapes are normalized by `decode.py`, and auto-derived defaults are applied by `auto_derive.py` (which also normalizes runtime channels), followed by interned-segment lowering in `native_lowering.py`. There is no native serializer: the loader reads native files but does not emit them. + +## Native file formats + +A native graph workload can be written as either: + +- **JSONL records**: one JSON object per non-empty line. Each object is a record with a `kind` field. +- **JSON/YAML document**: a single mapping with top-level `graph` and `traces` sections. YAML multi-document files may also use one `kind` record per document. + +Benchmark workload detection selects native parsing only when the caller chooses the native graph format, for example via `--graph-format native`; auto-detection excludes `native` so arbitrary JSONL/YAML chat datasets are not hijacked into graph mode. The lower-level `parse_graph(...)` helper can still auto-detect the native adapter from a `.yaml`, `.yml`, or `.jsonl` suffix when called directly by tooling. + +### JSONL record shape + +```jsonl +{"kind":"graph","version":"2.0","nodes":{"llm":{"node_type":"llm","prompt":["@messages"],"output":"messages"}}} +{"kind":"trace","id":"trace-1","initial_state":{"messages":[{"role":"user","content":"hi"}]}} +``` + +Supported record kinds: + +| `kind` | Body decodes as | Notes | +| --- | --- | --- | +| `graph` | `GraphRecord` | At most one. Must appear before any `trace` record. | +| `trace` | `TraceRecord` | If `kind` is omitted in JSONL or multi-document YAML, the parser defaults the record to `trace`. | + +Earlier revisions of this schema carried a rich node-kind zoo (`replay`, `tool_call`/`tool_result`, `subgraph`, `spawn`/`await`, `delay`, `barrier`, `compact`, `bootstrap`, `loop`), conditional edges, `mix` records, top-level subgraphs, and per-graph endpoint pools. All of these were retired with their runtime: every live producer (the weka and dynamo trace adapters, and the native lowering) emits flat `llm`-node topologies wired with static edges. Authored files that still carry one of these constructs fail at parse like any other invalid input — as an unknown node kind, unknown record kind, or unknown field (`branches`, `endpoints`). + +Parser ordering and cardinality rules: + +- A `graph` record must precede all `trace` records. +- More than one `graph` record is rejected. +- Unknown record kinds (anything other than `graph` and `trace`, including the former `mix` and `subgraph`) are rejected. +- A `graph` record carrying an `endpoints` block is rejected as an unknown field; use the global `--url`/`--model` run configuration. + +### Single-document JSON/YAML shape + +The same workload is usually easier to author as one JSON/YAML mapping: + +```yaml +graph: + version: "2.0" + provenance: + source: hand-authored + tool: manual + state: + messages: + type: messages + reducer: add_messages + nodes: + llm: + node_type: llm + prompt: + - "@messages" + output: messages + edges: + - source: START + target: llm + - source: llm + target: END + +traces: + - id: trace-1 + tags: [chat] + initial_state: + messages: + - role: user + content: hi +``` + +Single-document expansion maps: + +- `graph:` to a `kind: graph` record. +- Each item in `traces:` to a `kind: trace` record. + +A `mix:` or `subgraphs:` section is rejected at the top-level-key gate like any other unknown key (the record kinds they used to expand to were retired with their runtime). + +When authoring by hand, the single-document layout (a `graph:` block followed by `traces:`) is the recommended form; the equivalent JSONL form is an ordered `kind: graph` record followed by `kind: trace` records. + +## Top-level in-memory shape + +Parsing native files produces a `ParsedGraph`: + +| Field | Type | Authoring status | +| --- | --- | --- | +| `graph` | `GraphRecord` | Main/top-level graph. Native files author this under `graph:` or `kind: graph`. Defaults to an empty graph if absent. | +| `graphs` | map of graph key to `GraphRecord` | Per-trace top-level graphs for multi-graph workloads (Weka heterogeneous directories, per-trace native lowering), keyed by the value each `TraceRecord.graph_ref` names. Resolved via `resolve_trace_graph`. Not authored. | +| `traces` | list of `TraceRecord` | Native files author this under `traces:` or `kind: trace` records. | +| `segment_pool` | `SegmentPool` or null | Content-addressed segment pool that backs every `LlmNode`'s interned prompt; set by every live producer and drained into the unified store by the build plane. Do not author. | + +## Graph records + +A `GraphRecord` declares topology and channel state: + +| Field | Type | Default | Notes | +| --- | --- | --- | --- | +| `version` | string | `"2.0"` | Known schema version. | +| `provenance` | `ProvenanceSpec` | `{source: hand-authored, tool: manual}` | Origin metadata. Unknown provenance keys round-trip through `extra`. | +| `system` | string or null | null | Linear-chat shorthand system prompt. Used only when no explicit nodes are authored. | +| `state` | map of channel name to `ChannelSpec` | `{}` | Explicit channel declarations. Missing safe runtime channels may be auto-derived. | +| `nodes` | map of node id to node spec | `{}` | Node ids must not be `START`, `END`, or start with `_aiperf`. | +| `edges` | list of `StaticEdge` | `[]` | `START` and `END` are sentinel ids. | + +If a main graph has nodes but omits explicit sentinel edges, the native parser auto-injects `START -> ` for nodes with no incoming edge and ` -> END` for nodes with no outgoing edge. Existing explicit `START` or `END` edges are preserved. This auto-injection applies to the main graph parse path. + +If a graph has no nodes, `auto_derive.py` synthesizes a linear chat graph: + +- node id `_llm` +- prompt containing optional `graph.system` plus `@messages` +- `messages` channel with `type: messages`, `reducer: add_messages` +- `START -> _llm -> END` + +Trace-level `messages` are lifted into `initial_state.messages` for every workload (with or without explicit nodes), before linear-chat synthesis runs. + +## Channels and reducers + +`graph.state` maps channel names to `ChannelSpec`: + +| Field | Values | Default | Meaning | +| --- | --- | --- | --- | +| `type` | `text`, `messages` | `text` | Value modality. | +| `reducer` | `overwrite`, `add_messages` | `overwrite` | How writes merge. | + +Reducer semantics: + +- `overwrite`: one writer value wins; static validation catches multi-writer conflicts, and runtime catches duplicate overwrite writes from any writer that reaches the channel. +- `add_messages`: appends message lists and replaces prior messages with the same message `id`. Pair with `type: messages`. + +Auto-derived channels: + +- Node output channels, trace `initial_state` channels, and trace `replay_outputs` channels are declared if missing. +- A missing channel named `messages` defaults to `type: messages`, `reducer: add_messages`. +- Other missing runtime channels default to `type: text`, `reducer: overwrite`. +- LLM prompt references in the main graph at top-level prompt-array position, such as `"@messages"`, auto-declare missing channels as `messages/add_messages`. +- LLM prompt references in the main graph inside a message content block, such as `content: ["@context"]`, auto-declare missing channels as `text/overwrite`. + +Prompt channel references use `@channel`; `@@literal` escapes a leading at-sign. + +## Channel requirements and `count` semantics + +Every node inherits a common `inputs` field: + +```yaml +inputs: + - channel: draft_ready + count: 1 + - channel: fanout_done + count: all +``` + +`inputs` is an AND-fan-in gate: a node with non-empty `inputs` fires only after every requirement is satisfied. An empty `inputs` list keeps the legacy successor-walk behavior. + +`ChannelRequirement.count` is not a token count, byte count, or list length. It counts node-write arrivals on the named channel: + +- `count: N` waits until the Nth node write to that channel has arrived. `N` must be at least 1. +- `count: all` resolves to the static number of declared producers for that channel and waits for that many arrivals. +- `trace.initial_state` seeds a channel value but does not increment the arrival count. +- The read snapshot captures the first `N` node writes plus any initial seed as reducer input. +- If producer completion/cancellation makes the requested count unreachable, runtime raises an orphaned-channel error instead of waiting forever. + +Use `count: 1` for “fire when any one producer writes this gate channel” and `count: all` for “wait for every statically declared producer of this gate channel.” + +## Node common fields + +Every node inherits these optional fields: + +| Field | Type | Meaning | +| --- | --- | --- | +| `metadata` | mapping | Free-form span metadata, emitted under `aiperf.meta.*`. | +| `min_start_delay_us` | non-negative float or null | Minimum delay after all predecessors/inputs are satisfied before this node may start. | +| `arrival_offset_us` | non-negative integer or null | Recorded offset from trace arrival, used by adapters/snapshot replay. | +| `inputs` | list of `ChannelRequirement` | Explicit AND-fan-in channel gate. | + +`decode.py` accepts both `node_type` and legacy `kind` as the node discriminator; the only valid value is `llm`. Any other value fails as unknown, and a node with no discriminator decodes as `llm`. `messages` is accepted as an alias for `prompt`; a node that sets both is rejected. A node with neither fails to decode (`prompt` is required) — in particular, a prompt-less node carrying only `expected.input_tokens` is rejected with an explicit error: synth-token prompt fabrication is not supported on the native lowering path. Unknown node fields are rejected (`forbid_unknown_fields`), so a typo'd field name fails loudly instead of silently changing what is benchmarked; the same applies to `graph` and `trace` records, with `provenance` remaining the vendor-key catch-all (unknown provenance keys fold into `extra`). + +## Node kinds + +`llm` is the only node kind. + +### `llm` + +Dispatches an LLM request. + +Canonical fields: + +- `prompt`: list of prompt items (`messages` is accepted as an alias; a node without either fails to decode). +- `output`: output channel name. + +Important optional fields: + +- `streaming`: defaults to `true`. +- `expected`: expected token accounting (`input_tokens`, `output_tokens`, cache fields). +- `extra_body`: request-body overrides such as temperature or provider-specific keys (Turn naming; model / streaming / max_tokens / raw_tools / extra_headers are first-class node fields). + +Writes: `[output]`. + +## Edges + +Edges are `StaticEdge` records, tagged `edge_type: static` in canonical output; `decode.py` also accepts loose edge objects without the tag. An edge carrying a `branches` mapping (the former conditional-edge shape) is rejected at decode as an unknown field. + +```yaml +- source: START + target: llm + delay_after_predecessor_us: 1000 + min_start_delay_us: 500 +``` + +Fields: + +- `source`: source node id or `START` (a source that is neither is rejected by validator rule 56). +- `target`: target node id or `END` (a target that is neither is rejected by validator rule 56). +- `delay_after_predecessor_us`: optional non-negative idle/scheduling delay after predecessor completion. +- `delay_after_predecessor_start_us`: optional non-negative idle delay measured from the moment the predecessor DISPATCHES (the successor does not await the predecessor's completion or output). Mutually exclusive with `delay_after_predecessor_us`; an edge that sets both is rejected by validator rule 54. +- `delay_after_predecessor_first_token_us`: optional non-negative idle delay measured from the predecessor's OBSERVED FIRST TOKEN; the runtime gates the successor at `first_token_wall + this delay`, falling back to dispatch + `delay_after_predecessor_start_us` when the predecessor terminates without a first token. Only valid alongside `delay_after_predecessor_start_us` (the dispatch fallback), must not combine with `delay_after_predecessor_us`, and cannot be `START`-sourced — validator rule 55. The native lowering also enforces the missing-fallback case at parse time. +- `min_start_delay_us`: optional non-negative minimum wait on the successor after predecessors are satisfied. + +All delay values must be finite: `inf`/NaN are rejected at decode time (`GraphDecodeError`) and by validator rule 57 for graphs constructed directly from typed structs — a non-finite gate never clears and would hang the trace. + +## Trace records + +A `TraceRecord` supplies per-session data: + +| Field | Type | Notes | +| --- | --- | --- | +| `id` | string | Required stable trace id. | +| `tags` | list of strings | Opaque provenance labels, round-tripped through the codec/sidecar; no runtime consumer reads them. | +| `graph_ref` | string or null | Multi-graph selector naming a key in `ParsedGraph.graphs`; resolution goes through `resolve_trace_graph`. Set by adapters and the native lowering; native single-graph workloads leave this null. | +| `messages` | list of messages or null | Linear-chat shorthand, equivalent to `initial_state.messages`; lifted into `initial_state.messages` for every trace (explicit-node graphs included) unless `initial_state.messages` is already set, which wins. | +| `initial_state` | mapping | Initial channel values at trace start. Seeds reducers but does not satisfy `ChannelRequirement.count`. | +| `replay_outputs` | mapping of node id to channel values | Authorable native-file surface read by `auto_derive` for channel inference; adapters never populate it, and the structural sidecar strip (`graph_meta_sidecar`) clears it. | + +Channel values spliced by `@channel` prompt references must be plain data: a messages splice requires a list of `{role, content}` message dicts and a text splice requires a string. Typed/multimodal directive blocks and non-string content blocks are rejected by the unified-store lowering with an actionable error. + +## Advanced and internal fields boundary + +The graph IR intentionally preserves more information than the runtime needs to execute a hand-authored workload. Treat fields in three groups: + +### Stable authoring surface + +These are normal native workload fields: + +- `graph.version`, `provenance`, `system`, `state`, `nodes`, `edges` +- `traces` +- Node common fields, `llm` node fields, `inputs`, edge timing anchors + +### Advanced preserved metadata + +These may be authored, but are primarily produced by adapters or used for fidelity/analysis rather than core scheduling: + +- `expected` +- `extra_body` +- `arrival_offset_us` +- vendor-specific keys inside the `ProvenanceSpec` catch-all + +Unknown vendor keys round-trip only for models that explicitly preserve extras (`provenance.extra`). Native single-document parsing rejects unknown root keys with the offending key named and the nearest valid key suggested (only `graph` and `traces` are recognized). + +### Runtime/adapter internals + +Do not hand-author these in native graph files unless you are extending the loader/runtime: + +- `ParsedGraph.graphs` +- `ParsedGraph.segment_pool` + +## Validation highlights + +The validator enforces cross-field rules. Common authoring errors include: + +- cycles in static topology +- unreachable nodes +- multiple writers to an `overwrite` channel, with broader duplicate-write protection enforced at runtime +- node ids colliding with `START`/`END` or the `_aiperf` reserved prefix +- unknown `graph.version`, or a non-hand-authored graph missing `provenance.tool` (error) or still carrying the default `"manual"` (warning) +- an edge setting both `delay_after_predecessor_us` and `delay_after_predecessor_start_us` (mutually exclusive; validator rule 54) +- a `START`-sourced edge that is start-anchored (`delay_after_predecessor_start_us`); the `START` pseudo-node never dispatches, so its start-anchored successors would never fire — use `min_start_delay_us` for absolute offsets (validator rule 54) +- a first-token-anchored edge (`delay_after_predecessor_first_token_us`) that sets no `delay_after_predecessor_start_us` fallback, combines with `delay_after_predecessor_us`, or is `START`-sourced; the START pseudo-node never dispatches or streams a first token (validator rule 55) +- an edge whose `source`/`target` is neither a declared node nor the matching `START`/`END` sentinel (validator rule 56) +- a non-finite delay value on an edge or a node `min_start_delay_us` (validator rule 57; also rejected at decode time) + +`validate()` runs every rule over the main graph and over every per-trace graph in `ParsedGraph.graphs`; issues in a `graphs` entry are located under `graphs[].…`. + +Only parse-time checks run automatically before a benchmark: schema/type decoding (surfaced as `GraphParseError`) plus the record-ordering and singleton-graph rules (rule-19, rule-20) enforced in `parser.py`. The full semantic rule set (`validate()` in `validator.py`) is not invoked by the profile runtime — it is an adapter-output contract exercised by unit tests and available to tooling that lints a `ParsedGraph`. See [Graph IR Validation](./graph-ir-validation.md). diff --git a/docs/reference/graph-ir-validation.md b/docs/reference/graph-ir-validation.md new file mode 100644 index 0000000000..c7a2f384bd --- /dev/null +++ b/docs/reference/graph-ir-validation.md @@ -0,0 +1,158 @@ + + +# Graph IR Validation Reference + +This reference describes the static validation layer for AIPerf Graph IR workloads. The rule set is deliberately small: it is the contract for what the live trace adapters (weka, dynamo) emit — `LlmNode` topologies wired with `StaticEdge` timing anchors — plus the shared graph-record invariants (version, provenance, reserved names). Rules for constructs no live producer emits (conditional edges, mix records, subgraphs, loops, barriers, spawn/await, replay coverage, endpoint pools, channel directives, streaming reducers, placement hints, retry policies) were retired; see git history for the removal rationale. + +Primary implementation files: + +- `src/aiperf/dataset/graph/validator.py` +- `src/aiperf/dataset/graph/parser.py` + +## Validation lifecycle + +Graph IR validation happens after a workload has been parsed into a `ParsedGraph`. + +1. `parse_graph(...)` reads native Graph IR files or delegates to a graph adapter. +2. Native parsing enforces record-level shape and ordering, including parse-time rule-19 and rule-20 failures. +3. The parser decodes typed Graph IR models and wraps schema decode failures as `GraphParseError` messages. +4. `validate(parsed)` runs the validator rule set and returns all findings. + +The profile runtime does not invoke `validate(...)` automatically; the rule set is exercised as an adapter-output contract by unit tests and is available to tooling that wants to lint a `ParsedGraph`. + +The validator does not stop after the first problem. It accumulates every implemented `ValidationIssue` so related wiring mistakes can be fixed in one pass. + +## Finding format and severity + +Each validator finding is a `ValidationIssue` with these fields: + +| Field | Meaning | +|---|---| +| `rule_id` | Stable rule label such as `rule-1`, when the rule is implemented in source. | +| `severity` | `error` or `warning`. Errors block execution. Warnings are informational unless the caller elects to treat them as blocking. | +| `location` | Human-readable pointer to the offending graph, node, edge, or trace field. | +| `message` | Plain-English explanation of the problem. | +| `suggested_fix` | Optional remediation hint. | + +Use the `location` first when debugging. It is intentionally close to author-facing Graph IR paths, for example `graph.nodes.`, `graph.edges[->]`, `graph.provenance.tool`, or `graph.version`. + +## Rule categories + +The current validator dispatches the following verified rule IDs. Rule IDs not listed here should not be documented as implemented without checking source first. + +| Area | Verified rule IDs | What is checked | +|---|---:|---| +| Structural topology | `rule-1`, `rule-21` | Cycles (iterative DFS — safe on 100k+-node chains); reachability from `START`. | +| Writer conflicts | `rule-9` | Two or more nodes writing the same overwrite-reducer channel. | +| Versioning and provenance | `rule-11`, `rule-12`, `rule-13` | Reserved node names, known graph version, provenance tool (missing tool on a non-hand-authored graph is an error; the default `"manual"` on a non-hand-authored graph is a warning — it means the generating adapter never stamped itself). | +| Engine expectations | `rule-15` | Warning when `expected.cache_read_tokens` / `expected.cache_creation_tokens` are set — not all engines report cache fields. | +| Timing edge anchors | `rule-54`, `rule-55` | Edge delay-anchor exclusivity and first-token-anchor shape (dispatch fallback present, no completion anchor, no `START` source). | +| Edge endpoints | `rule-56` | Every edge `source` is a declared node or `START`; every edge `target` is a declared node or `END`. Dangling endpoints (typos, `END` as source, `START` as target) are errors. | +| Delay finiteness | `rule-57` | Every edge delay field (`delay_after_predecessor_us`, `min_start_delay_us`, `delay_after_predecessor_start_us`, `delay_after_predecessor_first_token_us`) and node `min_start_delay_us` must be finite — a non-finite gate never clears and hangs the trace. Catches typed-struct producers; the loose decoder rejects non-finite values at decode time. | +| Parse-time record ordering | `rule-19`, `rule-20` | Native record order and the singleton graph record, enforced by `parser.py` before `validate(...)` returns findings. | + +`validate(parsed)` runs every rule over the main `parsed.graph` **and** over every per-trace graph in `parsed.graphs` (multi-graph workloads: weka heterogeneous directories, per-trace native lowering). Issues found in a `parsed.graphs` entry are re-located under `graphs[].…` so the offending graph is identifiable; the aliased first entry (identical object to `parsed.graph`) is not double-reported. + +## Common failures and remediations + +### Graph topology errors + +Symptoms: + +- Cycle findings on `graph.edges` (rule-1). +- Nodes reported unreachable from `START` (rule-21). + +Fixes: + +- Pre-unroll loops into acyclic trace topology. +- Add missing `START -> node`, `node -> node`, or `node -> END` edges, or remove dead nodes. For adapter-emitted IR, fix the adapter's edge lowering rather than the generated file. + +### Writer conflicts + +Symptoms: + +- Multiple nodes write the same overwrite channel (rule-9). + +Fixes: + +- Rename outputs so overwrite channels have one writer. For adapter-emitted IR this indicates a node-id / output-channel collision in the lowering. + +### Versioning, provenance, and reserved names + +Symptoms: + +- A node id collides with `START`/`END` or begins with `_aiperf` (rule-11). +- `graph.version` is not a known major version (rule-12). +- A non-hand-authored graph is missing `provenance.tool` (error), or still carries the field default `"manual"` (warning — the generating adapter never stamped itself) (rule-13). + +Fixes: + +- Rename the node; set `version` to a known value; have the generating adapter stamp `provenance.tool` as `/`. + +### Timing anchor errors + +Symptoms: + +- An edge sets both `delay_after_predecessor_us` and `delay_after_predecessor_start_us`, or a `START`-sourced edge is start-anchored (rule-54). +- A first-token-anchored edge lacks its `delay_after_predecessor_start_us` dispatch fallback, combines with the completion anchor, or sources at `START` (rule-55). + +Fixes: + +- Keep exactly one anchor per edge; give first-token anchors their start-anchored fallback; use `min_start_delay_us` for absolute offsets from trace start. + +## Parse-time failures versus validation findings + +Not every invalid workload reaches `validate(...)`. + +Native Graph IR parsing raises `GraphParseError` for file and record-shape problems such as malformed JSON/YAML, non-object JSONL records, non-mapping YAML documents, malformed top-level `graph`/`traces` blocks, retired record kinds (`kind: mix`, `kind: subgraph`, or a `graph` record carrying an `endpoints` block), and unknown record kinds. A single-document workload with an unknown top-level key (anything other than `graph`/`traces`, including the retired `mix`/`subgraphs` sections) is rejected with the offending key named and the nearest valid key suggested; a document containing none of those sections is rejected rather than silently discarded. + +The loose decoder (`decode.py`, surfaced as `GraphDecodeError` / `GraphParseError`) additionally rejects: + +- unknown fields on `graph`, node, edge, and trace records (`forbid_unknown_fields` — a typo'd field name fails instead of silently changing what is benchmarked; vendor keys are still preserved under `provenance.extra`), +- a node setting both `prompt` and `messages` (aliases; keep exactly one), +- a prompt-less node carrying `expected.input_tokens` (synth-token prompt fabrication is not supported on the native lowering path), +- non-finite (`inf`/NaN) edge delay values and node `min_start_delay_us` (mirrored by rule-57 for typed-struct producers), +- a non-dict `provenance.extra`. + +The native lowering additionally rejects a first-token-anchored edge with no `delay_after_predecessor_start_us` fallback (the rule-55 shape) at parse time, because the profile pipeline lowers without invoking `validate(...)` and would otherwise treat the edge as a completion edge. + +The parser also enforces these verified rule IDs directly: + +- `rule-19`: a `kind: graph` record must precede trace records. +- `rule-20`: only one `kind: graph` record is allowed. + +Typed model decode failures are wrapped as `GraphParseError` with an `IR error` message that preserves the decoder's field path. + +## Unsupported-feature `NotImplementedError` location convention + +Validator findings are preferred for invalid author input that can be reported as a `ValidationIssue`. Use `NotImplementedError` only for unsupported constructs or guardrails that are outside the static rule set. + +When adding an unsupported-feature gate, follow the repository validator-gate convention: + +```text +: +``` + +The `` prefix must identify the exact author-owned construct. For Graph IR, use the same location style as `ValidationIssue.location`, for example: + +- `graph.nodes.: ` +- `graph.nodes..: ` +- `traces[].: ` + +This keeps unsupported-feature failures actionable without requiring users to grep the workload. Do not raise a bare `NotImplementedError` such as `unsupported graph shape` when the offending node, trace, or field is known. + +## Developer checklist for new rules + +When adding a Graph IR validator rule: + +- Add the rule implementation to `validator.py` as a `_rule_NN_` function. +- Return `list[ValidationIssue]`; do not raise for ordinary authoring errors. +- Include `rule_id`, `severity`, `location`, `message`, and a `suggested_fix` where practical. +- Add the new rule to `validate(...)` in `validator.py`. +- Keep `location` paths human-readable and specific enough to locate the offending workload field. +- Add or update tests that pin both positive and negative behavior against IR the live adapters can actually emit. +- Do not document a new exact rule ID until it exists in source. +- Only add a rule when a live producer can emit the construct it validates (or the rule guards an adapter-lowering contract); rules for hypothetical authoring surface get retired. diff --git a/docs/reference/graph-runtime-troubleshooting.md b/docs/reference/graph-runtime-troubleshooting.md new file mode 100644 index 0000000000..f62d74ac23 --- /dev/null +++ b/docs/reference/graph-runtime-troubleshooting.md @@ -0,0 +1,327 @@ + + +# Graph Runtime Troubleshooting + +Symptom-to-action runbook for the AIPerf graph v1 runtime used by Weka graph replay. + +This page focuses on runtime stalls and dropped-return diagnostics in the path: + +```text +GraphIRReplayStrategy + -> TraceExecutor + -> CreditDispatchAdapter + -> CreditIssuer.issue_graph_credit + -> worker request materialization + -> CreditCallbackHandler graph return observer + -> CreditDispatchAdapter.resolve +``` + +For the architecture view, see `docs/reference/graph-async-dataflow-runtime.md`. + +## First classify the stall + +| Symptom | Most likely layer | First action | +| --- | --- | --- | +| No new requests for a long time, but no error; corpus has captured idle time | Faithful idle-gap replay | Check whether the run has `--benchmark-duration`; see [Idle gaps](#idle-gaps). | +| No credit appears to have been issued for the stuck node | Executor pre-dispatch channel readiness | Enable `AIPERF_GRAPH_EXECUTOR_WATCHDOG_TIMEOUT`; see [Pre-dispatch deadlocks](#pre-dispatch-deadlocks). | +| A credit was issued, but the awaiting node eventually raises `TimeoutError` | Adapter waiter did not receive a matching return | Inspect graph return routing and worker health; see [Dispatch timeout](#dispatch-timeout). | +| Logs mention `unknown waiter key` or `unknown instance_id` | Late, duplicate, or unrouted return | Decide whether it is benign after cancellation or a routing bug; see [Unknown return or waiter](#unknown-return-or-waiter). | +| Worker returns an error for a graph credit, or materialized prompt/body is missing | Worker materialization / graph store / ordinal path | Verify `trace_id`, `node_ordinal`, `phase_variant`, and mmap stores; see [Worker materialization failure](#worker-materialization-failure). | + +## Idle gaps + +### Symptom + +The run appears idle for seconds or minutes with no new console output. There is no traceback, and the process is still alive. + +Typical cases: + +- Weka replay is running without `--benchmark-duration`. +- The input corpus has non-zero `min_start_delay_us`, `delay_after_predecessor_us`, or `delay_after_predecessor_start_us` values. +- `GraphIRReplayStrategy` logs a notice like: the corpus has a recorded inter-turn gap up to `Ns` and no `--benchmark-duration` is set. + +### Why it happens + +`TraceExecutor` faithfully honors recorded graph timing. Edge and node delay gates are not runtime-clamped. In count/session/bare graph runs, the phase can therefore span the slowest admitted trace's recorded wall time. + +This is expected behavior, not by itself a deadlock. + +`AIPERF_GRAPH_DISPATCH_TIMEOUT` does not bound idle gaps. It starts only after a node reaches `CreditDispatchAdapter.dispatch()` and a graph credit is issued or deferred. + +### Actions + +1. If you want a wall-clock bound, rerun with `--benchmark-duration `. When the duration budget elapses, in-flight executors are cancelled and already returned records are kept. +2. If you are doing a diagnostic speed run and do not care about faithful inter-turn timing, set `AIPERF_GRAPH_IGNORE_EDGE_DELAYS=1` to collapse recorded edge and node delays globally. +3. For warmup-only cache-pressure diagnostics, prefer the warmup-specific path when applicable: `--agentic-cache-warmup-duration `. It affects WARMUP pacing only. +4. Do not set a short `AIPERF_GRAPH_EXECUTOR_WATCHDOG_TIMEOUT` merely to shorten legitimate idle gaps. The watchdog wraps the whole executor run and can convert a faithful long wait into a failure. +5. If the process must be stopped, Ctrl+C should finalize and export partial results rather than waiting for every parked idle node to fire. + +## Pre-dispatch deadlocks + +### Symptom + +A trace never reaches the worker for one or more nodes. There is no `CreditDispatchAdapter` timeout because no graph credit was issued for the stuck node. + +Typical causes: + +- A `ChannelRequirement` can never be satisfied. +- Producer accounting is wrong for an AND-fan-in channel. +- A producer that terminated without writing its output channel did not advance the AND-fan-in count (`mark_producer_done` miscount). +- A graph shape has a liveness bug where every candidate entry waits on another node's output. + +### Why it happens + +The executor waits for channel readiness before dispatching an LLM/tool request. `AIPERF_GRAPH_DISPATCH_TIMEOUT` only protects the post-dispatch Future bridge. A node stuck in `VersionedChannelStore.await_inputs()` has no parked adapter Future to time out. + +### Actions + +1. Reproduce with a generous executor watchdog: + + ```bash + AIPERF_GRAPH_EXECUTOR_WATCHDOG_TIMEOUT=600 \ + aiperf ... + ``` + + Use a value comfortably above legitimate recorded idle gaps for the corpus. +2. If the watchdog raises `TimeoutError`, inspect the graph around the last scheduled node frontier: + - declared `inputs` on the blocked node, + - the AND-fan-in count each channel requires vs. the number of live producers that will ever write it, + - producers that exit without writing (a contained `GraphDispatchError` still marks output channels done; verify that path). +3. Keep `AIPERF_GRAPH_DISPATCH_TIMEOUT` unchanged while debugging this class. Raising or lowering it will not affect a pre-dispatch wedge. +4. For CI regression tests, enable the watchdog with a short timeout on intentionally wedged test graphs. Leave the production default unset to preserve faithful unbounded idle-gap replay unless the run already has a duration budget. + +## Dispatch timeout + +### Symptom + +A node reaches the adapter, `issue_graph_credit()` is called, but the awaiting dispatch raises `asyncio.TimeoutError` after `AIPERF_GRAPH_DISPATCH_TIMEOUT`. + +### Why it happens + +`CreditDispatchAdapter.dispatch()`: + +1. resolves the runtime node to a build-time `node_ordinal`, +2. mints a correlation key from `x_correlation_id` and `turn_index`, +3. parks an `asyncio.Future`, +4. issues a graph credit through `CreditIssuer.issue_graph_credit()`, +5. waits for `CreditCallbackHandler`'s unconditional graph-return observer to route the matching return back to `resolve()`. + +A timeout means the parked Future did not receive a matching return in time. + +### Actions + +1. Determine whether the credit was actually placed on the wire. + - If `issue_graph_credit()` refused because the stop gate, duration, request-count cap, or prefill acquisition blocked it, the adapter should raise `GraphDispatchError` promptly, not wait for the full dispatch timeout. + - If it timed out, treat it as an issued-or-deferred credit whose matching return was not observed. +2. Check worker-side logs for request failure, crash, or lost return. +3. Check callback routing: + - `CreditCallbackHandler` must have a graph return observer installed before credits are sent. + - The returned `credit.trace_id` must match a live `GraphIRReplayStrategy` adapter registry key. + - The returned `(credit.x_correlation_id, credit.turn_index)` must match a parked adapter waiter. +4. Search logs for related messages: + - `graph return for unknown instance_id=... dropped` + - `graph return for unknown waiter key ... dropped` (DEBUG level — requires verbose logging) + - `graph dispatch errored: ...` + - `graph dispatch cancelled by worker return` +5. Raise `AIPERF_GRAPH_DISPATCH_TIMEOUT` only when a single valid request can legitimately exceed the default timeout, for example very large prefill plus long generation. Do not use it to mask dropped returns. + +## Unknown return or waiter + +### Symptom + +Logs mention one of these return-routing diagnostics: + +```text +graph return for unknown instance_id=... dropped: no live adapter is registered +``` + +or: + +```text +graph return for unknown waiter key (...) dropped +``` + +### Meaning + +`unknown instance_id` is logged by `GraphIRReplayStrategy._on_graph_return()`. The callback handler delivered a graph credit return, but the strategy had no live adapter registered under `credit.trace_id`. + +`unknown waiter key` is logged by `CreditDispatchAdapter.resolve()`. The adapter was found, but the exact `(x_correlation_id, turn_index)` waiter was absent. + +### When it is benign + +A small number of unknown waiter logs can be expected after cancellation or timeout: + +- the adapter removes orphaned waiters when a dispatch times out or is cancelled; +- a late worker return can arrive after that cleanup; +- an already resolved or duplicate return is dropped. + +### When it is actionable + +Treat the message as a bug if it coincides with dispatch timeouts, missing records, or a phase that fails to drain. + +Actions: + +1. For `unknown instance_id`, verify adapter lifetime: + - `credit.trace_id` should be the per-instance id, for example `trace#lane.recycle`. + - The strategy should register the adapter before `TraceExecutor.run()` and keep it until the parent is done and no in-flight waiters remain. + - The adapter is retained until its instance's executor finishes AND no in-flight waiters remain (`inflight_count == 0`); a return arriving after that reap finds no adapter. +2. For `unknown waiter key`, verify correlation uniqueness and return fidelity: + - `x_correlation_id` includes the per-instance base, runtime trace id, node id, and phase variant. + - `turn_index` is a per-correlation fire counter. + - Worker returns must preserve both fields exactly. +3. If the log appears immediately after phase teardown, verify no old observer remains registered into a new phase. `GraphIRReplayStrategy.teardown_phase()` should detach the observer and clear retained adapters. +4. If the return carries the wrong `trace_id`, distinguish base trace id from instance id. The worker strips `#...` when looking up graph-store content, but the credit return router must preserve the full instance id for adapter de-mux. + +## Worker materialization failure + +### Symptom + +A graph credit reaches the worker, but the worker cannot rebuild the request body or returns an error for that graph node. The adapter converts non-overflow worker errors into `GraphDispatchError`; context-overflow errors are treated as a clean early trajectory termination. + +Related diagnostics can include: + +```text +no node_ordinal for trace=... key=...; credit will carry node_ordinal=None +``` + +or a worker-side error referencing a missing graph-store envelope, missing segment content, an unknown ordinal, or an unexpected phase variant. + +### Why it happens + +Worker-side graph request materialization is addressed by graph metadata on the credit: + +- `credit.trace_id`: the per-instance id used for return routing; graph-store lookup uses the base template trace id after recycle suffix stripping. +- `credit.node_ordinal`: build-time ordinal for the fired node. +- `credit.phase_variant`: `warmup` or `profiling`. + +For Weka segment-trie replay, the worker reconstructs the prompt/body from the unified segment store rather than from `node.output` in the executor. If the ordinal or phase variant is wrong, or the mmap stores are missing/stale, materialization fails before a useful model response can be produced. + +### Actions + +1. Check adapter ordinal resolution warnings; node ids resolve directly against the per-trace catalog. +2. If `node_ordinal=None` appears, inspect the build-time catalog for the base trace and the bare node id. A credit with `node_ordinal=None` is unlikely to materialize into the intended request. +3. Confirm the worker can open the graph store for the current benchmark id: + - the interned unified segment store (`aiperf_graph_segments_`), which every graph build (weka, dynamo, native) writes and the worker opens by on-disk existence. + - It resolves under `AIPERF_DATASET_MMAP_BASE_PATH` (or the system temp dir); in distributed runs that path must be shared with all workers. +4. Confirm phase variant: + - WARMUP uses `phase_variant="warmup"` and applies the warmup output-token cap during materialization. + - PROFILING uses `phase_variant="profiling"`. +5. If the error body is a context length / maximum context / `context_length_exceeded` style error, the adapter classifies it as context overflow and terminates the trajectory cleanly instead of continuing to dispatch later turns that would also overflow. +6. If materialization succeeds but downstream nodes still fail, remember that Weka LLM `node.output` is intentionally a placeholder in the executor. Downstream prompt content comes from the recorded unified-store content-pool handles (or, for slot-carrying nodes, the worker-local dynamic pool), not the live LLM output channel. + +## `did not carry GraphSegmentClientMetadata` / `graph_meta sidecar missing at ...` / `failed the unified-store index cross-check` + +Graph runs hard-fail at timing configure time if the structural sidecar is +not loadable — the DatasetManager broadcasts the store and sidecar locations +on the graph-typed `DatasetConfiguredNotification`, and the TimingManager and +workers use those exact paths (nothing re-parses the workload and nothing +re-derives paths from env conventions). + +- **Not graph-typed:** the broadcast carried no `GraphSegmentClientMetadata` + for a graph run — a build-plane bug or a version-skewed DatasetManager; + check the DatasetManager logs for the `graph_meta sidecar written:` line. +- **Missing:** the advertised path does not exist on this service's + filesystem. In multi-host or containerized deployments, point + `AIPERF_DATASET_MMAP_BASE_PATH` at a filesystem shared by the + DatasetManager, the TimingManager, and all workers. +- **Unreadable:** a truncated or foreign file at the sidecar path — remove + the `aiperf_graph_meta_/` directory and re-run. +- **Index cross-check failure:** the sidecar topology no longer matches the + unified store's manifests (stale artifacts from an earlier build under the + same benchmark id). Remove both `aiperf_graph_meta_/` and + `aiperf_graph_segments_/` and re-run. + +## Graph credit slot and completion checks + +Graph credits intentionally bypass linear session-slot lifecycle accounting: + +- `CreditIssuer.issue_graph_credit()` does not acquire a session slot. +- It still acquires a prefill slot per request. +- `CreditCallbackHandler._release_slots_for_return()` does not release a session slot for graph credits. +- `CreditCounter` does not let graph credits trip `is_final_credit`; graph phase completion is owned by `GraphIRReplayStrategy`. + +Actions when diagnosing apparent phase deadlocks: + +1. Do not look for a final graph credit to set linear completion. The strategy calls `mark_graph_sending_complete()` after executor admission/drain or duration cancellation. +2. If no credits were issued, the strategy sets the graph returned event directly. +3. If credits were issued, the callback handler increments returned counts and signals all-returned once every issued-and-not-cancelled graph credit has returned. +4. Session-slot underflow or session-slot leaks usually indicate a graph credit accidentally went through a non-graph path or lost its `trace_id`. + +## Useful knobs + +### Runtime knobs + +| Knob | Default | Use when | Caution | +| --- | --- | --- | --- | +| `AIPERF_GRAPH_DISPATCH_TIMEOUT` | `300.0` | Bound a node after it issued/deferred a graph credit and is waiting for a return. | Does not detect pre-dispatch channel deadlocks or shorten idle gaps. | +| `AIPERF_GRAPH_EXECUTOR_WATCHDOG_TIMEOUT` | unset | Convert executor frontier deadlocks into `TimeoutError`. | Set generously; it wraps the whole executor run and can trip on legitimate long replay gaps. | +| `AIPERF_GRAPH_IGNORE_EDGE_DELAYS` | `False` | Collapse graph edge and node timing for diagnostic/stress runs. | Not faithful to recorded Weka timing. | +| `--agentic-cache-warmup-duration` (config) | unset | Compress WARMUP delay gates for cache-pressure warmup. | Affects WARMUP pacing only, not PROFILING. | +| `AIPERF_GRAPH_IDLE_GAP_NO_DURATION_WARN_SECONDS` | `30.0` | Tune the no-duration idle-gap advisory threshold. | Advisory only; does not change replay timing. | +| `--trajectory-start-min-ratio` (config) | `0.0` | Lower bound for per-trace t-star sampling. | Window off by default (full replay). `--scenario inferencex-agentx-mvp` auto-applies `0.0`/`1.0`. | +| `--trajectory-start-max-ratio` (config) | `0.0` | Upper bound for per-trace t-star sampling. | Must be greater than or equal to min ratio; `0.0` keeps the window off (full replay). | +| `--random-seed` (config) | unset | Reproduce or sweep t-star snapshot selection. | The SAME seed drives synthesized content, so one seed reproduces the whole run. | +| `--burst-phase-starts` (config) | `False` | Collapse phase-start leading offsets so initial WARMUP/PROFILING frontier fires immediately. | Inter-turn delays after the first are still honored. | +| `AIPERF_GRAPH_WARMUP_MAX_OUTPUT_TOKENS` | `1` | Adjust warmup request generation cap. | Applies to WARMUP boundary turns, not PROFILING. | + +### Dataset and materialization knobs + +| Knob | Default | Use when | +| --- | --- | --- | +| `AIPERF_DATASET_MMAP_BASE_PATH` | unset | Put mmap artifacts (the interned unified store and the graph-meta sidecar) on a shared filesystem for multi-process or distributed workers. | +| `AIPERF_DATASET_WEKA_GRAPH_PARALLEL_WORKERS` | `0` | Force a specific Weka graph parse worker count (`0` auto-sizes). | +| `AIPERF_DATASET_WEKA_HF_SPLIT` | `train` | Bound the streamed rows ingested from a HuggingFace Weka corpus with a slice, e.g. `train[:100]`. | + +There is no store-shape or cache knob: the interned unified store is the sole +trie build path (chosen by graph structure), and there is no cross-run graph +cache. The store-selection and cache env flags described in older +revisions were retired. + +### CLI-level bounds + +| Option | Use when | +| --- | --- | +| `--benchmark-duration ` | Bound faithful idle-gap replay by wall time and cancel in-flight executors at the duration stop condition. | +| `--concurrency ` | Set graph lane fan-out / trace admission pressure. | +| `--num-conversations ` | Cap total admitted root graph instances. | +| `--request-count ` | Cap issued node requests through the graph issue gate. | +| Ctrl+C | Finalize and export partial results when a no-duration idle-gap run is intentionally stopped. | + +## Log messages and next steps + +| Message fragment | Interpretation | Next step | +| --- | --- | --- | +| `idle-gap corpus has a recorded inter-turn gap` | Faithful replay is sleeping through captured gaps without duration bound. | Add `--benchmark-duration`, wait, or Ctrl+C for partial export. | +| `TimeoutError` from `TraceExecutor.run` with watchdog enabled | Pre-dispatch frontier deadlock or watchdog too low for legitimate replay time. | Inspect channel requirements and producer accounting; raise watchdog if gaps are expected. | +| `graph dispatch errored:` | Worker returned an error for an issued graph credit. | Inspect worker error body and materialization inputs. | +| `graph dispatch cancelled by worker return` | Worker reported cancellation for the request. | Check cancellation policy, duration/request-count caps, and worker cancellation logs. | +| `graph credit refused by issuer` | Stop gate, request-count cap, duration/cancel gate, or prefill acquisition refused before a return could arrive. | Treat as clean trace stop unless unexpected; inspect stop condition values. | +| `graph return for unknown waiter key` | Adapter found but specific Future was already gone or never parked. | Benign after timeout/cancel; actionable if paired with missing records or dispatch timeouts. | +| `graph return for unknown instance_id` | Return could not find a live adapter for `credit.trace_id`. | Check adapter registry lifetime (parent-done + zero in-flight reap), phase teardown, and instance id preservation. | +| `no node_ordinal for trace` | Runtime node could not be mapped to a build-time catalog ordinal. | Inspect the per-trace catalog for the base trace id and confirm the t-star snapshot rewrite preserved the fired node's id (a rewritten node id will miss the catalog and yield `node_ordinal=None`). | + +## Quick decision tree + +```text +Is the process alive with no error? + | + +-- Is the corpus known to contain long captured gaps? + | | + | +-- Yes: add --benchmark-duration, wait, or use IGNORE_EDGE_DELAYS for diagnostics. + | +-- No: continue. + | + +-- Was a graph credit issued for the stuck node? + | + +-- No / unknown: enable EXECUTOR_WATCHDOG_TIMEOUT and inspect channel readiness. + | + +-- Yes: did a matching CreditReturn arrive? + | + +-- No: investigate worker crash, lost return, graph observer, adapter registry. + | + +-- Yes but dropped: inspect unknown instance/waiter logs and correlation fields. + | + +-- Yes with error: debug worker materialization or server error body. +``` diff --git a/docs/reference/graph-segment-unified-store.md b/docs/reference/graph-segment-unified-store.md new file mode 100644 index 0000000000..9e843af249 --- /dev/null +++ b/docs/reference/graph-segment-unified-store.md @@ -0,0 +1,305 @@ + + +# Graph Segment Unified Store (Internal Reference) + +Internal developer reference for the **unified segment store**, the sole on-disk +store shape for segment-trie graph builds (weka, dynamo, native, dag_jsonl). It +describes build-path mechanics that are not user-facing. See +`src/aiperf/dataset/graph_segment_unified_store.py` for the implementation. + +## What it is + +A graph build persists everything a worker needs to reconstruct a node's +request into ONE directory. Earlier revisions split this across two on-disk +stores (a separate content pool plus a per-node envelope store); those were +retired and the unified store is now the only graph store shape. Every graph +parse — weka, dynamo, native (via `lower_native_to_unified`), and dag_jsonl — +lowers onto it. + +The unified store folds content and per-node addressing into ONE directory so a +worker opens a single client that carries **both** faces (the envelope +addressing face `get_node_envelope` and the interned content face +`materialize_handles` / `build_request_body_handles`). +`GraphSegmentUnifiedClient` is the ONE object the worker resolves via +`Worker._graph_store_reader` and hands to the `worker_materialize` module's +functions (e.g. `materialize_graph_request_unified`). + +## On-disk layout + +The store lives in a directory named +`aiperf_graph_segments_/` (distinct from the graph-meta sidecar's +`aiperf_graph_meta_/`) and holds four files: + +```mermaid +flowchart TB + subgraph dir["aiperf_graph_segments_<benchmark_id>/"] + direction TB + cb["content.blob
concatenated segment wire blobs
(verbatim wire_json or derived {role,content})"] + ci["content.idx
packed array('Q') span pairs [off,size,...]"] + nb["nodes.blob
concatenated per-node envelope JSON"] + ni["nodes.idx
trace_id -> ordinal:variant -> [offset,size]"] + end + cb -.->|span offset+size| ci + nb -.->|envelope offset+size| ni +``` + +- **`content.blob`** — every unique segment's wire blob serialized once, + concatenated back-to-back. A raw-authored segment (one carrying + `Segment.wire_json`, e.g. a verbatim dag_jsonl message) persists its + `wire_json` VERBATIM — key order and extra keys preserved byte-for-byte; + only `wire_json`-less segments derive the normalized + `orjson.dumps({"role": ..., "content": ...})` blob (see + `GraphSegmentUnifiedBackingStore.put_segment`). Each blob is streamed + straight to this file at `put` time (the incremental spill, below), so the + encoded content is never accumulated in a RAM `bytearray`; the on-disk bytes + and their order are identical to an accumulate-then-flush build. +- **`content.idx`** — the index INTO `content.blob`, a raw `array('Q')` of + 64-bit unsigned integers laid out as `[off0, size0, off1, size1, ...]`, one + span pair per segment handle. The `i`-th segment drained into the pool gets + handle `i`. +- **`nodes.blob`** — every node's envelope, concatenated. Each envelope is + `{"handles": [...], "dispatch_overrides": {...}, "stream": ...}` plus the + optional dynamic-content keys `items` / `capture` for slot-carrying corpora + and the optional `extra_headers` map (per-node HTTP headers, e.g. dynamo's + `x-dynamo-*` session identity) the worker attaches to the request HEADERS + via the synthetic `Turn.extra_headers` — never the body. The stored values + are the RECORDED session ids; at dispatch the worker suffixes + `x-dynamo-session-id` / `x-dynamo-parent-session-id` with the credit's + phase variant and trace-instance suffix + (`worker_materialize.uniquify_dynamo_session_headers`, e.g. + `sess-X` → `sess-X#profiling.0.0`) so concurrent replay instances of one + trace open distinct server-side sessions — parent/child linkage within an + instance is preserved by applying the same suffix to both headers, and + `x-dynamo-session-final` is forwarded untouched (each instance closes only + its own session). Optional keys are omitted when unset, so envelopes for + corpora without them stay byte-identical. +- **`nodes.idx`** — an `orjson` map `trace_id -> ":" -> + [offset, size]` locating each node's envelope in `nodes.blob`. The inner key + is the store's `":"` form (`_encode_inner_key`). + +Blobs are `mmap`ed read-only (`ACCESS_READ`) so all workers share one physical +copy; a zero-length blob maps to `None` rather than an empty mmap. + +## Interned packed-handle format (A2-strict) + +The store is **interned only**: segments are addressed by their +**insertion-index int handle** (not by hex segment id), and `content.idx` is the +raw packed `array('Q')` span table described above — mmap-friendly with no JSON +parse. Node envelopes carry an int `handles` list, and the worker materializes +messages through the int-handle faces `materialize_handles(handles)` / +`build_request_body_handles(handles, ...)`. + +The reader is A2-strict: `GraphSegmentUnifiedClient._load_content_idx` peeks the +first byte of `content.idx`, and a legacy JSON (hex-composition) index — which +begins with `b"{"` — is rejected loudly with a `ValueError` (`"legacy +non-interned unified store (pre-v3) ...; re-parse required"`). There is no +runtime auto-detect between two formats and no dict/hex-id fallback inside the +unified store; an on-disk legacy shape is a re-parse, not a soft fallback. + +The build-time writer `GraphSegmentUnifiedBackingStore` always emits the packed +`content.idx` (its `finalize` writes the flat `array('Q')`); it takes no +`interned` argument. The hex->handle map (`_ids`) lives ONLY in the writer at +build time and never reaches disk or workers. + +## When the unified store is written + +There is no environment flag gating the unified store — it is the sole graph +store shape. `GraphStoreBuilder._build_graph_store_streaming` +(`src/aiperf/dataset/graph/store_build.py`) dispatches on the workload format +to ONE of two drains; both build the SAME on-disk unified store and each writes +its own mandatory content-free `graph_meta` sidecar: + +- **Weka — worker-pool payload stream (`_build_graph_store_streaming_trie`)**: + weka sources (local file, directory, or HF corpus id) stream worker-parsed + `TraceSegmentPayload`s into the unified store via + `build_unified_trie_store_from_payloads`, interned. Each per-row worker + serializes its trace's envelopes so the parent never decodes a full + `ParsedGraph` only to re-serialize the same content; streaming `put_segment` + dedup on the content-addressed id bounds RAM at corpus scale. Each row also + ships a content-free structural graph, which the parent merges + (`_merge_structural_graphs`) into the corpus structural graph that feeds BOTH + the sidecar and the per-node prefix-cache map. `_build_graph_store_streaming` + returns `(catalog, merged_structural)`. +- **Every other format — in-process interned drain + (`_build_interned_unified_store`)**: dynamo, native, and dag_jsonl parse ONCE + in-process (a whole-graph lowering — a dynamo capture lowers each session-tree + on its own node set, with the live write-through store pinning the store build + to the serial in-process path, and dag_jsonl expands whole trees, so the store + build has no per-item parse to fan out) and drain that SAME parse into the interned unified store + via `build_unified_trie_store_interned`. In-process there is no worker pool to + fan out to, so the payload round trip is pure overhead; the interned drain is + also the only one that persists dynamic-slot envelopes (native `@channel` + assembly items/capture, dag_jsonl live-reply lineage), so slot-carrying graphs + ride this route with no separate fallback. The mandatory sidecar is written + DIRECTLY from the stripped whole parse (`_write_graph_sidecar(parsed, ...)`), + so its traces are in PARSE order (the retired weka-style merge sorted them by + id). `_build_graph_store_streaming` returns `(catalog, parsed)` — the full + parse is the prefix-cache source. `build_unified_trie_store_interned` is also + the parity-test oracle for the weka payload-stream drain + (`tests/unit/dataset/test_dynamo_streaming_store_parity.py` and + `tests/unit/dataset/test_dag_jsonl_streaming_store_parity.py` pin + payload-stream == interned store byte-for-byte); the non-weka route face is + pinned by `tests/unit/dataset/test_nonweka_interned_route.py`. + +`graph_carries_assembly_slots` no longer routes the store build (every non-weka +format takes the interned drain regardless); it survives only as the +schedule-plane t\*-gate predicate (`workload_detect._gate_dynamic_slots_vs_tstar`). + +## Incremental content spill + +`put_segment` writes each segment's wire blob straight to an open +`content.blob` handle (opened in `__init__` for buffered binary write) and +advances a running `_content_bytes_written` counter, rather than appending to a +RAM `bytearray` and flushing the whole thing at `finalize`. The span offset a +segment records (`off = self._content_bytes_written` before the write) is +byte-for-byte the value the old `len(self._content_buf)` produced, so +`content.blob` and the packed `content.idx` are identical to the pre-spill +build — the spill is a **write-side-only reshaping** and the read-side +`GraphSegmentUnifiedClient` is untouched. + +Only `_ids` (dedup + handle resolution) and `_spans` (the `content.idx` source) +stay resident on the content side; the encoded content itself is on disk as it +is produced, so `finalize` flushes and closes the handle instead of writing a +`bytes(self._content_buf)` transient double. The handle is plain buffered +binary file I/O with NO event-loop coupling, so the store can be constructed on +the event-loop thread and drained / finalized inside `asyncio.run` on a worker +thread (`GraphStoreBuilder`'s two drains do exactly this). + +`_nodes_buf` is deliberately NOT spilled: the manifest region fills only in the +drain window (below the parse peak) and stays small (~0.2–0.3 GB even at 1M +nodes), so a second write handle would add complexity for no meaningful +peak-RAM win. It remains a RAM `bytearray` flushed once at `finalize`. + +**`abort()` (partial-file cleanup).** Because content is spilled as the drain +runs, a drain that raises before `finalize` would leave a half-written +`content.blob` on disk. `abort()` closes the write handle and unlinks the four +store files; it is idempotent and non-raising, and a `_finalized` guard skips +the unlink after a successful finalize so a complete store is never deleted. +`GraphStoreBuilder`'s interned and streaming drains both wrap their drain in +`try/except → unified.abort() + shutil.rmtree(data_dir) → re-raise`, so a +mid-drain failure leaves no store directory for a later open to trip on. + +## Dynamo direct write-through route (`StoreBackedSegmentPool`) + +The in-process interned drain normally fills the content side in TWO passes: the +adapter parses into an in-RAM `SegmentPool` (`segment_ir/pool.py`, one `Segment` +per unique message segment in `_by_id`), then `build_unified_trie_store_interned` +walks that pool and `put_segment`s each segment into the store. The **dynamo +route skips the second copy**: `GraphStoreBuilder._build_graph_store_streaming` +constructs the `GraphSegmentUnifiedBackingStore` BEFORE the parse and threads it +into `from_dynamo_trace(direct_store=...)`. The adapter then hands +`build_trie_ir` a `StoreBackedSegmentPool` +(`adapters/dynamo/store_backed_pool.py`) instead of a plain `SegmentPool`. + +`StoreBackedSegmentPool.add` computes the identical content-addressed +`segment_id(parent_id, role, tokens)` and calls `store.put_segment(sid, role, +content)` directly, so every prompt/response segment is interned STRAIGHT INTO +`content.blob` at parse time and the pool's `_by_id` stays empty. The store's +`_ids` first-occurrence dedup yields the same handle stream the eager drain +assigns — both intern in `build_trie_ir`'s content-loop first-occurrence order, +the single ordering authority — so the on-disk store is **byte-identical** to the +eager route. This is pinned three-way (direct == eager == streaming) by +`tests/unit/dataset/test_dynamo_streaming_store_parity.py` and by the golden +store-digest (`test_dynamo_store_golden_digest.py`). + +Because the shim's `_by_id` is empty, the returned `ParsedGraph.segment_pool` +no-ops the interned drain's put loop, and `strip_replay_text` replaces it with a +fresh empty `SegmentPool` before the content-free sidecar is msgpack-encoded (the +live shim is never encoded). `add()` is the ONLY real operation — the sole pool +call the dynamo content path makes; `add_text` / `add_raw_message` / `get` / +`materialize` all raise `NotImplementedError` naming the dynamo-only +write-through contract, so any non-dynamo adopter fails loud rather than silently +interning into a pool the store never sees. The shim lives in its own module so +`pool.py` stays a stdlib-only leaf (the store type is referenced only under +`TYPE_CHECKING`). Native and dag_jsonl keep the eager pool→drain path. + +Because the store is live before the parse, the dynamo branch wraps BOTH the +parse and the drain in the same `try/except → abort() + rmtree`, so a mid-parse +failure (e.g. a `DynamoISLMismatchError` on a block-inconsistent record) after +content has already spilled leaves no partial store directory. + +**Memory effect (measured).** The write-through empties the resident content +pool: on a corpus-scale synthetic parse the eager route holds ~17,850 live +`Segment` objects while the direct route holds **0** (measured by +`tests/unit/dataset/graph/adapters/test_dynamo_corpus_scale_memory.py::test_direct_route_content_tier_collapses`). +Freeing those Segments also stops pinning their `content` strings, so the +measured parse peak drops as well. It does NOT remove the per-node +segment-ADDRESSING cost (the hex `prompt_segment_ids` each node retains, allocated +by `segment_id`) — that is a distinct tier present on both routes. + +## Build-memory snapshot and buffer release + +`GraphSegmentUnifiedBackingStore.finalize` brackets its writes with two +build-memory concerns. + +**`build_stats` (finalize ENTRY).** As its FIRST act — before any remaining +file write — `finalize` computes a `GraphStoreBuildStats` snapshot +(`_compute_build_stats`) and stashes it on `self._build_stats`. The snapshot is +a cheap `O(traces) + O(1)` derivation from state the store already holds: +`content_bytes` reads the running `_content_bytes_written` counter (the seam the +`build_stats` docstring used to name as future work, now consumed by the spill), +`manifest_bytes` is a `len(self._nodes_buf)` read (the manifest region is not +spilled), plus one `len()` over the trace map and one `len()` per trace's inner +map — no pass over content. The counter is monotonic, so the finalize-ENTRY +snapshot still measures the full content footprint even though the blob was +streamed to disk during the write. It costs nothing yet turns a +pool/envelope-size regression into a visible build-log delta instead of mystery +RSS at corpus scale. Its counters (`segment_count`, `content_bytes`, +`node_manifest_count`, `manifest_bytes`, `trace_count`) reflect APPENDED +(write-side) totals: a duplicate `(trace, ordinal, variant)` write orphans the +earlier blob in `_nodes_buf` while `node_manifest_count` tracks live index +entries, so a future count/bytes divergence reads as that duplicate-write bug. +`peak_rss_mib` is a `RUSAGE_SELF`-only process peak (`None` on Windows; macOS +reports bytes, Linux KiB, each normalized to MiB) — it excludes pool workers and +is log-only. The `build_stats` property is `None` until `finalize` runs, and +both build-complete log lines +(`GraphStoreBuilder._build_graph_store_streaming_trie` for the weka +payload-stream drain and `_build_interned_unified_store` for the non-weka +in-process interned drain) render it through +the module-level `_format_store_build_stats`, which totals over `None` +(`build_stats=unavailable`) rather than crashing a pre-finalize success line. + +**Write-buffer release (finalize END).** Once every file is flushed, `finalize` +drops all five write-side buffers — `_ids`, `_spans`, `_content_buf`, +`_nodes_buf`, `_node_offsets`. Post-spill the load-bearing clears are `_ids` and +`_spans` (the content-side state that survived the drain); `_content_buf` is +already empty (content spilled at `put` time) so its clear is a formality. The +store object then lives on through the structural-merge / sidecar / prefix-cache +tail with zero readers of those buffers (the snapshot was already taken at +entry), so retaining them would just pin RAM. Because the buffers are gone, the +mutating/reading writer methods (`put_segment`, `add_node_manifest`, +`segment_handle`) all guard on `self._finalized` and raise loudly if called +post-finalize; `segment_handle` in particular raises rather than silently +returning `None` (its id map is released), directing callers to read handles via +`GraphSegmentUnifiedClient` instead. + +## Worker read path + +The worker selects the interned handle path whenever the unified store exists on +disk for the run's `benchmark_id` (existence-based, not flag-driven). It +pre-reads the node envelope once (`read_node_envelope`) and passes it into +whichever materialize function it selects, so the manifest is decoded only once +per credit. See +[Graph Worker Materialization](./graph-worker-materialization.md) for the full +selection logic: + +- `materialize_graph_request_unified` (dict path) reads the envelope and + resolves `handles` through `materialize_handles` (or `_assemble_items` for + slot-splice `items` programs), then applies dispatch overrides / warmup cap / + stream. The worker then layers run-level payload options and cache-bust + markers on the returned dict. +- `materialize_graph_request_unified_bytes` builds the pre-serialized body once + from mmap content-pool slices via `build_request_body_handles`, folding all + outer fields into the overrides tail. Taken only when + `endpoint.cache_bust == CacheBustTarget.NONE` (a pre-serialized body cannot + have a cache-bust marker prepended) AND the envelope carries no dynamic-slot + `items` program (slot-carrying nodes always take the dict path because their + messages are composed per-request from the dynamic pool). + +There is no fallback store: a graph credit whose unified store cannot be opened +fails loudly with a cached fatal `GraphStoreUnavailable` error on every credit +(see `Worker._graph_store_reader`), and a node whose envelope is absent from an +opened store fails with `GraphEnvelopeMissing`. diff --git a/docs/reference/graph-structural-handoff.md b/docs/reference/graph-structural-handoff.md new file mode 100644 index 0000000000..f5e03e90f6 --- /dev/null +++ b/docs/reference/graph-structural-handoff.md @@ -0,0 +1,222 @@ +# Graph Structural Sidecar Handoff (Phase 2) + +This document describes the `graph_meta.msgpack` sidecar — the mandatory +artifact that hands a pre-built structural `ParsedGraph` from the DatasetManager +build plane to the TimingManager schedule plane, so the schedule plane never +re-parses the graph workload file (weka, dynamo, native, or dag_jsonl). The +DatasetManager writes it on every graph build route and advertises its exact +path on the graph-typed `DatasetConfiguredNotification`; the TimingManager +ingests it from that broadcast path only. + +## Two-plane architecture + +A graph IR workload is processed by two independent services in separate +processes: + +| Plane | Service | Responsibility | +|-------|---------|----------------| +| Build | `DatasetManager` (build owned by `GraphStoreBuilder`, `src/aiperf/dataset/graph/store_build.py`) | Build the unified segment store via one of two drains — weka streams worker-parsed trace payloads and merges their per-trace structural graphs; dynamo/native/dag_jsonl parse ONCE in-process and drain that whole-graph parse through the interned builder — then write the mandatory sidecar (from the merged structural graphs on the weka route; DIRECTLY from the stripped parse, in parse order, on the non-weka route). The DatasetManager broadcasts the resulting store and sidecar paths on the graph-typed `DatasetConfiguredNotification` | +| Schedule | `TimingManager` | Ingest the sidecar from the broadcast path (hard fail if the broadcast is not graph-typed or the sidecar is unloadable), hand the `ParsedGraph` to `PhaseOrchestrator` | + +The sidecar bridges these planes: the build plane writes it once and advertises +its location on the broadcast, the schedule plane reads it later from that exact +path — never a second parse of the workload. + +```mermaid +sequenceDiagram + participant DM as DatasetManager + participant B as GraphStoreBuilder + participant FS as Filesystem + participant TM as TimingManager + + DM->>DM: validate_graph_endpoint_type(run) + DM->>B: GraphStoreBuilder(run).build(graph_path) + alt weka_trace (worker-pool payload stream) + B->>B: stream worker-parsed payloads (pool) + B->>FS: drain payloads into unified segment store + B->>B: _merge_structural_graphs(structural_sink) → merged structural ParsedGraph + else dynamo / native / dag_jsonl (in-process interned drain) + B->>B: parse_graph_workload(run, path) → whole-graph ParsedGraph + B->>FS: build_unified_trie_store_interned(parsed) → unified segment store + B->>B: strip_replay_text(parsed) → structural ParsedGraph (parse order) + end + B->>B: catalogs_match(structural, store_catalog)? + alt catalogs match + B->>FS: write_graph_meta_sidecar() → graph_meta.msgpack + B-->>DM: GraphStoreBuildResult(facet, sidecar_path, base_path) + DM->>TM: DatasetConfiguredNotification(GraphSegmentClientMetadata: store + sidecar path) + else catalogs diverge / empty or unmergeable stream + B->>B: raise DatasetError (run fails; the sidecar is mandatory) + end + + TM->>FS: read sidecar_path from GraphSegmentClientMetadata + alt graph-typed, present, decodable, index cross-check passes + FS-->>TM: graph_meta.msgpack bytes + TM->>TM: decode_graph_meta_sidecar() → structural ParsedGraph + TM->>TM: use sidecar graph (never re-parses) + else not graph-typed / missing / undecodable / index-divergent + TM->>TM: raise InvalidStateError (hard configure-time failure) + end +``` + +## The `graph_meta.msgpack` sidecar + +**File location:** `/aiperf_graph_meta_/graph_meta.msgpack` + +resolved by `sidecar_path_for` (the unified segment store lives in its own +`aiperf_graph_segments_/` directory). + +**Producer:** `GraphStoreBuilder._write_graph_sidecar` +(`src/aiperf/dataset/graph/store_build.py`), called on both build routes. The weka payload-stream drain writes it from the merged structural +graphs (see below). Every non-weka format (dynamo, native, dag_jsonl — +slot-free or slot-carrying alike) parses once in-process and takes the interned +drain, which writes the sidecar DIRECTLY from that whole-graph parse (its traces +stay in parse order). Slot-carrying graphs (live-reply dag_jsonl lineage, or +trie assembly items/capture from the native `@channel` lowering) ride this +non-weka route naturally — the interned drain is the only one that persists slot +envelopes, and weka corpora never carry slots. Either way the write is +mandatory: a catalog-divergent sidecar raises `DatasetError` (an unwritable +path fails the build with the underlying I/O error). + +**Consumer:** `TimingManager._load_graph_sidecar` (ingests from the +`GraphSegmentClientMetadata` on the graph-typed `DatasetConfiguredNotification`). + +**Content-free:** The sidecar encodes a **structural** `ParsedGraph` produced by +`strip_replay_text`: graph topology, node ids/types, edges, `arrival_offset_us`, +catalog keys, and the native dispatch fields (model / max_tokens / raw_tools / +extra_headers / theoretical prefix-cache counts) are preserved, but per-trace content is stripped: every +trace's `replay_outputs` (per-node recorded output channel values, +`node_id -> {channel: value}`) is cleared to empty. For the trie IR (`segment_pool is not None`) the strip is deeper: +the `SegmentPool` is emptied (kept non-`None` so the loaded graph still takes the +trie ordinal scheme), and each `LlmNode`'s inline `prompt` and `metadata["trie"]` +contents (`prompt_segment_ids`, raw `hash_ids`) are cleared — only the `"trie"` +marker key is kept. The real content lives in the unified segment store; the +sidecar never duplicates it. + +**Wire format** (`encode_graph_meta_sidecar` / `decode_graph_meta_sidecar` in +`codecs.py`): + +``` +[header, pg_bytes] +``` + +where `header` is +`{"kind": "parsed_graph", "schema_version": int, "source_fingerprint": dict}` +and `pg_bytes` is the canonical `encode_parsed_graph_msgpack` output (a nested +blob). The `kind` discriminator is **required**: `decode_graph_meta_sidecar` +rejects any frame whose `kind` is missing or not `"parsed_graph"` (raising +`ValueError`, which the TimingManager surfaces as a hard `InvalidStateError`), +which is why the schema version was bumped from 2 to 3 — a pre-v3 kind-less +sidecar decodes as invalid and fails the run until it is rebuilt. +`GRAPH_META_SCHEMA_VERSION` is now 4: the 3→4 bump marks the verbatim +raw-JSON `Segment.wire_json` variant, but unlike 2→3 it adds **no reader-side +gate** — the version is advisory provenance only. A pre-v4 blob (which +normalized every segment to `{"role", "content"}`, dropping key order and extra +keys) still decodes fine because `Segment.wire_json` defaults to `None` (a +normalized role/content segment); only the 2→3 `kind`-required transition +gates. The outer +`_SIDECAR_DECODER` is untyped so the frame is decoded +without knowing the inner type; `pg_bytes` is then decoded by the typed +`_PG_MSGPACK_DECODER`. + +## Promote-time catalog cross-check + +Before writing the sidecar, `GraphStoreBuilder` verifies the structural graph's +catalog matches the store's. On the non-weka in-process interned drain (dynamo, +native, dag_jsonl) the structural graph comes from `strip_replay_text` on the +real-content `ParsedGraph` (the same object that built the unified store): + +```python +structural = strip_replay_text(parsed) +if not catalogs_match(structural, catalog): + raise DatasetError("graph_meta sidecar catalog mismatch: ...") +write_graph_meta_sidecar(...) +``` + +On the weka payload-stream drain the same gate runs against the merged +structural graph — assembled from the per-trace content-free graphs emitted +alongside the store payloads by the weka pool workers (each worker calls +`iter_trace_segment_payloads` on its parsed item). + +`catalogs_match` (`graph_meta_sidecar.py`) calls `build_catalog_context` on the +structural graph and compares its `.catalog` dict to the content-build catalog. +Because both come from the same parse(s), a divergence would indicate a bug in +the strip/merge rather than a separate-parse race — the check is a safety +net. A mismatch raises `DatasetError` and fails the run: the sidecar is +mandatory, and one describing a different topology than the stored envelopes +would misschedule the workers. + +## Load-time index cross-check + +When the TimingManager loads an existing sidecar, `_sidecar_passes_index_check` +performs a best-effort verification against the unified store's per-node manifest +index (the store the worker will actually read): + +```python +sidecar_matches_index(graph, index_offsets) +``` + +`sidecar_matches_index` (`graph_meta_sidecar.py`) is a pure comparison: it checks +that every node ordinal in the sidecar's per-trace catalog is present in the +supplied store-index ordinal set. A missing ordinal means the sidecar's topology +diverged from the stored manifests, so it returns `False` and the TimingManager +raises `InvalidStateError`. The surrounding `TimingManager._sidecar_passes_index_check` +handles reachability: any I/O failure while opening the unified store — including +an absent store — is treated as "not reachable" and returns `True`, so the +sidecar is accepted rather than triggering a spurious hard failure. + +## Cache travel: retired + +The cross-run graph cache (and with it the sidecar's promote/install +travel) was removed: every run builds its stores fresh, and the sidecar is +written beside them in the run's own benchmark directory. There is no +persistent cache for the sidecar to travel through. + +## Build drains: weka payload stream vs non-weka interned + +`GraphStoreBuilder._build_graph_store_streaming` +(`src/aiperf/dataset/graph/store_build.py`) dispatches on the workload format +to ONE of two drains; both build the SAME unified store and write the mandatory +sidecar, but they source the structural graph differently. + +**Weka — worker-pool payload stream (`_build_graph_store_streaming_trie`).** Weka +sources — local `.json` file, directory of `.json` files, or HF corpus id — +stream worker-parsed payloads one trace at a time, so the parent never holds a +whole-corpus real-content `ParsedGraph`. Each streamed trace's content-free +structural graph is collected into a `structural_sink` and merged once +(`_merge_structural_graphs`, which raises `DatasetError` on an empty or +unmergeable stream); the merged structural graph feeds both the sidecar and the +per-node prefix-cache map. `_write_graph_sidecar` then writes `graph_meta.msgpack` +after the catalog cross-check (`catalogs_match`; a mismatch raises `DatasetError`). + +**Every other format — in-process interned drain (`_build_interned_unified_store`).** +dynamo, native, and dag_jsonl parse once in-process (whole-graph) and drain that +SAME parse through `build_unified_trie_store_interned`, with no payload round +trip (in-process there is no worker pool to fan out to). The sidecar is written +DIRECTLY from the stripped whole parse (`_write_graph_sidecar(parsed, ...)`), so +its traces are in PARSE order (the weka merge sorts by id), and the full parse +is the per-node prefix-cache source. This drain is also the only one that +persists dynamic-slot envelopes (live-reply dag_jsonl lineage, native trie +assembly items/capture), so slot-carrying graphs ride it with no separate +fallback; `graph_carries_assembly_slots` no longer routes the store build (it +survives as the schedule-plane t\*-gate predicate). + +## Key symbols + +| Symbol | Location | +|--------|----------| +| `strip_replay_text` | `src/aiperf/dataset/graph/graph_meta_sidecar.py` | +| `write_graph_meta_sidecar` | `src/aiperf/dataset/graph/graph_meta_sidecar.py` | +| `sidecar_path_for` | `src/aiperf/dataset/graph/graph_meta_sidecar.py` | +| `catalogs_match` | `src/aiperf/dataset/graph/graph_meta_sidecar.py` | +| `sidecar_matches_index` | `src/aiperf/dataset/graph/graph_meta_sidecar.py` | +| `encode_graph_meta_sidecar` / `decode_graph_meta_sidecar` | `src/aiperf/dataset/graph/codecs.py` | +| `GRAPH_META_SIDECAR_FILENAME` | `src/aiperf/dataset/graph/codecs.py` | +| `_load_graph_sidecar` | `src/aiperf/timing/manager.py` | +| `_sidecar_passes_index_check` | `src/aiperf/timing/manager.py` | +| `GraphSegmentClientMetadata` / `GraphDatasetMetadata` | `src/aiperf/common/models/dataset_models.py` | +| `DatasetManager._configure_graph_dataset` / `_configure_graph_workload` / `_build_graph_store` | `src/aiperf/dataset/dataset_manager.py` | +| `GraphStoreBuilder` / `GraphStoreBuildResult` | `src/aiperf/dataset/graph/store_build.py` | +| `GraphStoreBuilder._write_graph_sidecar` | `src/aiperf/dataset/graph/store_build.py` | +| `GraphStoreBuilder._merge_structural_graphs` | `src/aiperf/dataset/graph/store_build.py` | +| `iter_trace_segment_payloads` | `src/aiperf/dataset/graph/segment_ir/store_builder.py` | diff --git a/docs/reference/graph-trie-prompt-convention.md b/docs/reference/graph-trie-prompt-convention.md new file mode 100644 index 0000000000..18574a1e66 --- /dev/null +++ b/docs/reference/graph-trie-prompt-convention.md @@ -0,0 +1,56 @@ + + +# Graph Trie-Route `prompt=[]` Convention + +Internal developer reference for the inline-prompt convention on the +segment-trie IR route (the `dag_jsonl`, `dynamo`, and `weka` trace adapters). + +## Convention + +Trie-route adapters stamp an EMPTY inline prompt (`LlmNode.prompt == []`). +Prompt content lives ONLY in the run's content-addressed `SegmentPool`, reached +through the node's `metadata["trie"]["prompt_segment_ids"]` path. The inline +`prompt` field stays a required `LlmNode` field (authored native graphs still +carry real inline prompts), but on the trie route it is deliberately left empty. + +Producers (all stamp `prompt=[]`): + +- `src/aiperf/dataset/graph/adapters/dag_jsonl/lowering.py` +- `src/aiperf/dataset/graph/adapters/dynamo/trie_lowering.py` (`build_dynamo_llm_node`) +- `src/aiperf/dataset/graph/adapters/weka/trie_build.py` (`_build_llm_node`) + +## Why + +The inline prompt is dead weight on the trie route: it is one +`{"role", "content"}` dict per prompt message per node -- O(sum of path lengths) +across the graph, held for the entire store build -- while the deduplicated +`SegmentPool` already holds the content once. Nothing on the trie route reads +`node.prompt`: + +- The unified store build (`segment_ir/store_builder.py`) drains only the + segment pool plus the per-node trie envelope (`prompt_segment_ids`, + the native dispatch fields, `stream`). +- The `graph_meta` sidecar (`graph/graph_meta_sidecar.py`, `strip_replay_text`) + forces `prompt=[]` unconditionally. +- The worker materializes prompts from the mmap segment store, not the node. + +## How consumers reach content + +- Build/worker plane: walk `metadata["trie"]["prompt_segment_ids"]` against the + store / `SegmentPool`. +- In-process debugging: `segment_pool.materialize(read_prompt_segment_ids(node))` + (`segment_ir/envelope.py`). + +New trie-route consumers MUST go through the segment path; reading `node.prompt` +on a trie graph yields `[]`. + +## Invariant + +The persisted store bytes are a function of `(segment pool, trie envelope)` only +-- never the inline `node.prompt`. This is pinned by +`tests/unit/dataset/test_dynamo_streaming_store_parity.py::test_store_bytes_independent_of_inline_prompt`, +which builds byte-identical stores through both the eager and streaming drains +from a real-content parse and a sentinel-prompt copy. diff --git a/docs/reference/graph-worker-materialization.md b/docs/reference/graph-worker-materialization.md new file mode 100644 index 0000000000..f24b67ebf0 --- /dev/null +++ b/docs/reference/graph-worker-materialization.md @@ -0,0 +1,453 @@ + + +# Graph Worker Materialization Reference + +Internal developer reference for how a graph-IR credit becomes an inference-server +request on a worker. This page covers the worker-side materialization path, the +runtime identities stamped by the executor bridge, phase and warmup variants, +store opening and error behavior, and how the path relates to the async +`TraceExecutor`. + +Primary implementation files: + +- `src/aiperf/graph/worker_materialize.py` +- `src/aiperf/workers/worker.py` +- `src/aiperf/graph/credit_dispatch_adapter.py` +- `src/aiperf/timing/strategies/graph_ir_replay.py` +- `src/aiperf/dataset/graph_segment_unified_store.py` + +## End-to-end shape + +Graph replay intentionally keeps scheduling and payload reconstruction in separate +planes: + +1. `TraceExecutor` fires graph nodes according to dataflow readiness and recorded + timing. +2. `CreditDispatchAdapter.dispatch()` maps the fired node to a build-time + `node_ordinal`, mints a correlation identity, and issues a graph `TurnToSend` + through `CreditIssuer.issue_graph_credit()`. +3. The normal credit router delivers a `Credit` to any worker. +4. The worker detects `credit.trace_id is not None` and bypasses the linear + sticky-session path. +5. The worker reads mmap-backed graph stores, materializes the node request, sends + it through the normal `InferenceClient`, and returns the credit. +6. The graph return observer resolves the parked adapter `Future`, allowing the + executor node to complete. + +```mermaid +sequenceDiagram + participant E as TraceExecutor + participant A as CreditDispatchAdapter + participant I as CreditIssuer + participant W as Worker + participant M as worker_materialize + participant C as InferenceClient + participant R as Graph return observer + + E->>A: dispatch(node, request, ctx) + A->>A: resolve node_ordinal and mint correlation key + A->>I: issue_graph_credit(TurnToSend) + I->>W: Credit(trace_id, node_ordinal, phase_variant) + W->>M: materialize from graph stores + M-->>W: raw payload dict or pre-serialized bytes + W->>C: send_request(RequestInfo) + W-->>R: CreditReturn(error/cancel/success) + R->>A: resolve parked Future + A-->>E: placeholder string or exception +``` + +## Runtime identity carried by graph credits + +A graph worker request is addressed by three graph fields on `TurnToSend` and +`Credit`: + +| Field | Meaning | Producer | Worker use | +| --- | --- | --- | --- | +| `trace_id` | Runtime trace **instance** id, for example `t-1#0.0`. | `GraphIRReplayStrategy` / `CreditDispatchAdapter` | Return de-mux key and cache-bust marker salt. The worker strips the first `#...` suffix before reading graph stores, because stores are keyed by the base template trace id. | +| `node_ordinal` | Build-time ordinal of the graph node inside the base trace. | `CreditDispatchAdapter._resolve_ordinal()` using `CatalogContext` | Store lookup key for the node envelope. | +| `phase_variant` | Graph materialization variant, currently `profiling` or `warmup`. | `GraphIRReplayStrategy._phase_variant()` | Selects profiling bytes or warmup override behavior. | + +Other request metadata remains normal credit metadata: + +- `x_correlation_id` and `turn_index` are minted by the adapter as the parked + Future key. The id includes the runtime trace id, node id, and phase variant; + `turn_index` is a per-correlation monotonic counter, so two in-flight dispatches + of the same node never share a waiter. +- `conversation_id` and `agent_depth` are metadata labels. Every live producer + lowers to a flat `LlmNode` graph, so `_conversation_identity` returns the bare + instance id; `_dag_identity` supplies the per-node + `(agent_depth, parent_correlation_id)` pair (depth from a dag_jsonl + `metadata["dag"]` stamp, else `0`). These do not change the graph-store + lookup. +- `parent_correlation_id` is carried for DAG metadata, but graph credits bypass + the linear sticky session lifecycle. + +The worker's base-store lookup is: + +```python +base_trace_id = credit.trace_id.split("#", 1)[0] +``` + +For example, `credit.trace_id == "t-1#3.0"` materializes from store key `"t-1"`. +The full instance id is still used for return routing and cache-bust marker +rotation. + +## Materialization inputs + +The worker materializer consumes a node envelope plus, for trie-style stores, a +segment store. + +### Common envelope fields + +All materialization paths apply the node's own outer fields after rebuilding +`messages`: + +| Envelope field | Purpose | +| --- | --- | +| `dispatch_overrides` | Per-node request fields such as `model`, token caps, or provider-specific options. `max_output_tokens` is mapped to the endpoint wire token field. Other keys pass through verbatim. | +| `stream` | Recorded per-node stream flag (weka `"s"` → `True`, `"n"` → `False`; dynamo derived from recorded `ttft_ms`). The recorded per-node mode **wins**: the worker's extraction rule stamps wire `stream` from it, and only a payload that carries NO recorded value falls back to the global `endpoint.streaming` setting. | +| `handles` | Interned unified-store path: ordered integer segment handles that form the full prompt. This is the only trie-store shape; build-time `prompt_segment_ids` are resolved to handles at build time and never reach the worker. | +| `items` | Dynamic-content nodes only (omitted otherwise): the ordered assembly program that interleaves static segment handles with dynamic slots filled at run time. See "Dynamic content slots". | +| `capture` | Dynamic-content producers only (omitted otherwise): `true` when this node's response is spliced into a downstream node, so the worker pools it after dispatch. | + +### Unified store paths + +The unified interned store is the SOLE graph store shape: every graph build +(weka / dynamo / native / dag_jsonl — one streaming store build, with an +eager interned drain (slot fallback) for slot-carrying graphs) writes it, and +the worker reads nothing else. `_graph_unified_reader()` attempts to open one +`GraphSegmentUnifiedClient` for the run's `benchmark_id` on the first credit; the +client carries both the addressing face (`get_node_envelope`) and the content face +(`materialize_handles` / `build_request_body_handles`). There is no build-side +environment flag and no fallback store. The retired legacy ancestor-delta +store path was removed when native graphs started lowering onto the unified +store. + +The unified store is interned-only (A2-strict): every node envelope carries an int +`handles` list. The worker materializes it through the int-handle faces: + +- `materialize_graph_request_unified()` returns a `messages` dict via + `client.materialize_handles(handles)`, for the dict path (cache busting or + run-level options that need a mutable payload). +- `materialize_graph_request_unified_bytes()` builds the pre-serialized body once + from mmap content-pool slices via `client.build_request_body_handles(...)`, + taken when `endpoint.cache_bust == CacheBustTarget.NONE`. + +Both only proceed when the envelope carries `handles`; a `None` return is a +genuine address miss (every interned-store envelope carries `handles`). A legacy +JSON (hex) `content.idx` on disk is rejected loudly by the reader with a +`ValueError` ("re-parse required") rather than silently falling back. + +### Pre-serialized bytes path + +The bytes path (`materialize_graph_request_unified_bytes`) avoids decoding and +re-encoding the `messages` array. It reads the same unified envelope, only +proceeds for nodes carrying `handles`, and builds a complete request body from +mmap slices plus an encoded override tail. Because bytes cannot be mutated after +build, the bytes path folds all outer fields into the body up front: + +- mapped per-node `dispatch_overrides`, +- warmup token cap if applicable, +- the resolved wire `stream` (the recorded per-node value winning, else the + global `endpoint.streaming` fallback for a mode-less payload), endpoint + `extra`, and `stream_options.include_usage`. + +It returns the `(body, model, effective_stream)` triple. The worker stores `body` +on `Turn.raw_payload_bytes` and discards the returned model: `Turn.model` is +deliberately left unset on both the bytes and dict paths. The recorded per-node +model rides only the wire body (sent verbatim), while `record.model_name` falls +back to the run `--model` in `_finalize_request_record`, so tokenizer selection +behaves like plain aiperf — recorded deployment ids are usually not resolvable +tokenizer repos. The `effective_stream` value (the FINAL +stamped wire `stream` mode) is carried onto `RequestInfo.stream_override` so the +transport picks the matching wire mode per-request (`effective_streaming` in +`base_transports.py`). + +The bytes path is deliberately skipped when cache busting is enabled +(`endpoint.cache_bust != CacheBustTarget.NONE`). Cache busting mutates the first +user message content, which a pre-serialized body cannot do. + +## Dispatch overrides and run-level payload options + +Materialization happens in two layers. + +First, `worker_materialize` applies node-local fields: + +- `dispatch_overrides["max_output_tokens"]` maps to `max_tokens` when + `endpoint.use_legacy_max_tokens` is true, otherwise to + `max_completion_tokens`. +- Other dispatch override keys pass through verbatim, including `model` and + provider-specific fields. +- `stream` is stamped from the envelope's recorded `stream` value **only when the + envelope carries one**; a key-absent envelope leaves `stream` unset so the + worker's extraction resolves it to the global fallback (identical to the bytes + path, which reads the raw envelope). + +Second, `Worker._process_graph_credit()` applies run-level endpoint behavior with +`apply_run_level_payload_options()` unless the bytes path already folded it into +the body: + +1. `payload["stream"]` is stamped from the RECORDED per-node stream override when + the payload carries one; a mode-less payload falls back to the global + `endpoint.streaming`. The recorded per-node mode **wins**, so a recorded `"n"` + turn stays non-streaming inside an otherwise-streaming run (and a recorded + `"s"` turn streams even with the global flag off). +2. Each `(key, value)` pair in `endpoint.extra` is merged, with the user-provided + run-level value winning over any per-node key. +3. If the FINAL stamped `stream` is on and `endpoint.use_server_token_count` is + true, `stream_options.include_usage` defaults to `True` only when that key is + absent; an explicit existing `include_usage` value is preserved along with any + other `stream_options` keys. + +This mirrors the normal chat endpoint formatting that graph credits bypass by +sending a raw payload. + +## Dynamic content slots + +Recorded corpora (weka, dynamo) and static native graphs materialize a node's +prompt entirely from build-time content. A hand-authored native graph may +instead reference a predecessor node's **actual response** via an `@channel` +prompt reference whose channel is written by an upstream `LlmNode`. Such a +reference lowers to a dynamic slot; the successor's prompt is composed at run +time from the producers' pooled responses. See the dynamic-content-pool design +spec and the authoring guide (`docs/benchmark-modes/agentic.md`, "Static vs. +dynamic content") for the composition rules. + +### Envelope `items` program + +A slot-carrying node's envelope replaces `handles` with an ordered `items` +program (slot-less envelopes keep `handles` and are byte-identical to recorded +corpora). Each token is one of: + +| Token | Meaning | +| --- | --- | +| `{"h": }` | A static message: the interned segment handle, materialized like the non-dynamic path. | +| `{"s": {"src": }}` | An array-level splice slot: the producer node's pooled reply as a single assistant message — the verbatim recorded assistant message (`tool_calls` preserved) when the capture is structured, else `{"role": "assistant", "content": text}` — or nothing when the producer failed / returned no replayable content (omission). | +| `{"m": {"role": , "parts": [...]}}` | A composed message whose content concatenates `{"t": text}` static parts and `{"sv": }` slot texts; a failed / empty producer substitutes the empty string, so the role and static instruction survive. | + +The build plane (`store_builder._resolve_assembly_items`) resolves the +lowering's producer node ids to node ordinals and hex segment ids to int +handles, so the persisted envelope carries only the worker's native keys. + +An array-level `@channel` messages splice reconstructs the **full +user/assistant alternation** at the read point: the interleaved program emits, +for each upstream writer in completion order, that writer's authored user turn +(static `{"h"}` handles — its "delta", its prompt minus the `@channel` it read) +followed by its reply `{"s"}` slot. So each `{"s"}` slot is preceded by its +producer's user turn, and a downstream reader sees a well-formed conversation +rather than back-to-back assistant messages. The delta handles dedup with the +producer's own prompt interning (content-addressed), so this costs no extra +store bytes. + +### Worker pool and capture + +The worker keeps a per-worker `GraphDynamicPool` keyed +`(trace_id, phase_variant, node_ordinal)` holding each captured response as a +structured `GraphCapturedReply` (the reply's joined text plus, for chat +`tool_calls` / structured replies, the verbatim orjson-serialized assistant +message), `FAILED`, or `EMPTY`. The pool is the graph twin of the linear +path's worker-cached `UserSession` state; content never returns to the timing +plane. + +- **Capture** — after dispatching a `capture: true` node, the worker extracts + the assembled response via `endpoint.build_assistant_turn(record)` and pools + it: a `GraphCapturedReply` on success (a `raw_messages`-bearing Turn — + tool_calls / structured content — also carries the verbatim assistant + message JSON so the splice byte-matches the legacy child-seed rendering), + `EMPTY` for a successful response with no replayable content at all, + `FAILED` on a dispatch error, cancellation, or extraction failure. The pool + write lands strictly before the credit return. +- **Splice** — slot-carrying nodes always take the dict path; + `_assemble_items` composes `messages` from the pool. A `FAILED`/`EMPTY` value + omits (array slot) or substitutes empty text (composed part). A **missing** + value — broken stickiness (worker death re-route) or backstop eviction — sets + a `credit_context.error` with the `aiperf.graph.pool_missing:` prefix, which + the adapter raises as `GraphStickinessError`; the executor treats it as a + non-containable trace-stop (never a silent default). +- **Sticky routing** — dynamic content requires every credit of one trace + instance to reach the same worker. Graph credits key their router session + on `Credit.trace_id` (the instance id), so every trajectory of one + instance shares one session/worker; linear credits keep the legacy + `x_correlation_id` key. The strategy closes the instance session with ONE + explicit `GraphTraceEnd` at adapter-reap, which evicts the worker's pool + entry (deferred while the trace still has in-flight credits). +- **Recorded dynamo identity headers** — replayed with per-instance + uniquification (`uniquify_dynamo_session_headers`) by default. When a + `--session-routing` plugin is active it owns session identity, so the + recorded `x-dynamo-*` identity headers are STRIPPED + (`strip_dynamo_session_headers`) and the plugin stamps live identity at the + request chokepoint instead. +- **Lifecycle backstop** — `AIPERF_GRAPH_DYNAMIC_POOL_MAX_BYTES` + (default 64 MiB/worker) LRU-evicts whole trace entries; evicting a live trace + surfaces as the loud `pool_missing` error above, never a silent truncation. +- **Consecutive-user merge** — omitting a failed producer's assistant turn can + leave adjacent user messages, which some chat APIs reject. + `AIPERF_GRAPH_MERGE_CONSECUTIVE_USER` (default False) merges consecutive + user-role messages in a spliced prompt into one (contents newline-joined). + +The t\* snapshot window is incompatible with dynamic slots (a producer chopped +into warmup would leave its consumer's pool value undefined); a slot workload +loaded with `--trajectory-start-max-ratio > 0` is rejected at parse +(`workload_detect._gate_dynamic_slots_vs_tstar`), the single parse seam — the +DatasetManager is the only production parser; the TimingManager ingests the +broadcast sidecar and never re-parses. + +## Phase variants and warmup overrides + +`phase_variant` is separate from the router-level `CreditPhase` enum. It selects +the graph request variant used by the materializer. + +| Variant | Store lookup | Payload behavior | +| --- | --- | --- | +| `profiling` | Look up `profiling` envelope bytes. | Use recorded messages and recorded dispatch overrides. | +| `warmup` | Reuse `profiling` envelope bytes. | Use the same input prefix, then force a warmup output cap. | +| Other strings | Look up that exact variant. | No special override unless implemented by the caller/store. | + +Warmup intentionally does **not** require duplicate warmup envelopes in the store. +The materializer translates `phase_variant == "warmup"` to a `profiling` lookup: + +```python +lookup_variant = "profiling" if phase_variant == "warmup" else phase_variant +``` + +Then it applies AgentX-style warmup capping after per-node dispatch overrides: + +```python +request.pop("max_completion_tokens", None) +request["max_tokens"] = Environment.GRAPH.WARMUP_MAX_OUTPUT_TOKENS +``` + +`Environment.GRAPH.WARMUP_MAX_OUTPUT_TOKENS` defaults to `1`. Warmup therefore +uses the exact profiling input prefix but forces a one-token output cap. This cap +is independent of accelerated warmup pacing; the pacing knob controls delays in +`GraphIRReplayStrategy`, while the materializer controls payload token limits. + +## Cache busting + +Graph cache busting is applied after dict materialization and run-level options: + +```python +stamp_cache_bust_marker( + payload, + benchmark_id=benchmark_id, + trace_instance_id=trace_instance_id, + target=target, +) +``` + +The marker is deterministic per `(benchmark_id, trace_instance_id)` and is +prepended to the first user message. The full runtime instance id is used, so +recycled instances of the same base trace get distinct markers. A `NONE` target is +a no-op. + +Because this changes message content, cache busting forces the dict path. The +bytes path is only selected when `endpoint.cache_bust == CacheBustTarget.NONE`. + +## Store opening and error modes + +The worker lazily opens graph stores on the first graph credit for a benchmark. +All store paths are resolved from: + +```text +Environment.DATASET.MMAP_BASE_PATH or tempfile.gettempdir() +benchmark_id = worker.run.benchmark_id +``` + +### Store unavailable + +`Worker._graph_store_reader()` resolves the unified client via +`_graph_unified_reader()` (attempted exactly once, +`_graph_unified_open_attempted`, and reused across credits). If the store is +absent or corrupt the worker: + +- creates an `ErrorDetails` with type `GraphStoreUnavailable`, +- records an actionable message naming the expected store directory and mmap base + path (folding in the A2-strict "pre-v3, re-parse required" reason when the + store exists but was rejected), +- caches that error on the worker, and +- returns `None` so no inference request is sent. + +The cached open error prevents a missing shared filesystem or bad base path from +becoming a per-credit retry loop. Content and per-node addressing both live in +the one lazily-opened `GraphSegmentUnifiedClient` returned by +`_graph_store_reader()` (which delegates to `_graph_unified_reader()`); there is +no separate segment-content reader. + +### Node address miss + +If a materializer returns `None` for the addressed node, the worker records a clean +`GraphEnvelopeMissing` error and sends no request. The message includes: + +- base trace id, +- runtime instance id, +- node ordinal, +- phase variant. + +This is distinct from a store-open failure: the store exists, but the requested +node envelope is absent. + +### Materialization faults + +Malformed envelope JSON, corrupt read-time mmap slices, or segment ids/handles +that cannot be resolved are treated as faults rather than graceful misses. They +should be investigated as build/store consistency bugs. + +## Relationship to the executor + +`TraceExecutor` does not build HTTP payloads. Its job is to execute the graph: +wait for channel inputs, honor timing gates, run node dispatch handlers, publish +node outputs, and schedule successors. + +For `LlmNode` dispatches in graph replay: + +1. The dispatch handler builds a `DispatchRequest` and calls the injected + `credit_issuer.dispatch(...)`. +2. In production graph replay, that issuer is a per-instance + `CreditDispatchAdapter`. +3. The adapter maps executor runtime identity to store identity: + - Every live producer lowers to a flat `LlmNode` graph, so the fired node's + bare `node_id` is its catalog key (`_resolve_ordinal`), and the runtime trace + id is always the bare instance id. + - The base `trace_id` keys the catalog and graph stores. + - The instance id keys return routing and cache-bust marker scope. +4. The adapter issues one graph credit and awaits the corresponding return. +5. The worker materializes and sends the payload. +6. The graph return observer resolves the adapter Future. +7. The executor receives a placeholder string on success. Downstream LLM request + bodies are reconstructed from the recorded unified-store content-pool handles, + not from the live LLM output channel — except a hand-authored graph that + splices a predecessor's real response through a dynamic slot (see "Dynamic + content slots"). + +Errors propagate back through the same Future bridge: + +- worker-reported errors become `GraphDispatchError`, except recognized context + overflow, which is converted into expected early termination for the trajectory; +- cancelled returns become `GraphDispatchError`; +- refused issue attempts reject immediately rather than waiting for a return that + will never arrive; +- a missing return is bounded by `Environment.GRAPH.DISPATCH_TIMEOUT` once the + node has reached the adapter. + +## Test coverage map + +Representative tests protecting this behavior: + +| Behavior | Tests | +| --- | --- | +| Graph credits bypass linear sessions and use worker materialization | `tests/unit/graph/test_worker_graph_branch.py` | +| Trie materialization via unified handles | `tests/component_integration/graph/test_weka_trie_e2e_materialize.py` | +| Worker materialization paths (dict / bytes selection) | `tests/unit/graph/test_worker_materialize.py` | +| Unified store and interned handle parity | `tests/unit/graph/test_unified_interned_materialize.py` | +| Bytes path body parity and per-node model preservation | `tests/unit/graph/test_worker_bytes_path.py` | +| Run-level endpoint options and store-open error handling | `tests/unit/graph/test_worker_payload_features.py` | +| Warmup profiling-byte reuse and one-token cap | `tests/unit/graph/test_warmup_variants.py` | +| Adapter identity, correlation, and return handling | `tests/unit/graph/test_credit_dispatch_adapter.py`, `tests/unit/graph/test_graph_return_bridge.py` | +| Per-trace sticky routing and `GraphTraceEnd` lifecycle | `tests/unit/credit/test_graph_sticky_lifecycle.py`, `tests/unit/graph/test_graph_sticky_stamp.py` | +| Dynamic pool lifecycle (deferred eviction, LRU backstop) | `tests/unit/graph/test_dynamic_pool.py` | +| Worker response capture (structured reply / FAILED / EMPTY) | `tests/unit/graph/test_worker_graph_capture.py` | +| Dynamic slot lowering and composition gates | `tests/unit/dataset/graph/test_native_lowering_slots.py` | +| Dynamic slots end-to-end (splice, omission, pool-missing) | `tests/unit/graph/test_dynamic_slots_e2e.py` | diff --git a/docs/reference/yaml-config-roadmap.md b/docs/reference/yaml-config-roadmap.md index a34ce98a47..ed985fabe3 100644 --- a/docs/reference/yaml-config-roadmap.md +++ b/docs/reference/yaml-config-roadmap.md @@ -21,17 +21,17 @@ The v2 envelope is partway between single-config and the multi-phase / multi-dat What works today: -- **Multi-model selection is wired.** `benchmark.models` is a `ModelsAdvanced` block (`src/aiperf/config/models.py:113`) with `items: list[ModelItem]` and a `strategy` field — `round_robin`, `random`, or `weighted`. `modality_aware` is roadmap-only and is not accepted by the current validator. The singular `model:` shorthand is normalized into the items list (`src/aiperf/config/loader/normalizers.py:79-89`). Multi-model in one run is a real feature, not a roadmap item. +- **Multi-model selection is wired.** `benchmark.models` is a `ModelsAdvanced` block (`src/aiperf/config/models.py:107`) with `items: list[ModelItem]` and a `strategy` field — `round_robin`, `random`, or `weighted`. `modality_aware` is roadmap-only and is not accepted by the current validator. The singular `model:` shorthand is normalized into the items list (`src/aiperf/config/loader/normalizers.py:79-89`). Multi-model in one run is a real feature, not a roadmap item. - **`benchmark.phases: [...]`** is a list, validated as a discriminated union over phase types. The singular `phases: { type: ..., ... }` shorthand is normalized to a one-entry list named `profiling` (`src/aiperf/config/loader/normalizers.py:99-103`). Top-level `warmup:` / `profiling:` shorthand is normalized to a `[warmup, profiling]` list. - **Singular `dataset:`** is auto-promoted to a one-entry list with `name: "default"` (`src/aiperf/config/loader/normalizers.py:92-97`). -- **Sweep parameter paths** address phases and datasets by name. Path keying logic lives in `src/aiperf/config/sweep/expand.py`; see the `phases.profiling.` special case at `expand.py:472-477`. +- **Sweep parameter paths** address phases and datasets by name. Path keying logic lives in `src/aiperf/config/sweep/expand.py`; see the `phases.profiling.` special case at `expand.py:495-499`. What does **not** yet hold end-to-end: - **Phase names are fixed.** `BasePhaseConfig.name` is typed as `Literal["warmup", "profiling"]` (`src/aiperf/config/phases.py:71-80`). Multiple phases of the same kind are allowed, but they must reuse one of those two canonical names. Truly user-named phases are not plumbed through credit issuance, the timing manager, the records pipeline, or the report layout. -- **`benchmark.datasets` is hard-capped at one entry.** The field is `list[DatasetConfig]` with `min_length=1, max_length=1` (`src/aiperf/config/config.py:166-177`). The list shape exists only so the same schema can be shared between YAML and the `AIPerfSweep` CRD; the field's own description states "the runtime currently loads exactly one dataset." Multiple-dataset input is rejected at validation time, not at runtime. -- **Per-phase dataset selection is half-scaffolded.** `TimingResolver._validate_fixed_schedule_timing` reads a per-phase dataset via `getattr(phase, "dataset", None) or run.cfg.get_default_dataset_name()` (`src/aiperf/config/resolution/resolvers.py:353-355`), but no `dataset:` field exists on `BasePhaseConfig` yet, so the lookup always falls through to the default. The seam is anticipating a feature that hasn't landed. -- **A phase-vs-dataset compatibility checker exists, but only along two axes.** `check_phase_dataset_compatibility` (`src/aiperf/config/resolution/predicates.py:201-243`) currently rejects only two combinations: a phase that `requires_sequential_sampling` (today, just `fixed_schedule`) against a file dataset that doesn't use sequential sampling, and a phase that `requires_multi_turn` (today, just `user_centric`) against a non-multi-turn file dataset. Other compatibility axes — synthetic-vs-trace for `fixed_schedule`, dataset format mismatches — are not yet enforced here. +- **`benchmark.datasets` is hard-capped at one entry.** The field is `list[DatasetConfig]` with `min_length=1, max_length=1` (`src/aiperf/config/config.py:198-209`). The list shape exists only so the same schema can be shared between YAML and the `AIPerfSweep` CRD; the field's own description states "the runtime currently loads exactly one dataset." Multiple-dataset input is rejected at validation time, not at runtime. +- **Per-phase dataset selection is half-scaffolded.** `TimingResolver._validate_fixed_schedule_timing` reads a per-phase dataset via `getattr(phase, "dataset", None) or run.cfg.get_default_dataset_name()` (`src/aiperf/config/resolution/resolvers.py:361-363`), but no `dataset:` field exists on `BasePhaseConfig` yet, so the lookup always falls through to the default. The seam is anticipating a feature that hasn't landed. +- **A phase-vs-dataset compatibility checker exists, but only along three axes.** `check_phase_dataset_compatibility` (`src/aiperf/config/resolution/predicates.py:235-294`) currently rejects three combinations: a phase that requires a stop condition with none of requests/duration/sessions set against a non-graph dataset (graph workloads infer a single-corpus-pass stop), a phase that `requires_sequential_sampling` (today, just `fixed_schedule`) against a file dataset that doesn't use sequential sampling, and a phase that `requires_multi_turn` (today, just `user_centric`) against a non-multi-turn file dataset. Other compatibility axes — synthetic-vs-trace for `fixed_schedule`, dataset format mismatches — are not yet enforced here. The roadmap items below describe how each of those gaps closes. @@ -150,10 +150,10 @@ benchmark: ### Required wiring -1. **Lift the `max_length=1` cap on `BenchmarkConfig.datasets`** in `src/aiperf/config/config.py:166-177`, replacing the schema-share comment with a real multi-dataset contract. -2. **Add `dataset: ` to `BasePhaseConfig`** so the partial scaffolding at `src/aiperf/config/resolution/resolvers.py:353-355` becomes a real read instead of always falling through to `get_default_dataset_name()`. +1. **Lift the `max_length=1` cap on `BenchmarkConfig.datasets`** in `src/aiperf/config/config.py:198-209`, replacing the schema-share comment with a real multi-dataset contract. +2. **Add `dataset: ` to `BasePhaseConfig`** so the partial scaffolding at `src/aiperf/config/resolution/resolvers.py:361-363` becomes a real read instead of always falling through to `get_default_dataset_name()`. 3. **Validate that every `phase.dataset` resolves** to an entry in `benchmark.datasets`. Use the existing "did you mean?" hinting infrastructure for typos. -4. **Extend `check_phase_dataset_compatibility`** (`src/aiperf/config/resolution/predicates.py:201-243`). Today it only checks `requires_sequential_sampling` (file-dataset sampling strategy) and `requires_multi_turn` (file-dataset format). Add: synthetic-vs-trace mismatches for `fixed_schedule`, dataset-format compatibility per phase type, and any rules that fall out of multi-dataset semantics. The fixed-schedule timing-data check in `TimingResolver._validate_fixed_schedule_timing` (`src/aiperf/config/resolution/resolvers.py:347-362`) can move here once it has a real `phase.dataset` to read. +4. **Extend `check_phase_dataset_compatibility`** (`src/aiperf/config/resolution/predicates.py:235-294`). Today it only checks the required-stop rule (non-graph datasets need an explicit requests/duration/sessions stop), `requires_sequential_sampling` (file-dataset sampling strategy), and `requires_multi_turn` (file-dataset format). Add: synthetic-vs-trace mismatches for `fixed_schedule`, dataset-format compatibility per phase type, and any rules that fall out of multi-dataset semantics. The fixed-schedule timing-data check in `TimingResolver._validate_fixed_schedule_timing` (`src/aiperf/config/resolution/resolvers.py:357-375`) can move here once it has a real `phase.dataset` to read. 5. **Dataset preloading.** Today, the dataset manager prepares one dataset. With multiple datasets in play, prepare each up-front, key shared resources (tokenizer, prompt cache) by dataset name, and stream the right one to the credit issuer per phase. 6. **Reporting.** Per-phase JSON exports already partition by phase; once phases reference distinct datasets, include the dataset name in each phase's metadata block so downstream tools can group by it without re-deriving from the config. diff --git a/fern/docs.yml b/fern/docs.yml index d0f4f4d014..6f16d2d1e1 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -24,6 +24,10 @@ redirects: destination: "/aiperf/" - source: "/aiperf/dev/index.html" destination: "/aiperf/dev/" + - source: "/aiperf/benchmark-modes/graph-workload-benchmarks" + destination: "/aiperf/benchmark-modes/agentic-workload-benchmarks" + - source: "/aiperf/dev/benchmark-modes/graph-workload-benchmarks" + destination: "/aiperf/dev/benchmark-modes/agentic-workload-benchmarks" navbar-links: - type: github diff --git a/pyproject.toml b/pyproject.toml index f031e8d2f1..9c2714836a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -264,4 +264,4 @@ filterwarnings = [ [tool.codespell] skip = "*.pyc,*build*,tests/unit/transports/test_aiohttp_sse.py,tests/integration/assets/canary_reference_inputs.json,src/aiperf/server_metrics/units.py,src/aiperf/api/static/dashboard.html,src/aiperf/config/schema/aiperf-config.schema.json" -ignore-words-list = "timeslice,timeslices,optiona,disjointness,concurency" +ignore-words-list = "timeslice,timeslices,optiona,disjointness,concurency,streming" diff --git a/src/aiperf/api/api_service.py b/src/aiperf/api/api_service.py index c890743c73..4074113e9a 100644 --- a/src/aiperf/api/api_service.py +++ b/src/aiperf/api/api_service.py @@ -151,11 +151,13 @@ async def _start_api_server(self) -> None: self.info(f"AIPerf FastAPI started at {self._base_url}/") self.info( - lambda: " Routes: " - + " | ".join( - r.path - for r in self.app.routes - if hasattr(r, "methods") and r.path not in ("/openapi.json",) + lambda: ( + " Routes: " + + " | ".join( + r.path + for r in self.app.routes + if hasattr(r, "methods") and r.path not in ("/openapi.json",) + ) ) ) diff --git a/src/aiperf/cli.py b/src/aiperf/cli.py index 9436e5be59..26dff086f6 100644 --- a/src/aiperf/cli.py +++ b/src/aiperf/cli.py @@ -32,6 +32,7 @@ def _get_help_text() -> str: app.command("aiperf.cli_commands.analyze_trace:app", name="analyze-trace") app.command("aiperf.cli_commands.chat:app", name="chat") app.command("aiperf.cli_commands.config:app", name="config") +app.command("aiperf.cli_commands.dynamo:app", name="dynamo") app.command("aiperf.cli_commands.profile:app", name="profile") app.command("aiperf.cli_commands.plot:app", name="plot") app.command("aiperf.cli_commands.plugins:app", name="plugins") diff --git a/src/aiperf/cli_commands/_chat_stats.py b/src/aiperf/cli_commands/_chat_stats.py index 5200daa538..05e23686af 100644 --- a/src/aiperf/cli_commands/_chat_stats.py +++ b/src/aiperf/cli_commands/_chat_stats.py @@ -116,6 +116,9 @@ def build_record( start_perf_ns=start_ns, timestamp_ns=timestamp_ns, end_perf_ns=end_ns, + # chat always requests SSE streaming (the payload hardcodes + # stream=True), and streaming-only metrics (TTFT/ITL) gate on this. + streamed=True, ) return ParsedResponseRecord( request=request, diff --git a/src/aiperf/cli_commands/dynamo.py b/src/aiperf/cli_commands/dynamo.py new file mode 100644 index 0000000000..7da373257e --- /dev/null +++ b/src/aiperf/cli_commands/dynamo.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CLI subcommand family for Dynamo agent-trace tooling.""" + +import sys + +from cyclopts import App + +app = App( + name="dynamo", + help=( + "Dynamo agent-trace tooling. Use `aiperf dynamo trace-report` to " + "aggregate metrics from a captured trace." + ), +) + + +@app.default +def dynamo() -> None: + """Dynamo agent-trace tooling namespace.""" + app.help_print([]) + sys.exit(2) + + +app.command("aiperf.cli_commands.dynamo_trace_report:app", name="trace-report") diff --git a/src/aiperf/cli_commands/dynamo_trace_report.py b/src/aiperf/cli_commands/dynamo_trace_report.py new file mode 100644 index 0000000000..f37de162b3 --- /dev/null +++ b/src/aiperf/cli_commands/dynamo_trace_report.py @@ -0,0 +1,475 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CLI subcommand: `aiperf dynamo trace-report `. + +Aggregates Dynamo agent-trace metrics per session (`agent_context.session_id`). + +Streams an `AgentTraceRecord` JSONL/JSONL.gz/segmented trace via +`dynamo.trace_reader.iter_trace_records`, groups `request_end` records by +`agent_context.session_id` (the unit of work in the current +`dynamo.request.trace.v1` schema; `parent_session_id` links subagent sessions +to their parent), and emits per-session percentile aggregates for token +counts, timings, kv_hit_rate, and queue_depth in json/table/csv format. +Replay-only records (no `agent_context`) carry no session identity and are +skipped with a counter. Duplicated `request_end` records (dynamo's dual file +sinks can write the SAME record into two files of one capture dir) are folded +once and counted separately, matching the chain parser's dedup identity. +""" + +from __future__ import annotations + +import csv as _csv +import io +import math +import sys +from collections import defaultdict +from dataclasses import dataclass, field +from pathlib import Path +from typing import Annotated, Any, Literal + +import orjson +from cyclopts import App, Parameter +from rich import box +from rich.console import Console +from rich.table import Table + +from aiperf.dataset.graph.adapters.dynamo.trace import _record_identity +from aiperf.dataset.graph.adapters.dynamo.trace_reader import ( + AgentTraceRecord, + iter_trace_records, +) + +OutFormat = Literal["json", "table", "csv"] + +app = App( + name="trace-report", + help=( + "Aggregate Dynamo agent-trace metrics per session. Streams " + "`request_end` records, groups them by `agent_context.session_id` " + "(`parent_session_id` gives the subagent hierarchy), and prints " + "percentile aggregates for token counts, timings, kv_hit_rate, and " + "queue_depth." + ), +) + + +_NUMERIC_FIELDS: tuple[str, ...] = ( + "input_tokens", + "output_tokens", + "cached_tokens", + "prefill_wait_time_ms", + "prefill_time_ms", + "ttft_ms", + "total_time_ms", + "avg_itl_ms", + "kv_hit_rate", + "kv_transfer_estimated_latency_ms", + "queue_depth", +) + +_PCT_STATS: tuple[str, ...] = ( + "count", + "min", + "p50", + "p90", + "p95", + "p99", + "max", + "mean", +) + + +@dataclass(slots=True) +class _SessionAggregate: + """In-memory accumulator for one `agent_context.session_id` group.""" + + session_id: str + """The session id this group aggregates.""" + parent_session_id: str | None = None + """Parent session id (subagent hierarchy); None for root sessions.""" + parent_session_id_conflict: bool = False + """True when records of this session disagree on their parent session.""" + request_count: int = 0 + """Number of `request_end` records folded into this session.""" + child_session_count: int = 0 + """Number of OTHER aggregated sessions whose parent is this session.""" + model_set: set[str] = field(default_factory=set) + """Distinct model names seen (records with `model: null` contribute none).""" + decode_workers: set[int] = field(default_factory=set) + """Distinct decode worker ids seen.""" + prefill_workers: set[int] = field(default_factory=set) + """Distinct prefill worker ids seen.""" + time_lo: int | None = None + """Earliest `event_time_unix_ms` seen.""" + time_hi: int | None = None + """Latest `event_time_unix_ms` seen.""" + metrics: dict[str, list[float]] = field(default_factory=lambda: defaultdict(list)) + """Raw per-record samples per numeric field, for percentile aggregation.""" + replay_records: int = 0 + """Number of records carrying replay (KV block-hash) metadata.""" + unique_hashes: set[int] = field(default_factory=set) + """Distinct replay input-sequence block hashes seen.""" + + +@dataclass(slots=True) +class SessionTraceReport: + """Per-session aggregation result for one Dynamo agent trace.""" + + rows: list[dict[str, Any]] + """Per-session aggregate dicts, sorted by session_id.""" + skipped_no_agent_context: int + """`request_end` records skipped for lacking `agent_context` (replay-only).""" + duplicate_records: int + """Duplicated `request_end` records skipped (same session_id + request_id). + + Dynamo's dual file sinks can hold the SAME record twice in one capture dir; + the dedup identity matches the chain parser's (`trace._record_identity`). + """ + + +def _percentiles(values: list[float]) -> dict[str, float]: + """Compute count/min/p50/p90/p95/p99/max/mean. Empty input returns empty dict.""" + if not values: + return {} + sorted_vals = sorted(values) + n = len(sorted_vals) + + def pct(p: float) -> float: + # Nearest-rank percentile: the 1-based rank is ceil(p/100 * n). + return sorted_vals[max(0, math.ceil(p / 100.0 * n) - 1)] + + return { + "count": float(n), + "min": sorted_vals[0], + "p50": pct(50), + "p90": pct(90), + "p95": pct(95), + "p99": pct(99), + "max": sorted_vals[-1], + "mean": sum(sorted_vals) / n, + } + + +def _record_parent(agg: _SessionAggregate, parent: str | None) -> None: + """Track the session's parent link, flagging conflicting parent claims.""" + if parent is None: + return + if agg.parent_session_id is None: + agg.parent_session_id = parent + elif agg.parent_session_id != parent: + agg.parent_session_id_conflict = True + + +def _record_workers(agg: _SessionAggregate, rec: AgentTraceRecord) -> None: + """Fold model name and prefill/decode worker ids into the aggregate.""" + model = rec.request.model + if model is not None: + agg.model_set.add(model) + worker = rec.request.worker + if worker is None: + return + if worker.decode_worker_id is not None: + agg.decode_workers.add(worker.decode_worker_id) + if worker.prefill_worker_id is not None: + agg.prefill_workers.add(worker.prefill_worker_id) + + +def _record_metrics(agg: _SessionAggregate, rec: AgentTraceRecord) -> None: + """Fold the record's timestamps, numeric metrics, and replay hashes.""" + ts = rec.event_time_unix_ms + agg.time_lo = ts if agg.time_lo is None else min(agg.time_lo, ts) + agg.time_hi = ts if agg.time_hi is None else max(agg.time_hi, ts) + + for f in _NUMERIC_FIELDS: + v = getattr(rec.request, f, None) + if isinstance(v, (int, float)) and not isinstance(v, bool): + agg.metrics[f].append(float(v)) + + if rec.request.replay is not None: + agg.replay_records += 1 + agg.unique_hashes.update(rec.request.replay.input_sequence_hashes) + + +def _count_children(aggs: dict[str, _SessionAggregate]) -> None: + """Fold child-session counts onto each parent present in the aggregate set.""" + for agg in aggs.values(): + parent = agg.parent_session_id + if parent is None or parent == agg.session_id: + continue + parent_agg = aggs.get(parent) + if parent_agg is not None: + parent_agg.child_session_count += 1 + + +def _to_dict(agg: _SessionAggregate) -> dict[str, Any]: + """Serialize a `_SessionAggregate` into a JSON-safe dict.""" + return { + "session_id": agg.session_id, + "parent_session_id": agg.parent_session_id, + "parent_session_id_conflict": agg.parent_session_id_conflict, + "child_session_count": agg.child_session_count, + "request_count": agg.request_count, + "models": sorted(agg.model_set), + "decode_worker_count": len(agg.decode_workers), + "prefill_worker_count": len(agg.prefill_workers), + "decode_workers": sorted(agg.decode_workers), + "prefill_workers": sorted(agg.prefill_workers), + "time_range_ms": [agg.time_lo, agg.time_hi], + "replay_records": agg.replay_records, + "unique_replay_hashes": len(agg.unique_hashes), + "metrics": {f: _percentiles(vs) for f, vs in agg.metrics.items()}, + } + + +def aggregate_by_session( + path: str | Path, + *, + session_id: str | None = None, + limit: int | None = None, +) -> SessionTraceReport: + """Stream a Dynamo agent-trace and return per-session aggregates. + + Public-API library entry point. The CLI is a thin wrapper. + + Args: + path: Path to `.jsonl`, `.jsonl.gz`, segmented prefix, or directory. + session_id: When set, only records whose `agent_context.session_id` + matches are read (pushed down to the reader's parse-time filter). + limit: When set, stop accumulating new sessions once `limit` distinct + session ids have been seen. Records belonging to already-seen + sessions continue to accumulate; child-session counts only cover + the accumulated set. + + Returns: + A `SessionTraceReport` with per-session rows (sorted by session_id), + the count of skipped context-less (replay-only) records, and the count + of skipped duplicate records. + """ + aggs: dict[str, _SessionAggregate] = {} + skipped_no_agent_context = 0 + duplicate_records = 0 + seen: set[tuple] = set() + for rec in iter_trace_records( + path, event_types={"request_end"}, session_id=session_id + ): + if rec.request is None: + continue + ctx = rec.agent_context + if ctx is None: + skipped_no_agent_context += 1 + continue + # Dynamo's dual file sinks can write the SAME record into two files of + # one capture dir; fold each identity once, mirroring the chain path. + identity = _record_identity(rec) + if identity is not None: + if identity in seen: + duplicate_records += 1 + continue + seen.add(identity) + agg = aggs.get(ctx.session_id) + if agg is None: + if limit is not None and len(aggs) >= limit: + continue + agg = aggs[ctx.session_id] = _SessionAggregate(session_id=ctx.session_id) + agg.request_count += 1 + _record_parent(agg, ctx.parent_session_id) + _record_workers(agg, rec) + _record_metrics(agg, rec) + _count_children(aggs) + return SessionTraceReport( + rows=[_to_dict(aggs[k]) for k in sorted(aggs)], + skipped_no_agent_context=skipped_no_agent_context, + duplicate_records=duplicate_records, + ) + + +def _format_json(report: SessionTraceReport) -> str: + return orjson.dumps( + { + "sessions": report.rows, + "skipped_no_agent_context": report.skipped_no_agent_context, + "duplicate_records": report.duplicate_records, + }, + option=orjson.OPT_INDENT_2, + ).decode("utf-8") + + +def _format_csv(rows: list[dict[str, Any]]) -> str: + """One row per session; columns are flat session facts plus per-metric percentiles.""" + buf = io.StringIO() + base_cols = [ + "session_id", + "parent_session_id", + "parent_session_id_conflict", + "child_session_count", + "request_count", + "models", + "decode_worker_count", + "prefill_worker_count", + "time_range_ms_lo", + "time_range_ms_hi", + "replay_records", + "unique_replay_hashes", + ] + pct_cols = [f"{f}_{stat}" for f in _NUMERIC_FIELDS for stat in _PCT_STATS] + writer = _csv.writer(buf) + writer.writerow(base_cols + pct_cols) + for row in rows: + time_lo, time_hi = row["time_range_ms"] + flat = [ + row["session_id"], + row["parent_session_id"] or "", + row["parent_session_id_conflict"], + row["child_session_count"], + row["request_count"], + ";".join(row["models"]), + row["decode_worker_count"], + row["prefill_worker_count"], + time_lo if time_lo is not None else "", + time_hi if time_hi is not None else "", + row["replay_records"], + row["unique_replay_hashes"], + ] + metrics = row["metrics"] + for f in _NUMERIC_FIELDS: + mvals = metrics.get(f, {}) + for stat in _PCT_STATS: + flat.append(mvals.get(stat, "")) + writer.writerow(flat) + return buf.getvalue() + + +def _format_table(rows: list[dict[str, Any]], console: Console) -> None: + """Render a Rich table with one column-block per metric (p50/p95/mean).""" + if not rows: + console.print("[dim](no request_end records)[/dim]") + return + + table = Table( + show_header=True, + header_style="bold magenta", + box=box.SIMPLE_HEAVY, + title="Dynamo trace aggregate", + ) + table.add_column("session_id", style="cyan", overflow="fold") + table.add_column("parent", style="yellow", overflow="fold") + table.add_column("requests", justify="right") + table.add_column("children", justify="right") + table.add_column("models", overflow="fold") + table.add_column("decode_w", justify="right") + table.add_column("prefill_w", justify="right") + metric_columns = ( + ("ttft_ms", "p50"), + ("ttft_ms", "p95"), + ("total_time_ms", "p50"), + ("total_time_ms", "p95"), + ("kv_hit_rate", "mean"), + ("prefill_wait_time_ms", "p95"), + ("queue_depth", "p95"), + ) + for f, stat in metric_columns: + table.add_column(f"{f}.{stat}", justify="right") + table.add_column("replay/hashes", justify="right") + + for row in rows: + models = ",".join(row["models"]) + cells = [ + row["session_id"], + (row["parent_session_id"] or "-") + + (" [red](mixed)[/red]" if row["parent_session_id_conflict"] else ""), + str(row["request_count"]), + str(row["child_session_count"]), + models, + str(row["decode_worker_count"]), + str(row["prefill_worker_count"]), + ] + for f, stat in metric_columns: + v = row["metrics"].get(f, {}).get(stat) + cells.append("-" if v is None else _fmt_num(v)) + cells.append(f"{row['replay_records']}/{row['unique_replay_hashes']}") + table.add_row(*cells) + console.print(table) + + +def _fmt_num(v: float) -> str: + if abs(v) >= 100: + return f"{v:.1f}" + if abs(v) >= 1: + return f"{v:.2f}" + return f"{v:.4f}" + + +@app.default +def dynamo_trace_report( + path: Annotated[ + Path, + Parameter( + help=( + "Path to a Dynamo agent-trace `.jsonl`, `.jsonl.gz`, segmented " + "prefix, or directory of segments." + ) + ), + ], + *, + out_format: Annotated[ + Literal["json", "table", "csv"], + Parameter( + name="--format", + help="Output format. Default: table.", + ), + ] = "table", + session_id: Annotated[ + str | None, + Parameter( + name="--session-id", + help="Restrict to records matching this agent_context.session_id.", + ), + ] = None, + limit: Annotated[ + int | None, + Parameter( + name="--limit", + help="Stop accumulating new sessions once N distinct session ids are seen.", + ), + ] = None, +) -> None: + """Aggregate Dynamo agent-trace metrics per agent_context.session_id.""" + report = aggregate_by_session(path, session_id=session_id, limit=limit) + + if out_format == "json": + sys.stdout.write(_format_json(report)) + sys.stdout.write("\n") + return + if out_format == "csv": + sys.stdout.write(_format_csv(report.rows)) + if report.skipped_no_agent_context: + sys.stderr.write( + f"skipped {report.skipped_no_agent_context} request_end records " + "without agent_context (replay-only)\n" + ) + if report.duplicate_records: + sys.stderr.write( + f"skipped {report.duplicate_records} duplicate request_end " + "records (dual file sinks)\n" + ) + return + console = Console() + _format_table(report.rows, console) + if report.skipped_no_agent_context: + console.print( + f"[dim](skipped {report.skipped_no_agent_context} request_end " + "records without agent_context -- replay-only)[/dim]" + ) + if report.duplicate_records: + console.print( + f"[dim](skipped {report.duplicate_records} duplicate request_end " + "records -- dual file sinks)[/dim]" + ) + + +__all__ = [ + "SessionTraceReport", + "aggregate_by_session", + "app", + "dynamo_trace_report", +] diff --git a/src/aiperf/cli_runner/_aggregate.py b/src/aiperf/cli_runner/_aggregate.py index 3d36a3ef77..46e76ddd13 100644 --- a/src/aiperf/cli_runner/_aggregate.py +++ b/src/aiperf/cli_runner/_aggregate.py @@ -14,7 +14,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Final if TYPE_CHECKING: from pathlib import Path @@ -22,9 +22,52 @@ from aiperf.common.aiperf_logger import AIPerfLogger from aiperf.config import BenchmarkPlan from aiperf.orchestrator.aggregation.base import AggregateResult + from aiperf.orchestrator.models import RunResult from aiperf.orchestrator.strategies import ExecutionStrategy +SCENARIO_SUBMISSION_CARRIER_KEYS: Final[frozenset[str]] = frozenset( + { + "_scenario_name", + "_validator_submission_valid", + "_validator_submission_invalid_reasons", + "_total_responses", + "_context_overflow_count", + "_was_cancelled", + } +) +"""The underscore-prefixed carrier keys ``_stamp_scenario_submission_metadata`` +writes onto ``AggregateResult.metadata``. + +Single source of truth shared by the writer (this module) and every aggregate +exporter that must keep the keys out of its public output: the +``AggregateConfidenceJsonExporter`` pops them individually while folding the +submission verdict; the ``AggregateConfidenceCsvExporter`` (which performs no +fold) strips them via :func:`strip_scenario_submission_carrier_keys`. +""" + + +def strip_scenario_submission_carrier_keys(metadata: dict[str, Any]) -> dict[str, Any]: + """Return a copy of ``metadata`` without the scenario-submission carrier keys. + + Exporters that do not fold the submission verdict themselves (e.g. the + aggregate confidence CSV) use this so the internal carrier keys never leak + into user-facing output. + + Args: + metadata: The ``AggregateResult.metadata`` dict (not mutated). + + Returns: + A shallow copy with every :data:`SCENARIO_SUBMISSION_CARRIER_KEYS` + entry removed. + """ + return { + key: value + for key, value in metadata.items() + if key not in SCENARIO_SUBMISSION_CARRIER_KEYS + } + + async def aggregate_and_export( results: list, plan: BenchmarkPlan, @@ -52,6 +95,7 @@ async def aggregate_and_export( aggregation = ConfidenceAggregation(confidence_level=plan.confidence_level) aggregate_result = aggregation.aggregate(results) aggregate_result.metadata["cooldown_seconds"] = plan.cooldown_seconds + _stamp_scenario_submission_metadata(aggregate_result, results, plan) aggregate_dir = strategy.get_aggregate_path(base_dir) @@ -86,6 +130,172 @@ async def aggregate_and_export( print_aggregate_summary(aggregate_result, logger) +def _stamp_scenario_submission_metadata( + aggregate: AggregateResult, + results: list[RunResult], + plan: BenchmarkPlan, +) -> None: + """Stamp scenario-submission carrier keys onto ``aggregate.metadata``. + + The ``AggregateConfidenceJsonExporter`` + reader (``aggregate_confidence_json_exporter.py`` ``_build_submission_metadata``) + pops these underscore-prefixed keys to re-derive ``submission_valid`` / + ``submission_invalid_reasons`` via ``compute_submission_outcome``; without this + writer the reader sees no ``_scenario_name`` and emits an empty verdict. + + Stamps the carrier keys the reader expects EXACTLY: + ``_scenario_name`` / ``_validator_submission_valid`` / + ``_validator_submission_invalid_reasons`` / ``_total_responses`` / + ``_context_overflow_count`` / ``_was_cancelled`` -- the + :data:`SCENARIO_SUBMISSION_CARRIER_KEYS` set, which non-folding exporters + strip via :func:`strip_scenario_submission_carrier_keys`. + + The validator (lock-level) outcome is read from the REAL per-run + ``ScenarioOutcome`` carried on ``RunResult.submission_valid`` / + ``RunResult.submission_invalid_reasons`` (stamped by ``apply_scenario`` in + the parent via ``local_executor._carried_scenario_outcome``). This + is the lock-only verdict -- ``(True, [])`` for a clean lock (INCLUDING a clean + config that merely carries ``--unsafe-override`` with no violation), or + ``(False, [...])`` under ``--unsafe-override`` with violations. It is NOT + re-derived from the ``unsafe_override`` flag (that flag alone never determines + validity) and NOT read from the per-run JSON's ``submission_valid`` (already + runtime-folded). Runtime-only signals (cross-run context-overflow rate and + cancellation) are folded on top by the reader. No-op when no scenario is set. + + Args: + aggregate: The confidence ``AggregateResult`` whose ``metadata`` dict is + mutated in place with the carrier keys. + results: All per-run ``RunResult`` objects from the orchestrator. Only + successful runs contribute to the cross-run response sums; the + cancellation OR folds over ALL runs (a graceful Ctrl+C that completed + zero requests is ``success=False`` but still wrote ``was_cancelled``). + plan: The ``BenchmarkPlan`` carrying the static scenario name on its + first config. + """ + config = plan.configs[0] + scenario_name = getattr(config, "scenario", None) + if scenario_name is None: + return + + validator_submission_valid, validator_reasons = _carried_validator_outcome(results) + + successful_runs = [r for r in results if r.success] + total_responses, context_overflow_count = _sum_runtime_response_counts( + successful_runs + ) + json_name = config.artifacts.profile_export_json_file.name + was_cancelled = any(_read_run_was_cancelled(run, json_name) for run in results) + + carriers: dict[str, Any] = { + "_scenario_name": str(scenario_name), + "_validator_submission_valid": validator_submission_valid, + "_validator_submission_invalid_reasons": validator_reasons, + "_total_responses": total_responses, + "_context_overflow_count": context_overflow_count, + "_was_cancelled": was_cancelled, + } + if set(carriers) != SCENARIO_SUBMISSION_CARRIER_KEYS: + raise RuntimeError( + "scenario-submission carrier keys drifted from " + "SCENARIO_SUBMISSION_CARRIER_KEYS; update the constant so exporters " + "keep stripping every stamped key" + ) + aggregate.metadata.update(carriers) + + +def _carried_validator_outcome( + results: list[RunResult], +) -> tuple[bool, list[str]]: + """Resolve the lock-only validator verdict carried on the per-run results. + + Reads ``RunResult.submission_valid`` / ``RunResult.submission_invalid_reasons`` + -- the real ``ScenarioOutcome`` stamped by ``apply_scenario`` in the parent + and carried through ``local_executor``. The lock verdict is config- + deterministic across the trials of a single cell (every run shares the same + config), so the first run carrying a non-None ``submission_valid`` defines it. + Falls back to ``(True, [])`` when no run carries a verdict, so the + reader's runtime fold still applies on a clean lock. + + Args: + results: All per-run ``RunResult`` objects from the orchestrator. + + Returns: + A ``(validator_submission_valid, validator_reasons)`` tuple. The bool is + the lock-only verdict (True for a clean lock, including clean + + ``--unsafe-override``); the list is the lock-only reason tags. + """ + for run in results: + if run.submission_valid is not None: + return run.submission_valid, list(run.submission_invalid_reasons) + return True, [] + + +def _sum_runtime_response_counts( + successful_runs: list[RunResult], +) -> tuple[int, int]: + """Sum total responses and context-overflow counts across successful runs. + + Each ``RunResult.summary_metrics`` is a tag -> ``JsonMetricResult`` map; the + per-run total lives on the count metric's ``avg``. Total responses = + ``request_count`` (valid) + ``error_request_count`` (non-overflow failures) + + ``context_overflow_count`` (overflow records skipped from normal metrics), + matching the InferenceX AgentX spec §4.8 / §7 denominator (all responses + received). Returns ``(0, 0)`` when no successful runs exist. + + Args: + successful_runs: The successful ``RunResult`` objects to sum over. + + Returns: + A ``(total_responses, context_overflow_count)`` tuple of non-negative ints. + """ + total_responses = 0 + context_overflow_count = 0 + for result in successful_runs: + metrics = result.summary_metrics or {} + for tag in ("request_count", "error_request_count"): + metric = metrics.get(tag) + if metric is not None and metric.avg is not None: + total_responses += int(metric.avg) + overflow_metric = metrics.get("context_overflow_count") + if overflow_metric is not None and overflow_metric.avg is not None: + overflow_count = int(overflow_metric.avg) + context_overflow_count += overflow_count + total_responses += overflow_count + return total_responses, context_overflow_count + + +def _read_run_was_cancelled(run: RunResult, json_name: str) -> bool: + """Read the top-level ``was_cancelled`` flag from a run's profile export JSON. + + A run that exits 0 after a graceful Ctrl+C still writes its export file (with + partial metrics) and marks it ``was_cancelled: true``; scenario submissions + must treat such runs as invalid. ``RunResult`` carries no cancellation + flag, so the signal is read back from the + per-run JSON at ``run.artifacts_path / json_name``. + + Args: + run: The ``RunResult`` whose ``artifacts_path`` locates the export dir. + json_name: The profile-export JSON filename (honors ``--profile-export-prefix``). + + Returns: + True when the run was cancelled early, False otherwise (including when the + artifacts path or export file is missing or unreadable). + """ + import orjson + + artifacts_path = run.artifacts_path + if artifacts_path is None: + return False + json_file = artifacts_path / json_name + if not json_file.exists(): + return False + try: + data = orjson.loads(json_file.read_bytes()) + except (OSError, ValueError, orjson.JSONDecodeError): + return False + return bool(data.get("was_cancelled", False)) + + def _maybe_compute_detailed( plan: BenchmarkPlan, results: list ) -> AggregateResult | None: diff --git a/src/aiperf/common/base_component_service.py b/src/aiperf/common/base_component_service.py index 04f5eb6e0c..8fd705438a 100644 --- a/src/aiperf/common/base_component_service.py +++ b/src/aiperf/common/base_component_service.py @@ -72,7 +72,9 @@ async def _heartbeat_task(self) -> None: async def _register_service_on_start(self) -> None: """Register the service with the system controller on startup.""" self.debug( - lambda: f"Attempting to register service {self} ({self.service_id}) with system controller" + lambda: ( + f"Attempting to register service {self} ({self.service_id}) with system controller" + ) ) result = None command_message = RegisterServiceCommand( @@ -103,7 +105,9 @@ async def _register_service_on_start(self) -> None: ) else: self.debug( - lambda: f"Service {self.service_id} registered with system controller" + lambda: ( + f"Service {self.service_id} registered with system controller" + ) ) break diff --git a/src/aiperf/common/base_service.py b/src/aiperf/common/base_service.py index fb07311aea..5d2e29f073 100644 --- a/src/aiperf/common/base_service.py +++ b/src/aiperf/common/base_service.py @@ -181,7 +181,9 @@ async def _kill(self) -> None: ) except Exception as publish_error: self.debug( - lambda e=publish_error: f"Failed to publish BaseServiceErrorMessage during _kill (comms may already be down): {e!r}" + lambda e=publish_error: ( + f"Failed to publish BaseServiceErrorMessage during _kill (comms may already be down): {e!r}" + ) ) self.stop_requested = True self.stopped_event.set() diff --git a/src/aiperf/common/clock.py b/src/aiperf/common/clock.py new file mode 100644 index 0000000000..291a5bb415 --- /dev/null +++ b/src/aiperf/common/clock.py @@ -0,0 +1,364 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Time abstraction for AIPerf control-flow components. + +AIPerf historically reads time directly via ``time.perf_counter_ns()`` / +``time.monotonic_ns()`` / ``asyncio.sleep()`` everywhere a load loop, arrival +pacer, ramp, timeout, or trace-replay firing gate consults the clock. That +ties the control flow to wall-clock time, which is correct for live HTTP +transports but undermines two things: + +* **Virtual-time validation.** The graph ``TraceExecutor`` replays a recorded + trace by sleeping each node's incoming firing-edge delay. Driven by wall + time, validating that the executor reproduces the recorded timeline of a + multi-hour agentic trace would itself take multiple hours. Driven by a + virtual clock advanced by an external pump, the same replay completes in + milliseconds and is deterministic. +* **Deterministic wake order.** Sleepers parked on a virtual clock must wake + in a fixed, reproducible order when sim time crosses their deadlines, so a + control flow that samples ``now_ns()`` right after a wake is reproducible. + +The fix is a clock abstraction that every time-sensitive component consults +instead of reading ``time.*`` directly. In wall mode it defers to the standard +library (behavior unchanged); in virtual mode it is advanced by an external +driver. + +Usage: + +.. code-block:: python + + from aiperf.common.clock import AIPerfClock, WallClock, VirtualClock + + async def pace(clock: AIPerfClock, rate_hz: float) -> None: + interval_ns = int(1e9 / rate_hz) + while True: + await clock.sleep_ns(interval_ns) + do_one() + +The default clock is a :class:`WallClock`; behavior is identical to reading +``time.*`` directly. A :class:`VirtualClock` is installed by validation +harnesses that pump ``advance_to`` to fast-forward sim time to the next parked +waiter whenever the event loop goes idle. +""" + +from __future__ import annotations + +import asyncio +import heapq +import time +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Protocol, runtime_checkable + +__all__ = ["AIPerfClock", "WallClock", "VirtualClock"] + + +@runtime_checkable +class AIPerfClock(Protocol): + """Time source for AIPerf control flow. + + Three unit conventions on every clock -- pick the one that matches the + call site so there is no per-call boilerplate conversion: + + * **Seconds** (default, no suffix): ``now()``, ``sleep(s)``, + ``sleep_until(s)``. Float seconds. + * **Nanoseconds** (``_ns`` suffix): ``now_ns()``, ``sleep_ns(ns)``, + ``sleep_until_ns(ns)``. Int. Matches ``time.perf_counter_ns()`` / + ``time.monotonic_ns()``; use when integrating with code that already + passes ns. + * **Milliseconds** (``_ms`` suffix): ``now_ms()``, ``sleep_ms(ms)``, + ``sleep_until_ms(ms)``. Float. Useful when interfacing with timeout + configs / SLA thresholds expressed in ms. + + Two concrete implementations: + + * :class:`WallClock` -- defers to ``time.perf_counter_ns()`` + + ``asyncio.sleep()``. The default for live replay/transports. + * :class:`VirtualClock` -- virtual clock; ``advance_to`` is called by an + external driver pump to fast-forward sim time. Sleepers park on a + per-waiter ``asyncio.Event`` and wake when sim time crosses their + deadline. + + Implementations are NOT required to be thread-safe across the read + methods because AIPerf's control-flow path runs on a single asyncio event + loop. + """ + + # ----- Seconds (default) ----- + + def now(self) -> float: + """Return the current time in float seconds.""" + ... + + async def sleep(self, duration_s: float) -> None: + """Sleep for ``duration_s`` seconds. + + ``duration_s <= 0`` returns immediately (matches ``asyncio.sleep`` + semantics). + """ + ... + + async def sleep_until(self, deadline_s: float) -> None: + """Sleep until the clock reaches ``deadline_s`` seconds (absolute).""" + ... + + # ----- Nanoseconds ----- + + def now_ns(self) -> int: + """Return the current time in nanoseconds.""" + ... + + async def sleep_ns(self, duration_ns: int) -> None: + """Sleep for ``duration_ns`` nanoseconds.""" + ... + + async def sleep_until_ns(self, deadline_ns: int) -> None: + """Sleep until the clock reaches ``deadline_ns`` (absolute).""" + ... + + # ----- Milliseconds ----- + + def now_ms(self) -> float: + """Return the current time in float milliseconds.""" + ... + + async def sleep_ms(self, duration_ms: float) -> None: + """Sleep for ``duration_ms`` milliseconds.""" + ... + + async def sleep_until_ms(self, deadline_ms: float) -> None: + """Sleep until the clock reaches ``deadline_ms`` (absolute).""" + ... + + +class _UnitConversionsMixin: + """Provide the seconds + ms layers in terms of the ns ones. + + Concrete clocks inherit from this so they only implement the ns layer. + """ + + __slots__ = () + + # Seconds layer. + def now(self) -> float: + return self.now_ns() / 1_000_000_000 # type: ignore[attr-defined] + + async def sleep(self, duration_s: float) -> None: + await self.sleep_ns(int(duration_s * 1_000_000_000)) # type: ignore[attr-defined] + + async def sleep_until(self, deadline_s: float) -> None: + await self.sleep_until_ns(int(deadline_s * 1_000_000_000)) # type: ignore[attr-defined] + + # Milliseconds layer. + def now_ms(self) -> float: + return self.now_ns() / 1_000_000 # type: ignore[attr-defined] + + async def sleep_ms(self, duration_ms: float) -> None: + await self.sleep_ns(int(duration_ms * 1_000_000)) # type: ignore[attr-defined] + + async def sleep_until_ms(self, deadline_ms: float) -> None: + await self.sleep_until_ns(int(deadline_ms * 1_000_000)) # type: ignore[attr-defined] + + +class WallClock(_UnitConversionsMixin): + """Default real-time clock -- defers to the standard library. + + Stateless; one instance can be shared across all consumers. + """ + + __slots__ = () + + def now_ns(self) -> int: + return time.perf_counter_ns() + + async def sleep_ns(self, duration_ns: int) -> None: + if duration_ns <= 0: + # Yield once even for non-positive durations (asyncio.sleep + # semantics) so behind-schedule pacing loops cannot starve the + # event loop. + await asyncio.sleep(0) + return + await asyncio.sleep(duration_ns / 1e9) + + async def sleep_until_ns(self, deadline_ns: int) -> None: + await self.sleep_ns(deadline_ns - self.now_ns()) + + +@dataclass(slots=True, order=True) +class _ClockSleeper: + """Heap entry for a sleeper parked on a :class:`VirtualClock`.""" + + deadline_ns: int + """Absolute sim-time deadline the sleeper waits for (primary heap key).""" + + insertion_id: int + """Monotonic tie-break so equal deadlines wake in registration order.""" + + event: asyncio.Event = field(compare=False) + """Per-waiter wake event set by ``advance_to`` when the deadline crosses.""" + + alive: bool = field(default=True, compare=False) + """False once the sleeper exited (woke or was cancelled). Dead entries are + reaped lazily so a cancelled sleeper's phantom deadline cannot fast-forward + sim time via ``peek_min_waiter_ns``.""" + + +class VirtualClock(_UnitConversionsMixin): + """Virtual clock advanced by an external driver pump. + + ``advance_to(ns)`` is monotonic -- calls with ``ns <= now_ns()`` are + silently ignored, matching the semantics of monotonically-progressing sim + time. Each parked sleeper owns its own ``asyncio.Event`` and is keyed in a + ``(deadline, insertion_id)`` min-heap, so ``advance_to`` wakes sleepers in + strict ``(deadline, registration-order)`` priority rather than via + ``Condition.notify_all`` -- which would let asyncio dispatch ready + callbacks in arbitrary order and bake nondeterminism into any control flow + that samples ``now_ns()`` right after a wake. + + Waiter tracking: every parked sleeper registers a + :class:`_ClockSleeper` ``(deadline_ns, insertion_id, event, alive)`` in the + heap. ``advance_to`` pops every entry with ``deadline <= ns`` in heap + order, setting each live entry's event in turn -- the asyncio loop then + schedules the awaiting coroutines via ``call_soon`` in the same order. + Deterministic wake order downstream. A cancelled sleeper marks its entry + dead (``alive=False``); dead entries are skipped by ``advance_to`` and + reaped lazily by ``has_waiters`` / ``peek_min_waiter_ns`` so phantom + deadlines never fast-forward sim time. + + A driver pump fast-forwards sim time when the event loop goes idle: it + reads :meth:`peek_min_waiter_ns` and calls :meth:`advance_to`. Use + :meth:`set_on_waiter_parked` to make idle handling level-triggered (the + callback fires the instant a fresh waiter parks, so the pump never drops a + kick). + """ + + __slots__ = ( + "_now_ns", + "_lock", + "_waiters", + "_insertion_counter", + "_on_waiter_parked", + ) + + def __init__(self) -> None: + self._now_ns: int = 0 + # Protects the waiter heap. ``asyncio.Lock`` (not Condition) -- we use + # per-waiter ``asyncio.Event`` for wake instead of a shared condvar so + # wake order is heap-priority deterministic. + self._lock: asyncio.Lock = asyncio.Lock() + # Min-heap of _ClockSleeper(deadline_ns, insertion_id, ...). The insertion + # id is the secondary key so simultaneous-deadline parks wake in + # registration order -- also deterministic. + self._waiters: list[_ClockSleeper] = [] + self._insertion_counter: int = 0 + # Optional sync callback invoked right after a new waiter's entry is + # pushed onto the heap. Lets a driver react when a fresh waiter parks + # (e.g., to fast-forward sim time if the pump is currently idle and the + # kick would otherwise be dropped). + self._on_waiter_parked: Callable[[], None] | None = None + + def set_on_waiter_parked(self, cb: Callable[[], None] | None) -> None: + """Register a callback fired right after a new waiter is parked. + + Called synchronously from ``sleep_until_ns`` while the clock's + ``_lock`` is held. The callback MUST NOT await, MUST NOT acquire + ``self._lock``, and should complete in microseconds -- its purpose is + to flip a flag or set an ``asyncio.Event`` the driver pump observes. + Pass ``None`` to clear. + """ + self._on_waiter_parked = cb + + def now_ns(self) -> int: + return self._now_ns + + def _reap_dead_head(self) -> None: + """Pop cancelled sleepers' entries off the heap root. + + Safe without the lock: every heap mutation happens in a synchronous + (no-await) section on the single event loop, so this sync reap cannot + interleave with a mutation in flight. + """ + while self._waiters and not self._waiters[0].alive: + heapq.heappop(self._waiters) + + def has_waiters(self) -> bool: + """Return True if at least one live sleeper is parked.""" + self._reap_dead_head() + return bool(self._waiters) + + def peek_min_waiter_ns(self) -> int | None: + """Return the earliest live parked waiter's deadline, or None if empty. + + Lock-free read of the heap root (dead-head entries from cancelled + sleepers are reaped first so a phantom deadline never fast-forwards + sim time). ``advance_to`` reaps crossed entries under the lock; in the + rare race where this peek sees a stale root, the worst case is a no-op + fast-forward -- the next ``advance_to`` will set the now-crossed + entry's event either way. + """ + self._reap_dead_head() + if not self._waiters: + return None + head_deadline = self._waiters[0].deadline_ns + return head_deadline if head_deadline > self._now_ns else None + + async def advance_to(self, ns: int) -> None: + """Advance sim time to ``ns`` (no-op if ``ns <= now``). + + Pops every waiter with ``deadline <= ns`` in (deadline, insertion_id) + order and sets each waiter's individual ``Event``. asyncio schedules + the awaiting coroutines via ``call_soon`` in the same order they were + set, giving deterministic wake-up ordering for sleepers that crossed + their deadline on the same advance. + """ + async with self._lock: + if ns <= self._now_ns: + return + self._now_ns = ns + while self._waiters and self._waiters[0].deadline_ns <= ns: + waiter = heapq.heappop(self._waiters) + if waiter.alive: + waiter.event.set() + + async def sleep_ns(self, duration_ns: int) -> None: + if duration_ns <= 0: + # Yield once even for non-positive durations (asyncio.sleep + # semantics) so zero-delay replay loops cannot starve the loop. + await asyncio.sleep(0) + return + await self.sleep_until_ns(self._now_ns + duration_ns) + + async def sleep_until_ns(self, deadline_ns: int) -> None: + # Fast path: already crossed. Still yield once (asyncio.sleep + # semantics) so crossed-deadline loops cannot starve the event loop. + if self._now_ns >= deadline_ns: + await asyncio.sleep(0) + return + waiter: _ClockSleeper | None = None + async with self._lock: + # Re-check under the lock; advance_to could have run between the + # fast path and acquiring the lock. + if self._now_ns < deadline_ns: + self._insertion_counter += 1 + waiter = _ClockSleeper( + deadline_ns, self._insertion_counter, asyncio.Event() + ) + heapq.heappush(self._waiters, waiter) + # Synchronous notification -- driver pumps observe a fresh + # waiter without an extra round-trip. Must not await or + # acquire self._lock. + if self._on_waiter_parked is not None: + self._on_waiter_parked() + if waiter is None: + await asyncio.sleep(0) + return + try: + # Wait outside the lock. ``advance_to`` will set this event in + # (deadline, insertion_id) priority order. + await waiter.event.wait() + finally: + # Mark dead on ANY exit (wake or cancellation). A cancelled + # sleeper's entry stays in the heap until lazily reaped, but a + # dead entry never wakes and never counts as a pending deadline. + waiter.alive = False diff --git a/src/aiperf/common/constants.py b/src/aiperf/common/constants.py index 7419466592..926343c5b0 100644 --- a/src/aiperf/common/constants.py +++ b/src/aiperf/common/constants.py @@ -15,6 +15,7 @@ IS_WINDOWS_ARM: bool = IS_WINDOWS and _platform.machine() == "ARM64" NANOS_PER_SECOND = 1_000_000_000 +MICROS_PER_SECOND = 1_000_000 NANOS_PER_MILLIS = 1_000_000 MILLIS_PER_SECOND = 1000 BYTES_PER_MIB = 1024 * 1024 @@ -39,3 +40,12 @@ GOOD_REQUEST_COUNT_TAG = "good_request_count" """GoodRequestCount metric tag""" + +STREAMED_REQUEST_TAG = "streamed_request" +"""Hidden per-record streaming predicate tag: its presence in a record's +``MetricRecordDict`` gates every per-record streaming metric. Kept out of +metric modules so non-metric consumers (e.g. gate checks) need not import them.""" + +STREAMED_REQUEST_COUNT_TAG = "streamed_request_count" +"""Visible aggregate streamed-request-count tag (the streaming denominator +displayed beside Request Count).""" diff --git a/src/aiperf/common/enums/__init__.py b/src/aiperf/common/enums/__init__.py index ef0251a3f4..106acc03ad 100644 --- a/src/aiperf/common/enums/__init__.py +++ b/src/aiperf/common/enums/__init__.py @@ -15,6 +15,7 @@ from aiperf.common.enums.enums import ( AIPerfLogLevel, AudioFormat, + CacheBustTarget, CommAddress, CommandResponseStatus, CommandType, @@ -35,6 +36,7 @@ IPVersion, LifecycleState, MediaType, + MemoryMapFormat, MessageType, ModelSelectionStrategy, OptimizationDirection, @@ -95,6 +97,7 @@ "BaselineKind", "BasePydanticBackedStrEnum", "BasePydanticEnumInfo", + "CacheBustTarget", "CaseInsensitiveStrEnum", "CommAddress", "CommandResponseStatus", @@ -151,6 +154,7 @@ "SSEEventType", "SSEFieldType", "ServerMetricsDiscoveryMode", + "MemoryMapFormat", "ServerMetricsFormat", "ServiceCapability", "ServiceRegistrationStatus", diff --git a/src/aiperf/common/enums/enums.py b/src/aiperf/common/enums/enums.py index fe90a64667..dc90716f56 100644 --- a/src/aiperf/common/enums/enums.py +++ b/src/aiperf/common/enums/enums.py @@ -44,6 +44,25 @@ class AudioFormat(CaseInsensitiveStrEnum): """MP3 format. Compressed audio, smaller file sizes, good quality.""" +class CacheBustTarget(CaseInsensitiveStrEnum): + """Where (and how) to inject a per-session cache-bust marker. + + A unique marker injected into a session's first-turn prompt PREFIX makes the + inference server's KV cache see a distinct prefix per session, defeating + cross-session prefix-cache hits that would otherwise inflate cache-hit + metrics. The graph snapshot-replay path + needs only the ``FIRST_TURN_PREFIX`` and ``NONE`` variants. + """ + + NONE = "none" + """Cache-bust disabled: sessions send their recorded bytes verbatim.""" + + FIRST_TURN_PREFIX = "first_turn_prefix" + """Prepend the marker to the first turn's prompt prefix (token 0), so the + whole prompt's KV-cache prefix diverges per session. Later turns inherit the + changed prefix naturally; only the first turn is independently busted.""" + + class CommAddress(CaseInsensitiveStrEnum): """Enum for specifying the address type for communication clients. This is used to lookup the address in the communication config.""" @@ -140,9 +159,9 @@ class PrerequisiteKind(CaseInsensitiveStrEnum): """Types of conditions that can gate a turn's dispatch. Extensible: v1 orchestrator only honors SPAWN_JOIN; the remaining values - are reserved and rejected at load time by - ``validate_for_orchestrator_v1``. Each deferred value is pinned to a - future orchestrator capability in the DAG prereq-gating design doc. + are reserved for future orchestrator capabilities and rejected at load + time by ``validate_for_orchestrator_v1``, so datasets authored against + them fail loud instead of silently ignoring the gate. """ SPAWN_JOIN = "spawn_join" @@ -481,6 +500,18 @@ class ServerMetricsFormat(CaseInsensitiveStrEnum): Includes cumulative deltas from reference point for counters and histograms.""" +class MemoryMapFormat(CaseInsensitiveStrEnum): + """Format options for memory-mapped backing stores. + + Identifies the on-disk layout of the concatenated-blob data file plus its + offset-index sidecar so readers can select the matching index model. + """ + + DATASET = "dataset" + """Conversation dataset format: one blob per conversation keyed by session ID. + See MemoryMapDatasetBackingStore / MemoryMapDatasetClient.""" + + class ServiceRegistrationStatus(CaseInsensitiveStrEnum): """Defines the various states a service can be in during registration with the SystemController.""" diff --git a/src/aiperf/common/environment.py b/src/aiperf/common/environment.py index 31f32f9e43..dd86f309ea 100644 --- a/src/aiperf/common/environment.py +++ b/src/aiperf/common/environment.py @@ -7,27 +7,33 @@ All settings can be configured via environment variables with the AIPERF_ prefix. Structure: - Environment.ACCURACY.* - Accuracy benchmark settings - Environment.API_SERVER.* - API server settings - Environment.COMPRESSION.* - Compression settings for streaming file transfers - Environment.DATASET.* - Dataset management - Environment.DEV.* - Development and debugging settings - Environment.GPU.* - GPU telemetry collection - Environment.HTTP.* - HTTP client socket and connection settings - Environment.LOGGING.* - Logging configuration - Environment.METRICS.* - Metrics collection and storage - Environment.MLFLOW.* - MLflow export settings - Environment.OTEL.* - OTel metrics streaming - Environment.RECORD.* - Record processing - Environment.SEARCH_PLANNER.* - Adaptive-search planner tunables - Environment.SERVER_METRICS.* - Server metrics collection - Environment.SERVICE.* - Service lifecycle and communication - Environment.TIMING.* - Timing manager settings - Environment.TOKENIZER.* - Tokenizer pre-warm and loading - Environment.UI.* - User interface settings - Environment.WANDB.* - Weights & Biases export settings - Environment.WORKER.* - Worker management and scaling - Environment.ZMQ.* - ZMQ communication settings + Environment.ACCURACY.* - Accuracy benchmark settings + Environment.AGENTX.* - InferenceX AgentX scenario runtime settings + Environment.API_SERVER.* - API server settings + Environment.CLI_RUNNER.* - CLI runner post-run callback behavior + Environment.COMPRESSION.* - Compression settings for streaming file transfers + Environment.DAG.* - DAG benchmark mode (dag_jsonl) settings + Environment.DATASET.* - Dataset management + Environment.DEV.* - Development and debugging settings + Environment.DYNAMO.* - Dynamo agent-trace adapter tunables + Environment.GPU.* - GPU telemetry collection + Environment.GRAPH.* - Graph-mode runtime tunables + Environment.HTTP.* - HTTP client socket and connection settings + Environment.LOGGING.* - Logging configuration + Environment.METRICS.* - Metrics collection and storage + Environment.MLFLOW.* - MLflow export settings + Environment.NETWORK_LATENCY.* - Network latency calibration + Environment.OTEL.* - OTel metrics streaming + Environment.RECORD.* - Record processing + Environment.SEARCH_PLANNER.* - Adaptive-search planner tunables + Environment.SERVER_METRICS.* - Server metrics collection + Environment.SERVICE.* - Service lifecycle and communication + Environment.TIMING.* - Timing manager settings + Environment.TOKENIZER.* - Tokenizer pre-warm and loading + Environment.UI.* - User interface settings + Environment.WANDB.* - Weights & Biases export settings + Environment.WORKER.* - Worker management and scaling + Environment.ZMQ.* - ZMQ communication settings Examples: # Via environment variables: @@ -99,6 +105,45 @@ class _AccuracySettings(BaseSettings): ) +class _AgentXSettings(BaseSettings): + """InferenceX AgentX scenario runtime settings. + + Runtime knobs for the locked AgentX scenario submission contract. The + context-overflow substring allowlist and rate threshold flip a scenario + run to ``submission_valid=false`` when the runtime overflow rate exceeds + the limit (RFC §7). Mirrors ``Environment.AGENTX`` on ``ajc/agentx``. + """ + + model_config = SettingsConfigDict(env_prefix="AIPERF_AGENTX_") + + CONTEXT_OVERFLOW_SUBSTRINGS: list[str] = Field( + default=[ + "context length", + "maximum context", + "context_length_exceeded", + "prompt is too long", + ], + description="Case-insensitive substring allowlist used to classify a " + "server error response as a context-overflow event. Matched against " + "the raw response body and the OpenAI-style nested 'error.message' " + "field. Extend via AIPERF_AGENTX_CONTEXT_OVERFLOW_SUBSTRINGS to " + "support additional inference-server vocabularies (vLLM, TGI, " + "TensorRT-LLM, ...). Empty list disables runtime detection.", + ) + CONTEXT_OVERFLOW_RATE_LIMIT: float = Field( + ge=0.0, + le=1.0, + default=0.01, + description="Strict upper bound on the per-run context-overflow rate " + "(context_overflow_count / total_responses) before a scenario " + "submission is flipped to submission_valid=false with reason " + "'context_overflow_rate_exceeded'. Default 0.01 (1%) matches the " + "scenario spec RFC §7. Comparison is strictly greater-than: rate " + "exactly equal to the limit is accepted. Has no effect on " + "non-scenario runs (no --scenario flag) or runs with zero responses.", + ) + + class _APIServerSettings(BaseSettings): """API server settings. @@ -266,6 +311,129 @@ class _DatasetSettings(BaseSettings): "value, the config loader logs a warning suggesting the user move the " "dataset to a JSONL file. No hard cap.", ) + WEKA_GRAPH_PARALLEL_THRESHOLD: int = Field( + ge=0, + default=8, + description="Minimum number of Weka graph rows/files required before the " + "graph adapter switches from serial parsing to the multi-process " + "forkserver parse path. Set to 0 to force the pool path for any non-empty " + "corpus.", + ) + WEKA_GRAPH_PARALLEL_WORKERS: int = Field( + ge=0, + le=256, + default=0, + description="Number of worker processes for Weka graph parsing. 0 = auto " + "(min(cpu_count - 1, WEKA_GRAPH_PARALLEL_AUTO_MAX_WORKERS, item_count " + "when known)), matching AgentX's full-corpus reconstruction default. Set " + "to 1 to force a single worker.", + ) + WEKA_GRAPH_PARALLEL_AUTO_MAX_WORKERS: int = Field( + ge=1, + le=256, + default=16, + description="Upper bound used when WEKA_GRAPH_PARALLEL_WORKERS=0 selects " + "auto worker sizing for Weka graph parsing.", + ) + WEKA_GRAPH_PARALLEL_PREFETCH_MULTIPLIER: int = Field( + ge=1, + le=64, + default=16, + description="Ordered pool submit-window multiplier for Weka graph parsing. " + "The adapter keeps at most workers * multiplier rows in flight so heavy " + "traces do not idle fast workers while still bounding parent-side memory. " + "Window-sizing rule: the window (workers * multiplier) must cover the rows " + "remaining behind the single heaviest trace, or fast workers stall " + "head-of-line while that one trace drains; a multiplier of 16 gives a " + "window of 256 at the auto 16 workers, which covers this 393-row corpus " + "whose heaviest row sits at index 140. Memory tradeoff: the parent buffers " + "up to workers * multiplier completed payloads plus the raw rows queued " + "behind them; measured on the 062126 full corpus at 16 workers this costs " + "~7.3 GiB parent VmHWM (~17.5 GiB process tree), vs ~2.8 GiB parent at " + "multiplier 4.", + ) + WEKA_GRAPH_PARALLEL_ITEM_TIMEOUT_SECONDS: float = Field( + gt=0.0, + default=600.0, + description="Per-item timeout in seconds for one Weka graph parse pool " + "result (a single trace file/row). A raw multiprocessing Pool never " + "completes the in-flight task of a worker killed mid-parse (OOM kill / " + "external SIGKILL), so an unbounded wait presents as a silent indefinite " + "hang; this bound turns it into a clear RuntimeError instead. Raise it " + "for corpora whose single heaviest trace legitimately parses slower than " + "the default.", + ) + DYNAMO_GRAPH_PARALLEL_THRESHOLD: int = Field( + ge=0, + default=8, + description="Minimum number of Dynamo session-trees in ONE " + "`from_dynamo_trace` call required before the graph adapter switches " + "from the serial tree-scoped build to the multi-process forkserver " + "build path. The parallel unit is TREES (a root session plus its " + "subagent descendants), so even one large file with many trees " + "parallelizes. At or below this count the build stays serial (no pool " + "spawn). Set to 0 to force the pool path for any multi-tree capture.", + ) + DYNAMO_GRAPH_PARALLEL_WORKERS: int = Field( + ge=0, + le=256, + default=0, + description="Number of worker processes for Dynamo session-tree graph " + "building. 0 = auto (min(cpu_count - 1, " + "DYNAMO_GRAPH_PARALLEL_AUTO_MAX_WORKERS, tree_count)). Set to 1 to force " + "the serial path (single-tree captures always build serially).", + ) + DYNAMO_GRAPH_PARALLEL_AUTO_MAX_WORKERS: int = Field( + ge=1, + le=256, + default=16, + description="Upper bound used when DYNAMO_GRAPH_PARALLEL_WORKERS=0 " + "selects auto worker sizing for Dynamo session-tree graph building.", + ) + DYNAMO_GRAPH_PARALLEL_PREFETCH_MULTIPLIER: int = Field( + ge=1, + le=64, + default=16, + description="Ordered pool submit-window multiplier for Dynamo " + "session-tree graph building, doubling as the batches-per-worker " + "factor: the trees are shuffled into at most workers * multiplier " + "CONTIGUOUS weight-balanced per-batch temp files (balanced by a " + "cumulative record-line byte-length proxy resolved WITHOUT parsing the " + "recorded `input_sequence_hashes`) and at most that many batches are " + "kept in flight. Contiguous batching preserves the global sorted-by-root " + "tree order so the parent's in-order union is byte-identical to the " + "serial build; a larger multiplier gives finer batches so a few heavy " + "trees do not idle fast workers, at the cost of more parent-buffered " + "batch results.", + ) + DYNAMO_GRAPH_PARALLEL_ITEM_TIMEOUT_SECONDS: float = Field( + gt=0.0, + default=600.0, + description="Per-item timeout in seconds for one Dynamo session-tree " + "batch build pool result. A raw multiprocessing Pool never completes " + "the in-flight task of a worker killed mid-build (OOM kill / external " + "SIGKILL), so an unbounded wait presents as a silent indefinite hang; " + "this bound turns it into a clear RuntimeError instead. Raise it for " + "captures whose single heaviest batch legitimately builds slower than " + "the default.", + ) + WEKA_HF_SPLIT: str = Field( + default="train", + description="HuggingFace `datasets` split loaded when a weka graph " + "workload is given as a HuggingFace `org/name` repo id (e.g. " + "`semianalysisai/cc-traces-weka-062126`) instead of a local path. The " + "published weka corpora ship a single `traces.jsonl` exposed as the " + "`train` split, so the default is `train`. Datasets slice syntax is " + "accepted (e.g. `train[:100]`) to bound a streamed slice without " + "forcing a full-split download.", + ) + WEKA_HF_REVISION: str | None = Field( + default=None, + description="Optional HuggingFace dataset revision (commit SHA, branch, " + "or tag) pinned when loading a weka corpus by repo id. Pinning a commit " + "SHA makes a streamed ingest reproducible across re-runs even if the " + "dataset's `main` branch is updated. None loads the default revision.", + ) class _DagSettings(BaseSettings): @@ -330,6 +498,27 @@ class _DeveloperSettings(BaseSettings): ) +class _DynamoSettings(BaseSettings): + """Dynamo agent-trace adapter tunables. + + Bounds the `parent_link` chain depth (cycle guard) walked while flattening + the Dynamo trace into the shared segment-trie IR + (`_guard_chain_forest` in + `src/aiperf/dataset/graph/adapters/dynamo/trace.py`). + """ + + model_config = SettingsConfigDict(env_prefix="AIPERF_DYNAMO_", extra="ignore") + + MAX_SUBAGENT_DEPTH: int = Field( + default=16, + ge=1, + description=( + "Maximum `parent_link` chain depth walked by the Dynamo trie-IR " + "cycle guard. Chains deeper than this raise DynamoTraceAdapterError." + ), + ) + + class _GPUSettings(BaseSettings): """GPU telemetry collection configuration. @@ -396,6 +585,156 @@ class _GPUSettings(BaseSettings): ) +class _GraphSettings(BaseSettings): + """Graph-mode runtime tunables. + + Controls async-dataflow graph runtime guardrails and adapter behavior. + Every live producer lowers to flat LlmNode graphs. + """ + + model_config = SettingsConfigDict(env_prefix="AIPERF_GRAPH_", extra="ignore") + + DYNAMIC_POOL_MAX_BYTES: int = Field( + default=64 * 1024 * 1024, + ge=1, + description=( + "Per-worker byte cap for the graph dynamic-content pool (captured " + "assistant responses keyed by trace/variant/ordinal, spliced into " + "successor prompts under per-trace sticky routing). LRU-evicts " + "whole trace entries when exceeded; an evicted live trace surfaces " + "as a loud pool-missing error at its consumer, never a silent " + "truncation. Load-bearing backstop for lost GraphTraceEnd " + "notifications." + ), + ) + MERGE_CONSECUTIVE_USER: bool = Field( + default=False, + description=( + "When a dynamic-content producer fails or returns no replayable " + "content, its assistant turn is OMITTED from the successor's " + "spliced prompt " + "(dynamic-content-pool spec 6.1), which can leave consecutive " + "user-role messages (e.g. an init user turn followed by a static " + "user turn). Some chat APIs reject consecutive same-role messages. " + "When True, the worker merges consecutive user messages in a " + "spliced prompt into one (contents joined by a newline). Default " + "False sends the messages as-is." + ), + ) + IGNORE_EDGE_DELAYS: bool = Field( + default=False, + description=( + "When true, the executor skips honoring " + "StaticEdge.delay_after_predecessor_us, " + "StaticEdge.delay_after_predecessor_start_us and " + "StaticEdge.min_start_delay_us — captured inter-node idle gaps " + "are collapsed and successor nodes fire as soon as their " + "channel-readiness inputs are satisfied. Use for " + "compressed-time replays where the user wants to stress the " + "server independent of the original trace's wall-clock pacing. " + "Default (false): honor edge delays exactly as captured." + ), + ) + WARMUP_MAX_OUTPUT_TOKENS: int = Field( + default=1, + ge=1, + description=( + "Per-request output-token cap applied UNCONDITIONALLY to WARMUP " + "boundary turns (via a `max_tokens` dispatch override in " + "`rewrite_for_warmup` / the warmup delta envelope), so the warmup " + "wire payload always " + "carries `max_tokens == 1`. Warmup is cache-priming history: it " + "primes each live chain's input prefix, and the (1-token) generated " + "output is discarded (warmup records are excluded from metrics). " + "Does NOT affect the PROFILING phase, which replays each turn's full " + "recorded output length." + ), + ) + HANDOFF_RESIDUAL_CAP: float | None = Field( + default=60.0, + gt=0, + description=( + "Upper bound in seconds on any single extended-warmup handoff " + "residual delay (the recorded inter-turn gap minus drain time a " + "resumed profiling frontier waits before firing; see " + "`chop_trie_at_frontier`). Graph edge delays are NOT bounded by " + "the build-plane " + "idle-gap warp when another stream's active interval covers the " + "gap, so this explicit clamp keeps a handoff lane from parking " + "for minutes at the profiling start; the 60s default matches the " + "recorded idle-gap cap. None disables the clamp " + "(faithful-but-unbounded residuals)." + ), + ) + PRESSURE_DRAIN_GRACE_CAP: float = Field( + default=300.0, + gt=0, + description=( + "Upper bound in seconds on the extended (cache-pressure) warmup " + "phase's drain grace period. When a pressure duration is set and " + "no explicit --warmup-grace-period overrides it, the warmup phase " + "waits min(duration, this cap) for in-flight returns after " + "sending completes -- instead of the boundary-priming warmup's unbounded " + "wait, so a wedged or lost return cannot hang the run. On grace " + "expiry the in-flight requests are cancelled; cancelled turns are " + "excluded from the handoff ledger (not-executed, profiling " + "refires them), so a drain whose cancellations land yields a " + "VALID handoff, and a drain that force-completes with unreturned " + "credits trips the completeness gate and skips the re-cut " + "(profiling then starts from the plain t* plans)." + ), + ) + DISPATCH_TIMEOUT: float = Field( + default=300.0, + gt=0, + description=( + "Deadlock guard for the CreditDispatchAdapter Future bridge. The " + "executor's `dispatch(...)` issues a graph credit and awaits the " + "correlated CreditReturn via an asyncio.Future; this caps that wait. " + "If a return is dropped (worker crash, lost message, unrouted " + "return) the Future is rejected with TimeoutError instead of " + "hanging the trace's TaskGroup forever. Sized to comfortably exceed " + "a slow single dispatch round-trip (large prefill + long output); " + "raise it for very long generations, lower it to fail faster in CI." + ), + ) + EXECUTOR_WATCHDOG_TIMEOUT: float | None = Field( + default=None, + gt=0, + description=( + "Optional wall-clock deadlock guard on TraceExecutor.run's firing " + "loop. DISPATCH_TIMEOUT only bounds a node that ISSUED a credit; a " + "node wedged on an unsatisfiable channel input (a producer-accounting " + "bug in mark_producer_done / fan-in gating) never dispatched and so " + "has no per-dispatch guard. When set, each executor run is wrapped in " + "asyncio.wait_for(timeout=this) so such a wedge raises TimeoutError " + "instead of hanging. Default None (disabled): bare/count/session runs " + "replay recorded idle gaps faithfully with no wall-clock bound; " + "a --benchmark-duration already bounds the " + "phase. Set a generous value in CI/dev to convert an executor deadlock " + "into a diagnosable failure without truncating a legitimate long gap." + ), + ) + IDLE_GAP_NO_DURATION_WARN_SECONDS: float = Field( + default=30.0, + ge=0.0, + le=86400.0, + description=( + "Threshold (seconds) for the once-per-run PROFILING-phase advisory " + "GraphIRReplayStrategy logs when an idle-gap recorded corpus is run " + "WITHOUT `--benchmark-duration`. A faithful recorded-trace replay honors every " + "recorded inter-turn idle gap verbatim (each per-gap-warped to the " + "`--synthesis-idle-gap-cap`), so a count/session/bare run spans " + "the slowest admitted trace's full recorded wall time -- sending " + "completes only after every scheduled idle gap elapses. When any inter-turn " + "edge delay exceeds this threshold and no duration budget is set, the " + "strategy emits a NOTICE advising `--benchmark-duration` (the bound) or " + "Ctrl+C (graceful partial export). Set 0.0 to always advise; raise to " + "silence the advisory for corpora with shorter gaps." + ), + ) + + class _HTTPSettings(BaseSettings): """HTTP client socket and connection configuration. @@ -1458,6 +1797,10 @@ class _Environment(BaseSettings): default_factory=_AccuracySettings, description="Accuracy benchmark settings (dataset version pins, etc.)", ) + AGENTX: _AgentXSettings = Field( + default_factory=_AgentXSettings, + description="InferenceX AgentX scenario runtime settings (context-overflow allowlist and rate threshold)", + ) API_SERVER: _APIServerSettings = Field( default_factory=_APIServerSettings, description="API server settings", @@ -1486,10 +1829,18 @@ class _Environment(BaseSettings): default_factory=_DeveloperSettings, description="Development and debugging settings", ) + DYNAMO: _DynamoSettings = Field( + default_factory=_DynamoSettings, + description="Dynamo agent-trace adapter tunables", + ) GPU: _GPUSettings = Field( default_factory=_GPUSettings, description="GPU telemetry collection settings", ) + GRAPH: _GraphSettings = Field( + default_factory=_GraphSettings, + description="Graph-mode runtime tunables (async-dataflow guardrails and adapter behavior)", + ) HTTP: _HTTPSettings = Field( default_factory=_HTTPSettings, description="HTTP client socket and connection settings", diff --git a/src/aiperf/common/hash_id_random_generator.py b/src/aiperf/common/hash_id_random_generator.py new file mode 100644 index 0000000000..5e8cdab14b --- /dev/null +++ b/src/aiperf/common/hash_id_random_generator.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Hash-ID-based random generator for parallel processing with reproducibility. + +Enables parallel processing of traces with hash_ids while maintaining +reproducibility. Each (trace_id, hash_id) pair produces a deterministic random +sequence regardless of worker count or processing order. + +Architecture: + Global Seed -> Base RNG -> (trace_id, hash_id) -> Deterministic tokens + +The trace_id (typically a content hash of the trace file) ensures that different +trace files with overlapping hash_id values produce different content, while the +same trace file always produces identical results. +""" + +import hashlib + +from aiperf.common.random_generator import RandomGenerator + +__all__ = ["HashIdRandomGenerator"] + + +class _DisabledNumpyRNG: + """Raises on any attribute access to prevent NumPy RNG usage.""" + + def __getattr__(self, name): + raise RuntimeError( + "HashIdRandomGenerator does not support NumPy RNG operations. " + "Use Python RNG methods (randrange, choice, etc.) instead." + ) + + +class HashIdRandomGenerator(RandomGenerator): + """RandomGenerator that re-seeds deterministically per (trace_id, hash_id). + + Designed for parallel processing where multiple workers need to generate + identical content for the same hash_id within a trace file. + + Thread Safety: + NOT thread-safe. Each worker process must have its own instance. + """ + + @classmethod + def from_base_rng(cls, base_rng: RandomGenerator) -> "HashIdRandomGenerator": + """Create from a base RandomGenerator (typically from rng.derive()). + + Reads ``base_rng.seed`` without consuming RNG state; only a seedless + (non-deterministic) base falls back to drawing one. Seed 0 is a legal + seed and must not trigger the consuming fallback. + """ + base_seed = ( + base_rng.seed if base_rng.seed is not None else base_rng.randrange(0, 2**64) + ) + return cls(base_seed, _internal=True) + + def __init__(self, base_seed: int, *, _internal: bool = False): + super().__init__(base_seed, _internal=_internal) + self._numpy_rng = _DisabledNumpyRNG() + self._trace_id: str = "" + + def set_trace_id(self, trace_id: str) -> None: + """Set trace identifier to scope hash_ids to a specific trace file. + + Args: + trace_id: Content hash or unique identifier for the trace file. + Different trace files must use different trace_ids. + """ + self._trace_id = trace_id + + def reseed_for_hash_id(self, hash_id: int, trace_id: str | None = None) -> None: + """Re-seed RNG deterministically for a specific hash_id. + + After calling, all random operations use the derived seed until + the next reseed_for_hash_id call. + + Args: + hash_id: KV block hash ID from trace data. + trace_id: Explicit trace scope for this reseed. Overrides the + instance-level ``set_trace_id`` state without mutating it, so + a shared generator can serve interleaved trace scopes. ``None`` + falls back to the instance state (empty string when unset -- + the GLOBAL namespace the linear mooncake path and dynamo's + content-global recorded hashes rely on). + """ + scope = self._trace_id if trace_id is None else trace_id + seed_bytes = hashlib.sha256(f"{self.seed}:{scope}:{hash_id}".encode()).digest() + self._python_rng.seed(int.from_bytes(seed_bytes[:8], "big")) diff --git a/src/aiperf/common/mixins/command_handler_mixin.py b/src/aiperf/common/mixins/command_handler_mixin.py index ae96900c61..bdcce57bcd 100644 --- a/src/aiperf/common/mixins/command_handler_mixin.py +++ b/src/aiperf/common/mixins/command_handler_mixin.py @@ -74,12 +74,16 @@ async def _process_command_message(self, message: CommandMessage) -> None: """ self.trace_or_debug( lambda: f"Received command message: {message}", - lambda: f"Received '{message.command}' command from '{message.service_id}' (id: {message.command_id})", + lambda: ( + f"Received '{message.command}' command from '{message.service_id}' (id: {message.command_id})" + ), ) if message.command_id in self._processed_command_ids: self.trace_or_debug( lambda: f"Received duplicate command message: {message}. Ignoring.", - lambda: f"Received duplicate command message '{message.command}' from '{message.service_id}' (id: {message.command_id}). Ignoring.", + lambda: ( + f"Received duplicate command message '{message.command}' from '{message.service_id}' (id: {message.command_id}). Ignoring." + ), ) # If we receive a duplicate command message, we just send an acknowledged response. await self._publish_command_acknowledged_response(message) @@ -91,7 +95,9 @@ async def _process_command_message(self, message: CommandMessage) -> None: # In the case of a broadcast command, you will receive a command message from yourself. # We ignore these messages. self.debug( - lambda: f"Received broadcast command message from self: {message}. Ignoring." + lambda: ( + f"Received broadcast command message from self: {message}. Ignoring." + ) ) return @@ -317,7 +323,9 @@ async def _process_command_response_message(self, message: CommandResponse) -> N """ self.trace_or_debug( lambda: f"Received command response message: {message}", - lambda: f"Received {message.status} response for command '{message.command}' from '{message.service_id}' (id: {message.command_id})", + lambda: ( + f"Received {message.status} response for command '{message.command}' from '{message.service_id}' (id: {message.command_id})" + ), ) # If the command ID is in the single response futures, we set the result of the future. @@ -328,7 +336,9 @@ async def _process_command_response_message(self, message: CommandResponse) -> N future.set_result(message) else: self.debug( - lambda: f"Already received response for command {message.command_id}, ignoring duplicate from {message.service_id}" + lambda: ( + f"Already received response for command {message.command_id}, ignoring duplicate from {message.service_id}" + ) ) return @@ -343,7 +353,9 @@ async def _process_command_response_message(self, message: CommandResponse) -> N future.set_result(message) else: self.debug( - lambda: f"Already received response for command {message.command_id} from {message.service_id}, ignoring duplicate" + lambda: ( + f"Already received response for command {message.command_id} from {message.service_id}, ignoring duplicate" + ) ) else: if message.command != CommandType.PROFILE_CANCEL: @@ -355,5 +367,7 @@ async def _process_command_response_message(self, message: CommandResponse) -> N # If we reach here, we received a command response that we were not tracking. It is # safe to ignore. self.debug( - lambda: f"Received command response for untracked command: {message}. Ignoring." + lambda: ( + f"Received command response for untracked command: {message}. Ignoring." + ) ) diff --git a/src/aiperf/common/mixins/hooks_mixin.py b/src/aiperf/common/mixins/hooks_mixin.py index 795951af5a..baed7433ff 100644 --- a/src/aiperf/common/mixins/hooks_mixin.py +++ b/src/aiperf/common/mixins/hooks_mixin.py @@ -69,7 +69,9 @@ def __init__(self, **kwargs) -> None: ) self.trace( - lambda: f"Provided hook types: {self._provided_hook_types} for {self.__class__.__name__}" + lambda: ( + f"Provided hook types: {self._provided_hook_types} for {self.__class__.__name__}" + ) ) def _check_hook_type_is_provided(self, hook_type: HookType) -> None: @@ -152,8 +154,9 @@ def for_each_hook_param( ) for param in params: self.trace( - lambda param=param, - type=param_type: f"param: {param}, param_type: {type}" + lambda param=param, type=param_type: ( + f"param: {param}, param_type: {type}" + ) ) if not isinstance(param, param_type): raise ValueError( diff --git a/src/aiperf/common/mixins/message_bus_mixin.py b/src/aiperf/common/mixins/message_bus_mixin.py index 9f7dc2b645..57aa913265 100644 --- a/src/aiperf/common/mixins/message_bus_mixin.py +++ b/src/aiperf/common/mixins/message_bus_mixin.py @@ -50,7 +50,9 @@ def _add_to_subscription_map(hook: Hook, message_type: MessageTypeT) -> None: subscribe_all for efficiency """ self.debug( - lambda: f"Adding subscription for message type: '{message_type}' for hook: {hook}" + lambda: ( + f"Adding subscription for message type: '{message_type}' for hook: {hook}" + ) ) subscription_map.setdefault(message_type, []).append(hook.func) diff --git a/src/aiperf/common/mixins/progress_tracker_mixin.py b/src/aiperf/common/mixins/progress_tracker_mixin.py index edbc5aa623..4e0e308676 100644 --- a/src/aiperf/common/mixins/progress_tracker_mixin.py +++ b/src/aiperf/common/mixins/progress_tracker_mixin.py @@ -80,7 +80,9 @@ def _update_phase_progress( pct = getattr(stats, f"{prefix}_progress_percent") _logger.debug( - lambda: f"Updating {prefix} stats for phase '{stats.phase.title()}': progress_percent: {pct}, finished: {finished}" + lambda: ( + f"Updating {prefix} stats for phase '{stats.phase.title()}': progress_percent: {pct}, finished: {finished}" + ) ) if not pct or finished == 0: diff --git a/src/aiperf/common/mixins/pull_client_mixin.py b/src/aiperf/common/mixins/pull_client_mixin.py index 0918def9dc..2a7f4324bd 100644 --- a/src/aiperf/common/mixins/pull_client_mixin.py +++ b/src/aiperf/common/mixins/pull_client_mixin.py @@ -46,7 +46,9 @@ async def _setup_pull_handler_hooks(self) -> None: def _register_pull_callback(hook: Hook, message_type: MessageTypeT) -> None: self.debug( - lambda: f"Registering pull callback for message type: {message_type} for hook: {hook}" + lambda: ( + f"Registering pull callback for message type: {message_type} for hook: {hook}" + ) ) self.pull_client.register_pull_callback( message_type=message_type, diff --git a/src/aiperf/common/mixins/reply_client_mixin.py b/src/aiperf/common/mixins/reply_client_mixin.py index 16e5dce229..f678cdf896 100644 --- a/src/aiperf/common/mixins/reply_client_mixin.py +++ b/src/aiperf/common/mixins/reply_client_mixin.py @@ -36,11 +36,13 @@ def __init__( @on_init async def _setup_request_handler_hooks(self) -> None: - """Configure the reply client to handle requests for all @request_handler hook decorators.""" + """Configure the reply client to handle requests for all @on_request hook decorators.""" def _register_request_handler(hook: Hook, message_type: MessageTypeT) -> None: self.debug( - lambda: f"Registering request handler for message type: {message_type} for hook: {hook}" + lambda: ( + f"Registering request handler for message type: {message_type} for hook: {hook}" + ) ) self.reply_client.register_request_handler( service_id=self.id, diff --git a/src/aiperf/common/models/dataset_models.py b/src/aiperf/common/models/dataset_models.py index d71f47fb66..0f13e45753 100644 --- a/src/aiperf/common/models/dataset_models.py +++ b/src/aiperf/common/models/dataset_models.py @@ -77,6 +77,54 @@ class MemoryMapClientMetadata(DatasetClientMetadata): ) +class GraphSegmentClientMetadata(DatasetClientMetadata): + """Client metadata for graph unified segment-store access. + + Graph runs have no conversation store: workers materialize per-node + payloads from the unified segment store addressed by + ``(trace_id, node_ordinal)``, and the TimingManager plans from the + graph_meta structural sidecar. Both locations travel here so no service + re-derives them from env conventions. + """ + + client_type: DatasetClientStoreType = DatasetClientStoreType.GRAPH_SEGMENT + + store_base_path: Path = Field( + ..., + description="Base directory containing the per-benchmark " + "aiperf_graph_segments_/ unified store the workers open.", + ) + benchmark_id: str = Field( + ..., + description="Benchmark id keying the per-run store and sidecar directories.", + ) + sidecar_path: Path = Field( + ..., + description="Exact path of the graph_meta.msgpack structural sidecar " + "the TimingManager ingests for scheduling. Mandatory for graph runs.", + ) + + +class GraphDatasetMetadata(AIPerfBaseModel): + """Graph facet of the dataset structure for graph IR workloads. + + Replaces the former per-trace stub conversations: the trace universe and + the per-node theoretical prefix-cache map are first-class here instead of + riding fabricated ConversationMetadata. + """ + + trace_ids: list[str] = Field( + default_factory=list, + description="Every trace id in the built graph store, in catalog order.", + ) + prefix_cache_by_trace: dict[str, dict[str, list[int]]] = Field( + default_factory=dict, + description="Per-trace {node_id: [theoretical_hit_blocks, total_blocks]} " + "map stamped by the trie build; consumed by the records-plane " + "theoretical prefix-cache accumulator.", + ) + + class Media(AIPerfBaseModel): """Base class for all media fields. Contains name and contents of the media data.""" @@ -143,6 +191,24 @@ class TurnMetadata(AIPerfBaseModel): "``ConversationMetadata`` can reach prereqs without holding the full " "Turn list.", ) + theoretical_prefix_cache_hit_blocks: int | None = Field( + default=None, + ge=0, + description="Number of leading hash-id blocks that would be prefix-cache " + "hits for this turn under an infinite per-trace cache. None when the " + "dataset loader did not provide hash-block metadata. Mirrors " + "``Turn.theoretical_prefix_cache_hit_blocks`` so the " + "``theoretical_prefix_cache`` accumulator can read it off " + "``DatasetMetadata`` without holding the full Turn list.", + ) + theoretical_prefix_cache_total_blocks: int | None = Field( + default=None, + ge=0, + description="Number of hash-id blocks considered for theoretical " + "prefix-cache hit accounting for this turn. Pairs with " + "``theoretical_prefix_cache_hit_blocks``; None when the loader did not " + "provide hash-block metadata.", + ) class Turn(AIPerfBaseModel): @@ -209,6 +275,15 @@ class Turn(AIPerfBaseModel): "Mutually exclusive with normal turn-content fields in spirit, but no " "validator enforces that — loaders construct one or the other.", ) + raw_payload_bytes: bytes | None = Field( + default=None, + description="Pre-serialized API request body bytes for verbatim replay. " + "When set, takes precedence over raw_payload: the bytes are sent to the " + "transport unchanged (no orjson.dumps re-encode) and recorded verbatim as " + "payload_bytes for ISL/raw-export. Set by the graph-IR worker bytes path " + "(GraphSegmentUnifiedClient.build_request_body_handles) when no content " + "mutation is needed; None on every other path.", + ) extra_body: dict[str, Any] | None = Field( default=None, description="Non-native per-turn request-body fields (temperature, " @@ -239,6 +314,22 @@ class Turn(AIPerfBaseModel): description="Duration of the audio content in seconds. Used by ASR-specific " "metrics like RTFx. Set by ASR dataset loaders.", ) + theoretical_prefix_cache_hit_blocks: int | None = Field( + default=None, + ge=0, + description="Number of leading hash-id blocks that would hit an infinite " + "per-trace prefix cache for this turn. Set by trace loaders that walk " + "hash_ids during reconstruction (the WEKA graph-IR prepass); None for " + "all other dataset types. Consumed by the ``theoretical_prefix_cache`` " + "results accumulator via ``TurnMetadata``.", + ) + theoretical_prefix_cache_total_blocks: int | None = Field( + default=None, + ge=0, + description="Number of hash-id blocks considered for theoretical " + "prefix-cache hit accounting for this turn. Pairs with " + "``theoretical_prefix_cache_hit_blocks``; None for non-trace datasets.", + ) def metadata(self) -> TurnMetadata: """Get the metadata of the turn.""" @@ -247,6 +338,12 @@ def metadata(self) -> TurnMetadata: delay_ms=self.delay, branch_ids=list(self.branch_ids), prerequisites=list(self.prerequisites), + theoretical_prefix_cache_hit_blocks=( + self.theoretical_prefix_cache_hit_blocks + ), + theoretical_prefix_cache_total_blocks=( + self.theoretical_prefix_cache_total_blocks + ), ) def copy_with_stripped_media(self) -> "Turn": @@ -294,10 +391,17 @@ def copy_with_stripped_media(self) -> "Turn": for vid in self.videos ], raw_payload=self.raw_payload, + raw_payload_bytes=self.raw_payload_bytes, extra_body=dict(self.extra_body) if self.extra_body is not None else None, prerequisites=list(self.prerequisites), branch_ids=list(self.branch_ids), audio_duration_seconds=self.audio_duration_seconds, + theoretical_prefix_cache_hit_blocks=( + self.theoretical_prefix_cache_hit_blocks + ), + theoretical_prefix_cache_total_blocks=( + self.theoretical_prefix_cache_total_blocks + ), ) @@ -367,6 +471,11 @@ class DatasetMetadata(AIPerfBaseModel): "Set by the loader based on dataset format semantics. " "Individual conversations can override this via their own context_mode field.", ) + graph: GraphDatasetMetadata | None = Field( + default=None, + description="Graph facet for graph IR workloads (trace universe + " + "prefix-cache map); None for conversation datasets.", + ) @field_validator("default_context_mode") @classmethod @@ -483,6 +592,12 @@ def metadata(self) -> ConversationMetadata: for bid in t.branch_ids ), prerequisites=list(t.prerequisites), + theoretical_prefix_cache_hit_blocks=( + t.theoretical_prefix_cache_hit_blocks + ), + theoretical_prefix_cache_total_blocks=( + t.theoretical_prefix_cache_total_blocks + ), ) for t in self.turns ] diff --git a/src/aiperf/common/models/export_models.py b/src/aiperf/common/models/export_models.py index a89ad59691..d80811d706 100644 --- a/src/aiperf/common/models/export_models.py +++ b/src/aiperf/common/models/export_models.py @@ -243,6 +243,31 @@ class RunInfo(AIPerfBaseModel): "the run was constructed without a CLI context." ), ) + scenario_name: str | None = Field( + default=None, + description=( + "The locked `--scenario` name applied by the ScenarioResolver, or " + "None when no scenario was set. From " + "`BenchmarkRun.resolved.scenario_outcome.scenario_name`." + ), + ) + submission_valid: bool | None = Field( + default=None, + description=( + "Whether this run satisfies the locked scenario's invariants: True " + "when all invariants held (after auto-fills), False under " + "`--unsafe-override` with downgraded violations, None when no " + "scenario was set." + ), + ) + submission_invalid_reasons: list[str] | None = Field( + default=None, + description=( + "Short tags explaining why `submission_valid` is False (e.g. " + "`unsafe_override`). None / empty when the submission is valid or no " + "scenario was set." + ), + ) @classmethod def from_run(cls, run: BenchmarkRun | None) -> RunInfo | None: @@ -255,6 +280,10 @@ def from_run(cls, run: BenchmarkRun | None) -> RunInfo | None: if run is None: return None variation = run.variation + outcome = run.resolved.scenario_outcome + scenario_name = getattr(outcome, "scenario_name", None) + submission_valid = getattr(outcome, "submission_valid", None) + invalid_reasons = getattr(outcome, "submission_invalid_reasons", None) or None return cls( benchmark_id=run.benchmark_id, sweep_id=run.sweep_id, @@ -265,6 +294,9 @@ def from_run(cls, run: BenchmarkRun | None) -> RunInfo | None: variation_index=variation.index if variation is not None else None, variation_values=dict(variation.values) if variation is not None else None, cli_command=run.cli_command, + scenario_name=scenario_name, + submission_valid=submission_valid, + submission_invalid_reasons=invalid_reasons, ) diff --git a/src/aiperf/common/models/model_endpoint_info.py b/src/aiperf/common/models/model_endpoint_info.py index e5e98bb62b..faf2910e9e 100644 --- a/src/aiperf/common/models/model_endpoint_info.py +++ b/src/aiperf/common/models/model_endpoint_info.py @@ -14,6 +14,7 @@ from pydantic import Field, field_serializer from aiperf.common.enums import ( + CacheBustTarget, ConnectionReuseStrategy, ModelSelectionStrategy, RequestContentType, @@ -122,6 +123,23 @@ def _redact_headers(self, value: list[tuple[str, str]]) -> list[tuple[str, str]] default=EndpointDefaults.USE_SERVER_TOKEN_COUNT, description="Use server-reported token counts from API usage fields instead of client-side tokenization.", ) + cache_bust: CacheBustTarget = Field( + default=CacheBustTarget.NONE, + description="Where to inject a per-trace-instance cache-bust marker on the " + "graph-IR replay path. NONE (default) sends recorded bytes verbatim; " + "FIRST_TURN_PREFIX prepends a `[rid:<12hex>]` marker to the first user turn, " + "shared across the trace instance's turns but distinct per instance and " + "reset on recycle, so the inference server's KV cache sees a distinct prefix " + "per trace instance, defeating cross-instance prefix-cache hits.", + ) + session_routing: str | None = Field( + default=None, + description="Selected session-routing plugin name (None = off).", + ) + session_routing_opts: dict[str, Any] = Field( + default_factory=dict, + description="Raw opts for the routing plugin's Options model.", + ) connection_reuse_strategy: ConnectionReuseStrategy = Field( default=EndpointDefaults.CONNECTION_REUSE_STRATEGY, description="Transport connection reuse strategy.", @@ -207,6 +225,11 @@ def from_run(cls, run: BenchmarkRun) -> ModelEndpointInfo: api_key=ep.api_key, use_legacy_max_tokens=ep.use_legacy_max_tokens, use_server_token_count=ep.use_server_token_count, + cache_bust=getattr(ep, "cache_bust", CacheBustTarget.NONE), + session_routing=( + str(ep.session_routing) if ep.session_routing is not None else None + ), + session_routing_opts=dict(ep.session_routing_opts), connection_reuse_strategy=ep.connection_reuse, download_video_content=ep.download_video_content, request_content_type=ep.request_content_type, diff --git a/src/aiperf/common/models/record_models.py b/src/aiperf/common/models/record_models.py index bb6f958235..6046d51de6 100644 --- a/src/aiperf/common/models/record_models.py +++ b/src/aiperf/common/models/record_models.py @@ -227,6 +227,16 @@ class MetricRecordMetadata(AIPerfBaseModel): description="The x_correlation_id of the parent session that spawned this record's session via a " "DAG subagent fork. None for root sessions. Use to group sibling branches of the same DAG.", ) + context_overflow_skip: bool = Field( + default=False, + description="True iff the record was classified as a context-overflow event AND the active workload " + "is a graph-IR (weka/dynamo trace) run. Set by ``RecordProcessor`` and consumed by ``RecordsManager``: " + "the record still increments the records-side success counter (so the completion barrier stays in lockstep " + "with the credit-side ``final_requests_completed`` and converges), and its ``context_overflow_count`` is " + "still aggregated for the submission-rate gate, but it is excluded from the error tracker and from the " + "performance-metric accumulators (ISL/OSL/TTFT/ITL/latency). Net effect: an overflow event never pollutes " + "a user-facing performance metric, while the run still terminates cleanly.", + ) class TimesliceResult(AIPerfBaseModel): @@ -761,12 +771,36 @@ class RequestInfo(RecordContext): description="Whether this is the final turn in the conversation. " "Used by per-conversation connection strategy to release the connection lease.", ) + root_correlation_id: str | None = Field( + default=None, + description="The x_correlation_id of the depth-0 root of this request's session TREE. " + "Stable across the whole tree (root + every descendant subagent at any depth); equals " + "x_correlation_id for a root session. Sourced from the originating Credit. Per-request " + "transport context read by the session-routing transform to stamp tree-scoped identity " + "on the outbound request; it stays worker-side and is not carried onto the exported record.", + ) + is_parent_final: bool | None = Field( + default=None, + description="Parent conversation had already returned its final turn at " + "credit-issue time; None for roots or when not determinable. Sourced from " + "the originating Credit.", + ) + is_tree_final: bool = Field( + default=False, + description="Provably the last request this session tree will send " + "(conservative False when indeterminate). Sourced from the originating Credit.", + ) url_index: int | None = Field( default=None, ge=0, description="Index of the URL to use when multiple --url values are configured. " "None means use the default (first) URL. Used for round-robin load balancing.", ) + stream_override: bool | None = Field( + default=None, + description="Per-request wire streaming mode for graph credits; None follows " + "model_endpoint.endpoint.streaming.", + ) class RequestRecord(AIPerfBaseModel): @@ -822,6 +856,21 @@ class RequestRecord(AIPerfBaseModel): default=None, description="The error details if the request failed.", ) + context_overflow: bool = Field( + default=False, + description="True iff this request's error response was classified " + "as a server-side context-overflow event by " + "``aiperf.common.scenario.is_context_overflow_response`` " + "(InferenceX AgentX scenario, RFC §7). Set on the record-parser side " + "at response-parsing time; consumed by the " + "``ContextOverflowCountMetric`` aggregate counter.", + ) + streamed: bool = Field( + default=False, + description="Whether this request was sent (and read) in streaming mode " + "on the wire; stamped by the inference client from the per-request " + "effective mode. False for records that never reached the transport.", + ) credit_drop_latency: int | None = Field( default=None, description="The latency of the credit drop in nanoseconds from when it was first received by a Worker to when the inference request was actually sent. " diff --git a/src/aiperf/common/random_generator.py b/src/aiperf/common/random_generator.py index 2132e7e389..ecd00e559c 100644 --- a/src/aiperf/common/random_generator.py +++ b/src/aiperf/common/random_generator.py @@ -59,6 +59,7 @@ "derive", "init", "reset", + "root_seed", ] _logger = AIPerfLogger(__name__) @@ -509,6 +510,16 @@ def reset() -> None: _manager = None +def root_seed() -> int | None: + """Return the global manager's root seed, or None when uninitialized or unseeded. + + Read-only introspection for callers that need to resolve an effective seed + (e.g. the graph adapters' content-seed ladder) without reaching into the + manager's privates. + """ + return _manager._root_seed if _manager is not None else None + + def derive_variation_seed( root_seed: int | None, variation_label: str, diff --git a/src/aiperf/common/scenario/__init__.py b/src/aiperf/common/scenario/__init__.py new file mode 100644 index 0000000000..f43c502426 --- /dev/null +++ b/src/aiperf/common/scenario/__init__.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from aiperf.common.scenario.base import ( + ScenarioLockError, + ScenarioOutcome, + ScenarioSpec, + ScenarioViolation, + TrajectoryWarmupFailedError, + UnknownScenarioError, +) +from aiperf.common.scenario.context_overflow import is_context_overflow_response +from aiperf.common.scenario.inferencex_agentx_mvp import INFERENCEX_AGENTX_MVP +from aiperf.common.scenario.registry import SCENARIOS, get_scenario +from aiperf.common.scenario.submission_outcome import ( + CONTEXT_OVERFLOW_REASON, + compute_submission_outcome, +) +from aiperf.common.scenario.validator import apply_scenario + +__all__ = [ + "CONTEXT_OVERFLOW_REASON", + "INFERENCEX_AGENTX_MVP", + "SCENARIOS", + "ScenarioLockError", + "ScenarioOutcome", + "ScenarioSpec", + "ScenarioViolation", + "TrajectoryWarmupFailedError", + "UnknownScenarioError", + "apply_scenario", + "compute_submission_outcome", + "get_scenario", + "is_context_overflow_response", +] diff --git a/src/aiperf/common/scenario/_env_locks.py b/src/aiperf/common/scenario/_env_locks.py new file mode 100644 index 0000000000..056493c877 --- /dev/null +++ b/src/aiperf/common/scenario/_env_locks.py @@ -0,0 +1,251 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Config apply-or-lock + submission scenario locks for graph-IR runtime knobs. + +The trajectory-start (t*) window (``cfg.trajectory_start_min/max_ratio``) and +the per-trace idle-gap cap (``synthesis.idle_gap_cap_seconds``) are per-run +config: their locks AUTO-APPLY the scenario value when the user left the field +unset and raise a violation only when a user-explicit value differs (mirroring +AgentX's semantics). Scenario application performs NO process-global writes. + +This module also holds the two scenario locks that do not fit +``validator.py``'s per-``_apply_*`` config-mutation shape: the concurrency-sweep +rejection (reads the run's ``SweepVariation``) and the canonical weka HF-repo +pin (sniffs the resolved input path). All are split out of ``validator.py`` to +keep that module under the file-size budget. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from aiperf.common.scenario.base import ScenarioSpec, ScenarioViolation + +if TYPE_CHECKING: + from aiperf.config.resolution.plan import BenchmarkRun + +_logger = logging.getLogger(__name__) + + +def apply_trajectory_ratios( + run: BenchmarkRun, + spec: ScenarioSpec, + violations: list[ScenarioViolation], + applied: list[str], +) -> None: + """Apply-or-lock for the trajectory-start (t*) window on the run config. + + The window lives at ``cfg.trajectory_start_min_ratio`` / + ``cfg.trajectory_start_max_ratio`` (``--trajectory-start-min-ratio`` / + ``--trajectory-start-max-ratio``), per-run config threaded natively to the + Dataset/Timing services. A user-explicit value is LOCKED (violation on + mismatch); an unset field is auto-applied from the spec. A violated bound + never blocks its unset sibling from being auto-applied (under + ``--unsafe-override`` the run proceeds with the mixed window). + """ + checks = ( + ( + "trajectory_start_min_ratio", + "--trajectory-start-min-ratio", + spec.default_trajectory_start_min_ratio, + ), + ( + "trajectory_start_max_ratio", + "--trajectory-start-max-ratio", + spec.default_trajectory_start_max_ratio, + ), + ) + cfg = run.cfg + any_checked = False + added = False + for field, flag, required in checks: + if required is None: + continue + any_checked = True + if field in cfg.model_fields_set: + actual = getattr(cfg, field) + if actual != required: + added = True + violations.append( + ScenarioViolation( + flag=flag, + current_value=actual, + required_value=required, + message=( + f"scenario {spec.name!r} locks the trajectory-start " + f"window; {flag} must equal {required}" + ), + ) + ) + continue + setattr(cfg, field, required) + _logger.info(f"Scenario {spec.name!r}: auto-set {flag}={required} (was unset).") + if any_checked and not added: + applied.append("trajectory_start_ratios") + + +def apply_trace_idle_gap_cap( + run: BenchmarkRun, + spec: ScenarioSpec, + violations: list[ScenarioViolation], + applied: list[str], +) -> None: + """Apply-or-lock for the per-trace idle-gap cap. + + The cap lives at ``synthesis.idle_gap_cap_seconds`` + (``--synthesis-idle-gap-cap``). A user-explicit value — including an + explicit null (warp disabled) — is LOCKED: a violation is raised when it + differs from ``spec.trace_idle_gap_cap_seconds``. When the field is unset + the scenario auto-applies the spec value onto the run config, so the bare + 60s adapter default never leaks into a scenario run (agentx-on-main + parity: apply when unset, lock when explicit). + """ + if spec.trace_idle_gap_cap_seconds is None: + return + dataset = run.cfg.get_default_dataset() + if dataset is None or "synthesis" not in type(dataset).model_fields: + # Synthetic datasets have no trace idle-gap concept; nothing to lock. + return + synthesis = getattr(dataset, "synthesis", None) + explicit = ( + synthesis is not None and "idle_gap_cap_seconds" in synthesis.model_fields_set + ) + if explicit: + actual = synthesis.idle_gap_cap_seconds + if actual != spec.trace_idle_gap_cap_seconds: + violations.append( + ScenarioViolation( + flag="--synthesis-idle-gap-cap", + current_value=actual, + required_value=spec.trace_idle_gap_cap_seconds, + message=( + f"scenario {spec.name!r} locks the per-trace idle-gap cap " + f"to {spec.trace_idle_gap_cap_seconds}; " + "--synthesis-idle-gap-cap must match" + ), + ) + ) + else: + applied.append("trace_idle_gap_cap") + return + if synthesis is None: + from aiperf.config.dataset import SynthesisConfig + + dataset.synthesis = SynthesisConfig( + idle_gap_cap_seconds=spec.trace_idle_gap_cap_seconds + ) + else: + synthesis.idle_gap_cap_seconds = spec.trace_idle_gap_cap_seconds + _logger.info( + f"Scenario {spec.name!r}: auto-set " + f"--synthesis-idle-gap-cap={spec.trace_idle_gap_cap_seconds} (was unset)." + ) + applied.append("trace_idle_gap_cap") + + +# Final dotted-path segments whose presence in a run's SweepVariation marks a +# concurrency sweep (a scenario locks ONE config; a sweep would multiply it). +_CONCURRENCY_SWEEP_SUFFIXES: tuple[str, ...] = ( + "concurrency", + "prefill_concurrency", + "warmup_concurrency", + "warmup_prefill_concurrency", +) + +# Canonical SemiAnalysis weka HF dataset org/repo prefix. The published weka +# corpora are all ``semianalysisai/cc-traces-weka-[-256k]`` repos. +# Workload detection cannot distinguish a specific dated repo (it sniffs +# only "is this a weka graph workload?"), so the lock pins the +# org/repo PREFIX: an HF-id input must live under this prefix to be a recognized +# canonical weka workload. A non-canonical HF id (a foreign org, or a repo that +# merely carries the "weka" marker) is flagged submission_invalid. +_WEKA_HF_REPO_PREFIX = "semianalysisai/cc-traces-weka" +# The corpus repo the scenario contract requires; surfaced as the +# example/required value so the violation message points at the canonical corpus. +_WEKA_HF_CANONICAL_REPO = "semianalysisai/cc-traces-weka-062126" + + +def apply_concurrency_sweep( + run: BenchmarkRun, + spec: ScenarioSpec, + violations: list[ScenarioViolation], + applied: list[str], +) -> None: + """Reject a swept ``--concurrency`` (or related) under a fixed-spec scenario. + + A scenario locks ONE fixed configuration; a swept concurrency + (``--concurrency 10,20,30``) would multiply it into N runs with diverging + settings, so it must be rejected outright. A raw list never reaches this + check, though: the sweep is expanded + one level UP -- ``--concurrency 10,20,30`` is promoted into a grid sweep and + each ``BenchmarkRun`` carries a single concrete concurrency plus a + ``SweepVariation`` whose ``values`` record the swept dotted-path key + (e.g. ``phases.profiling.concurrency``). Detecting that variation key here + flags every run minted from a concurrency sweep + as a violation (a ``ScenarioLockError``, downgradable under + ``--unsafe-override``). A single non-swept run has ``variation=None`` (or a + variation with no concurrency key) and is untouched. + """ + variation = getattr(run, "variation", None) + values = getattr(variation, "values", None) + if not isinstance(values, dict): + return + swept = [ + key for key in values if key.rsplit(".", 1)[-1] in _CONCURRENCY_SWEEP_SUFFIXES + ] + if not swept: + return + violations.append( + ScenarioViolation( + flag="--concurrency", + current_value=sorted(swept), + required_value="single value (no sweep)", + message=( + f"scenario {spec.name!r} does not support parameter sweeps; " + "pass a single --concurrency value instead of a list " + "(a sweep multiplies the locked config into N diverging runs)" + ), + ) + ) + + +def pin_weka_hf_repo( + run: BenchmarkRun, + spec: ScenarioSpec, + violations: list[ScenarioViolation], +) -> bool: + """Pin a HuggingFace-id weka input to the canonical org/repo prefix. + + Returns True (and appends a violation) when the input is an HF id NOT under + :data:`_WEKA_HF_REPO_PREFIX`; returns False (no violation) for a local-file + weka workload or a canonical HF id. Local imports avoid pulling the heavy + dataset graph loader at ``aiperf.common.scenario`` import time. + """ + from aiperf.dataset.graph.adapters.weka.trace import ( + _hf_dataset_id_str, + _looks_like_hf_dataset_id, + ) + from aiperf.dataset.graph.workload_detect import resolve_graph_workload + + ref = resolve_graph_workload(run) + if ref is None: + return False + repo = _hf_dataset_id_str(ref.path) + if not _looks_like_hf_dataset_id(repo): + return False + if repo.lower().startswith(_WEKA_HF_REPO_PREFIX): + return False + violations.append( + ScenarioViolation( + flag="--input-file (hf-repo)", + current_value=repo, + required_value=f"{_WEKA_HF_REPO_PREFIX}-* (e.g. {_WEKA_HF_CANONICAL_REPO})", + message=( + f"scenario {spec.name!r} only allows the canonical SemiAnalysis " + f"weka HF corpora ({_WEKA_HF_REPO_PREFIX}-*); the resolved " + f"dataset {repo!r} is not a recognized weka workload" + ), + ) + ) + return True diff --git a/src/aiperf/common/scenario/base.py b/src/aiperf/common/scenario/base.py new file mode 100644 index 0000000000..5d2d8d1225 --- /dev/null +++ b/src/aiperf/common/scenario/base.py @@ -0,0 +1,194 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from typing import Any + +from pydantic import ConfigDict, Field + +from aiperf.common.enums import CacheBustTarget +from aiperf.common.models import AIPerfBaseModel +from aiperf.plugin.enums import TimingMode + + +class ScenarioSpec(AIPerfBaseModel): + """Frozen declaration of a benchmark scenario's invariants.""" + + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid", frozen=True) + + name: str = Field(description="Scenario identifier, e.g. 'inferencex-agentx-mvp'.") + timing_mode: TimingMode = Field( + description="Required timing mode for this scenario." + ) + require_ignore_eos: bool = Field( + description="Inject ignore_eos=true into extra_inputs; error on explicit false." + ) + require_streaming: bool = Field( + default=False, + description=( + "Force --streaming=true (auto-enabled when unset; error on explicit " + "--no-streaming). Streaming is required for the per-token latency " + "metrics (TTFT, ITL) that are core to this benchmark; without it a " + "run would silently report no first-token signal." + ), + ) + forbid_input_truncation: bool = Field( + description=( + "Reject client-side input-length truncation. Currently checks " + "`--synthesis-max-isl` (which drops traces whose input length " + "exceeds the cap)." + ) + ) + require_loader: str | tuple[str, ...] = Field( + description=( + "Required loader plugin name (e.g. 'weka_trace'), or a tuple of " + "equivalent loader names. The detected loader must match any one " + "of them — useful when several loader plugins produce byte-identical " + "data (e.g. file-based vs HF-hosted variants)." + ) + ) + min_benchmark_duration_seconds: int = Field( + gt=0, description="Floor on --benchmark-duration in seconds." + ) + default_benchmark_duration_seconds: int | None = Field( + default=None, + gt=0, + description=( + "Value auto-filled into --benchmark-duration when the user leaves " + "it unset. Explicit user values are honored (subject to the " + "min_benchmark_duration_seconds floor). None disables auto-fill." + ), + ) + default_trajectory_start_min_ratio: float | None = Field( + default=None, + ge=0.0, + le=1.0, + description=( + "Trajectory-start (t*) window lower bound the scenario runs with, " + "living at cfg.trajectory_start_min_ratio " + "(--trajectory-start-min-ratio). apply_trajectory_ratios " + "AUTO-APPLIES this value onto the run config when the field is " + "unset; a user-explicit value differing from this raises " + "ScenarioLockError naming --trajectory-start-min-ratio " + "(downgradable to a warning via --unsafe-override). None disables " + "the check." + ), + ) + default_trajectory_start_max_ratio: float | None = Field( + default=None, + ge=0.0, + le=1.0, + description=( + "Trajectory-start (t*) window upper bound the scenario runs with, " + "living at cfg.trajectory_start_max_ratio " + "(--trajectory-start-max-ratio). apply_trajectory_ratios " + "AUTO-APPLIES this value onto the run config when the field is " + "unset; a user-explicit value differing from this raises " + "ScenarioLockError naming --trajectory-start-max-ratio " + "(downgradable to a warning via --unsafe-override). None disables " + "the check." + ), + ) + trace_idle_gap_cap_seconds: float | None = Field( + default=None, + ge=0, + description=( + "Hard ceiling (seconds) for idle gaps within each root trace. For " + "recorded graph replay (weka, dynamo), parent + subagent " + "request-start timestamps are compressed " + "per-trace before per-turn delays are derived." + ), + ) + require_cache_bust: CacheBustTarget | None = Field( + default=None, + description=( + "When set, endpoint.cache_bust must equal this value. " + "Mismatch is rejected unless --unsafe-override is also set " + "(which stamps submission_valid=false)." + ), + ) + + +class ScenarioViolation(AIPerfBaseModel): + """A single conflict between user config and a locked scenario invariant.""" + + flag: str = Field( + description="The user-facing flag or config field that conflicts." + ) + current_value: Any = Field(description="The value the user provided.") + required_value: Any = Field(description="The value the scenario requires.") + message: str = Field(description="Human-readable explanation of the conflict.") + + def __str__(self) -> str: + return ( + f"{self.flag}: got {self.current_value!r}, " + f"required {self.required_value!r} ({self.message})" + ) + + +class ScenarioOutcome(AIPerfBaseModel): + """Result of applying a scenario lock against a resolved ``BenchmarkRun``. + + Produced by ``aiperf.common.scenario.apply_scenario`` and stored on + ``run.resolved.scenario_outcome``. ``submission_valid`` is ``None`` when no + scenario is set, ``True`` when all invariants are satisfied (after + auto-fills), and ``False`` under ``--unsafe-override`` when violations were + downgraded to warnings. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + scenario_name: str | None = Field( + default=None, + description="The applied scenario name, or None when no --scenario was set.", + ) + applied_locks: list[str] = Field( + default_factory=list, + description="Short tags for each invariant lock that was applied " + "(auto-filled or validated), e.g. 'timing_mode', 'streaming', " + "'cache_bust'. Order reflects application order.", + ) + violations: list[ScenarioViolation] = Field( + default_factory=list, + description="All scenario invariant conflicts collected in one pass. " + "Non-empty only under --unsafe-override (otherwise a ScenarioLockError " + "is raised).", + ) + submission_valid: bool | None = Field( + default=None, + description="True when the scenario lock is satisfied, False under " + "--unsafe-override with violations, None when no scenario is set.", + ) + submission_invalid_reasons: list[str] = Field( + default_factory=list, + description="Short tags explaining why submission_valid is False " + "(e.g. 'unsafe_override').", + ) + + +class ScenarioLockError(ValueError): + """Raised when a scenario lock is violated and --unsafe-override is not set.""" + + def __init__(self, violations: list[ScenarioViolation]) -> None: + self.violations = violations + joined = "\n - ".join(str(v) for v in violations) + super().__init__( + f"Scenario invariants violated ({len(violations)} conflict" + f"{'s' if len(violations) != 1 else ''}):\n - {joined}\n" + "Pass --unsafe-override to convert to warnings (run will be marked submission_valid=false)." + ) + + +class TrajectoryWarmupFailedError(RuntimeError): + """Raised when WARMUP has terminal failures across trajectories and PROFILING cannot honestly start.""" + + def __init__(self, failed_trace_ids: list[str]) -> None: + self.failed_trace_ids = failed_trace_ids + super().__init__( + f"Trajectory warmup failed for {len(failed_trace_ids)} trace(s): " + f"{', '.join(failed_trace_ids)}. Run aborted to preserve metrics integrity." + ) + + +class UnknownScenarioError(ValueError): + """Raised when --scenario references a name not in the registry.""" diff --git a/src/aiperf/common/scenario/context_overflow.py b/src/aiperf/common/scenario/context_overflow.py new file mode 100644 index 0000000000..951f02c11c --- /dev/null +++ b/src/aiperf/common/scenario/context_overflow.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Runtime context-overflow detection for the InferenceX AgentX scenario. + +Substring allowlist lives on ``Environment.AGENTX`` so users can extend +without code changes. +""" + +from typing import Any + +import orjson + +from aiperf.common.environment import Environment + + +def is_context_overflow_response( + *, + body: str | bytes | None, + substrings: list[str] | None = None, +) -> bool: + """Classify whether an error response indicates a context-overflow. + + Performs a case-insensitive substring match against: + 1. The raw response body text. + 2. The OpenAI-style nested ``error.message`` field, when the body + parses as JSON. Falls through silently on non-JSON bodies (e.g. + vLLM's ``{"detail": "..."}`` shape -- which is still caught by the + raw-body match in step 1). + + Callers are expected to pre-filter to error responses (the + ``InferenceResultParser`` only invokes this on records with + ``has_error=True``); status-code gating lives at the call site so this + function stays a pure body-based classifier. + + Args: + body: The raw response body. ``str`` or ``bytes``; ``None`` returns + False. Empty body returns False. + substrings: Override the allowlist for tests. ``None`` reads + ``Environment.AGENTX.CONTEXT_OVERFLOW_SUBSTRINGS`` at call time + so test settings overrides take effect. + + Returns: + True iff at least one substring (case-insensitive) appears in + either the raw body text or the parsed ``error.message`` field. + """ + if body is None: + return False + + # Resolve substring list lazily so tests can override the env setting. + candidates: list[str] = ( + substrings + if substrings is not None + else list(Environment.AGENTX.CONTEXT_OVERFLOW_SUBSTRINGS) + ) + if not candidates: + return False + + text: str = ( + body.decode("utf-8", errors="replace") if isinstance(body, bytes) else body + ) + + if not text: + return False + + lowered = text.lower() + needles = [s.lower() for s in candidates if s] + + # 1. Raw body match. + for needle in needles: + if needle in lowered: + return True + + # 2. OpenAI-style {"error": {"message": "..."}} match. + nested_message = _extract_openai_error_message(text) + if nested_message: + nested_lower = nested_message.lower() + for needle in needles: + if needle in nested_lower: + return True + + return False + + +def _extract_openai_error_message(text: str) -> str | None: + """Return the OpenAI-style ``error.message`` field from a JSON body. + + Returns ``None`` when the body doesn't parse as JSON, or when the + expected ``{"error": {"message": ...}}`` shape isn't present. Tolerates + a string-shaped ``error`` field (some servers return ``"error": + "..."``) by using it as the message. + """ + try: + parsed: Any = orjson.loads(text) + except (orjson.JSONDecodeError, ValueError, TypeError): + return None + if not isinstance(parsed, dict): + return None + err = parsed.get("error") + if isinstance(err, dict): + msg = err.get("message") + if isinstance(msg, str): + return msg + elif isinstance(err, str): + return err + return None diff --git a/src/aiperf/common/scenario/inferencex_agentx_mvp.py b/src/aiperf/common/scenario/inferencex_agentx_mvp.py new file mode 100644 index 0000000000..471b65c0d3 --- /dev/null +++ b/src/aiperf/common/scenario/inferencex_agentx_mvp.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from aiperf.common.enums import CacheBustTarget +from aiperf.common.scenario.base import ScenarioSpec +from aiperf.plugin.enums import TimingMode + +INFERENCEX_AGENTX_MVP = ScenarioSpec( + name="inferencex-agentx-mvp", + timing_mode=TimingMode.GRAPH_IR, + require_ignore_eos=True, + require_streaming=True, + forbid_input_truncation=True, + require_loader=( + "semianalysis_cc_traces_weka_with_subagents", + "semianalysis_cc_traces_weka_with_subagents_256k", + "semianalysis_cc_traces_weka_with_subagents_060226", + "semianalysis_cc_traces_weka_with_subagents_060226_256k", + "semianalysis_cc_traces_weka_with_subagents_060526", + "semianalysis_cc_traces_weka_with_subagents_060526_256k", + "semianalysis_cc_traces_weka_with_subagents_060826", + "semianalysis_cc_traces_weka_with_subagents_060826_256k", + "semianalysis_cc_traces_weka_061326", + "semianalysis_cc_traces_weka_061326_256k", + "semianalysis_cc_traces_weka_061526", + "semianalysis_cc_traces_weka_061526_256k", + "semianalysis_cc_traces_weka_062126", + "semianalysis_cc_traces_weka_062126_256k", + "weka_trace", + "weka_hf", + ), + min_benchmark_duration_seconds=900, + default_benchmark_duration_seconds=1800, + default_trajectory_start_min_ratio=0.0, + default_trajectory_start_max_ratio=1.0, + trace_idle_gap_cap_seconds=10.0, + require_cache_bust=CacheBustTarget.FIRST_TURN_PREFIX, +) diff --git a/src/aiperf/common/scenario/registry.py b/src/aiperf/common/scenario/registry.py new file mode 100644 index 0000000000..86f0b81961 --- /dev/null +++ b/src/aiperf/common/scenario/registry.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from aiperf.common.scenario.base import ScenarioSpec, UnknownScenarioError +from aiperf.common.scenario.inferencex_agentx_mvp import INFERENCEX_AGENTX_MVP + +SCENARIOS: dict[str, ScenarioSpec] = { + INFERENCEX_AGENTX_MVP.name: INFERENCEX_AGENTX_MVP, +} + + +def get_scenario(name: str) -> ScenarioSpec: + """Return the registered :class:`ScenarioSpec` for ``name``. + + Raises :class:`UnknownScenarioError` (listing the valid names) when ``name`` + is not a registered scenario. + """ + if name not in SCENARIOS: + valid = ", ".join(sorted(SCENARIOS.keys())) + raise UnknownScenarioError( + f"Unknown scenario {name!r}. Valid scenarios: {valid}" + ) + return SCENARIOS[name] diff --git a/src/aiperf/common/scenario/submission_outcome.py b/src/aiperf/common/scenario/submission_outcome.py new file mode 100644 index 0000000000..879422547a --- /dev/null +++ b/src/aiperf/common/scenario/submission_outcome.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Combine the static scenario-lock outcome with runtime threshold checks. + +The validator-side outcome (``ScenarioOutcome`` from ``apply_scenario``) covers +static config violations -- invariant-lock conflicts and ``--unsafe-override``. +This helper folds in runtime-only signals that are only knowable post-run: the +context-overflow response rate (InferenceX AgentX spec §7) and early +cancellation. +""" + +from aiperf.common.environment import Environment + +CONTEXT_OVERFLOW_REASON = "context_overflow_rate_exceeded" +RUN_CANCELLED_REASON = "run_cancelled" + + +def compute_submission_outcome( + *, + scenario_name: str | None, + validator_submission_valid: bool | None, + validator_reasons: list[str] | None = None, + total_responses: int = 0, + context_overflow_count: int = 0, + was_cancelled: bool = False, +) -> tuple[bool | None, list[str]]: + """Combine validator outcome with runtime threshold checks into a verdict. + + The validator-side outcome covers static config violations (handled by + ``apply_scenario`` and stored on ``run.resolved.scenario_outcome``). This + helper folds in runtime-only signals that are only knowable post-run -- + the >1% context-overflow rate per the InferenceX AgentX spec §7, and early + cancellation (Ctrl+C): + a cancelled run produces partial metrics and is never a valid submission. + + Rate semantics: strictly greater than + ``Environment.AGENTX.CONTEXT_OVERFLOW_RATE_LIMIT`` (default 0.01 per the + InferenceX AgentX spec §7, override via + ``AIPERF_AGENTX_CONTEXT_OVERFLOW_RATE_LIMIT``) flips + ``submission_valid`` to False; equal-to is accepted (boundary behavior + pinned by tests). When ``total_responses == 0`` the rate is treated as 0 + (undefined / no successful responses), so the overflow rule does not flip + submission validity in that case. + + When ``scenario_name`` is None this is a no-scenario run and the function + returns ``(None, [])`` -- callers should drop the ``submission_valid`` + field from the output entirely. + + Args: + scenario_name: Active scenario, or None for a non-scenario run. + validator_submission_valid: Outcome from ``apply_scenario`` -- True if + the static lock was satisfied, False under ``--unsafe-override`` + with violations, None for a non-scenario run. + validator_reasons: Reason codes already collected by the validator + (e.g. ``"unsafe_override"``). + total_responses: Total responses received during the run (successes + + overflow + other failures). + context_overflow_count: Count of context-overflow responses during the + run. + was_cancelled: Whether the run was cancelled early (graceful Ctrl+C). + True flips ``submission_valid`` to False with reason + ``"run_cancelled"``. + + Returns: + A ``(submission_valid, reasons)`` tuple. ``submission_valid`` is + ``None`` when ``scenario_name`` is None. + """ + if scenario_name is None: + return None, [] + + reasons: list[str] = list(validator_reasons or []) + valid: bool = ( + bool(validator_submission_valid) + if validator_submission_valid is not None + else True + ) + + if total_responses > 0: + rate = context_overflow_count / total_responses + if rate > Environment.AGENTX.CONTEXT_OVERFLOW_RATE_LIMIT: + valid = False + if CONTEXT_OVERFLOW_REASON not in reasons: + reasons.append(CONTEXT_OVERFLOW_REASON) + + if was_cancelled: + valid = False + if RUN_CANCELLED_REASON not in reasons: + reasons.append(RUN_CANCELLED_REASON) + + return valid, reasons diff --git a/src/aiperf/common/scenario/validator.py b/src/aiperf/common/scenario/validator.py new file mode 100644 index 0000000000..f6b451d6d5 --- /dev/null +++ b/src/aiperf/common/scenario/validator.py @@ -0,0 +1,402 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Scenario invariant lock applied as a config-resolver step (graph-IR path). + +``apply_scenario(run)`` reads ``run.cfg.scenario``; when set, it looks up the +:class:`ScenarioSpec`, auto-fills unset defaults (info log) and validates each +invariant against ``run.cfg`` / ``run.resolved``. Explicit user conflicts raise +:class:`ScenarioLockError` unless ``run.cfg.unsafe_override`` downgrades them to +warnings and stamps ``submission_valid=False``. + +The validator derives +``TimingMode.GRAPH_IR`` from weka graph-workload detection (no per-phase +``timing_mode`` override consumer), stores cache-bust as a bare +``CacheBustTarget`` on ``EndpointConfig.cache_bust`` (no ``CacheBustConfig`` +wrapper), and treats the trajectory-start window + idle-gap cap as per-run +config -- their locks (``_env_locks.py``) auto-apply the spec value when the +field is unset and raise only on a user-explicit mismatch. + +User-explicit vs default is read from ``model_fields_set`` membership on the +LIVE converted config (the resolver chain mutates ``run`` in place), so +membership is faithful here -- no explicit-set sentinel is required. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any + +from aiperf.common.enums import CacheBustTarget +from aiperf.common.scenario._env_locks import ( + apply_concurrency_sweep, + apply_trace_idle_gap_cap, + apply_trajectory_ratios, + pin_weka_hf_repo, +) +from aiperf.common.scenario.base import ( + ScenarioLockError, + ScenarioOutcome, + ScenarioSpec, + ScenarioViolation, +) +from aiperf.common.scenario.registry import get_scenario + +if TYPE_CHECKING: + from aiperf.config.resolution.plan import BenchmarkRun + + +def _is_graph_workload(run: BenchmarkRun) -> bool: + """Lazily import + delegate to the memoized graph-workload resolution. + + The import is local to avoid a module-import cycle: the dataset graph + loader pulls heavy deps, and ``aiperf.common.scenario`` is imported early + by the config/resolver chain. + """ + from aiperf.dataset.graph.workload_detect import resolve_graph_workload + + return resolve_graph_workload(run) is not None + + +_logger = logging.getLogger(__name__) + +# Synthetic loader identity returned for any weka graph workload. +_WEKA_GRAPH_LOADER = "weka_trace" + + +def apply_scenario(run: BenchmarkRun) -> ScenarioOutcome: + """Apply the locked scenario invariants to ``run.cfg`` / ``run.resolved``. + + Reads ``run.cfg.scenario``; when None returns a no-op outcome. Otherwise + looks up the spec, auto-fills defaults + validates each invariant, and + raises :class:`ScenarioLockError` on a conflict unless + ``run.cfg.unsafe_override`` downgrades to warnings + ``submission_valid=False``. + The result is stored on ``run.resolved.scenario_outcome`` and returned. + """ + scenario_name = getattr(run.cfg, "scenario", None) + if scenario_name is None: + outcome = ScenarioOutcome() + run.resolved.scenario_outcome = outcome + return outcome + + spec = get_scenario(scenario_name) + violations: list[ScenarioViolation] = [] + applied: list[str] = [] + + _apply_timing_mode(run, spec, violations, applied) + _apply_require_streaming(run, spec, violations, applied) + _apply_ignore_eos(run, spec, violations, applied) + _apply_forbid_input_truncation(run, spec, violations, applied) + _apply_require_loader(run, spec, violations, applied) + _apply_require_cache_bust(run, spec, violations, applied) + apply_concurrency_sweep(run, spec, violations, applied) + _apply_duration(run, spec, violations, applied) + apply_trajectory_ratios(run, spec, violations, applied) + apply_trace_idle_gap_cap(run, spec, violations, applied) + + # There is deliberately NO random_seed auto-fill here (no secrets.randbits + # fallback when unset). The graph-IR config has no input.random_seed field: the run + # seed lives on BenchmarkRun.random_seed, threaded deterministically by the + # orchestrator (resolve_run_seed -> variation_seeds / derive_variation_seed) + # BEFORE this resolver step runs. Synthesized content is already seed- + # invariant via resolve_graph_content_seed (which returns --random-seed + # verbatim, with no weka-specific fallback). A non-deterministic + # secrets.randbits here would collide with the + # orchestrator's deterministic derivation and break the content-determinism + # contract (every parse of the same run must synthesize identical bytes); + # the observable behavior (a seed is always assigned, the run is + # reproducible) is identical-or-better without it. + + unsafe = bool(getattr(run.cfg, "unsafe_override", False)) + if violations and not unsafe: + raise ScenarioLockError(violations) + + if violations and unsafe: + for v in violations: + _logger.warning("Scenario violation (override active): %s", v) + outcome = ScenarioOutcome( + scenario_name=spec.name, + applied_locks=applied, + violations=violations, + submission_valid=False, + submission_invalid_reasons=["unsafe_override"], + ) + run.resolved.scenario_outcome = outcome + return outcome + + outcome = ScenarioOutcome( + scenario_name=spec.name, + applied_locks=applied, + violations=[], + submission_valid=True, + ) + run.resolved.scenario_outcome = outcome + return outcome + + +def _is_falsy_extra_input(value: Any) -> bool: + """True when ``value`` is an explicit falsy ``ignore_eos`` extra-input.""" + if isinstance(value, bool): + return not value + if isinstance(value, str): + return value.strip().lower() in ("false", "0", "no") + if isinstance(value, (int, float)): + return value == 0 + return False + + +def _has_config_field(model: Any, field: str) -> bool: + """True when ``field`` is a real declared Pydantic field on ``model``. + + Used to keep dormant locks honest: a check whose backing config field does + not exist must NOT claim to be "applied". ``getattr`` alone can't tell a + real field from an absent one, so this inspects the model's declared fields. + """ + return field in getattr(type(model), "model_fields", {}) + + +def _apply_timing_mode( + run: BenchmarkRun, + spec: ScenarioSpec, + violations: list[ScenarioViolation], + applied: list[str], +) -> None: + """Verify the run is a graph workload (graph-IR is derived, not stamped). + + The target derives ``GRAPH_IR`` from ``resolve_graph_workload(run)`` with + no per-phase ``timing_mode`` consumer, so this verifies the workload IS + graph and raises otherwise rather than stamping phases. + """ + if _is_graph_workload(run): + applied.append("timing_mode") + return + violations.append( + ScenarioViolation( + flag="--input-file (timing_mode)", + current_value="non-graph workload", + required_value=str(spec.timing_mode), + message=( + f"scenario {spec.name!r} requires a weka graph workload " + f"(timing_mode={spec.timing_mode}); the input is not one" + ), + ) + ) + + +def _apply_require_streaming( + run: BenchmarkRun, + spec: ScenarioSpec, + violations: list[ScenarioViolation], + applied: list[str], +) -> None: + """Auto-enable ``--streaming`` when unset; violation on explicit ``False``.""" + if not spec.require_streaming: + return + endpoint = run.cfg.endpoint + if endpoint.streaming: + applied.append("streaming") + return + explicit = "streaming" in endpoint.model_fields_set + if explicit: + violations.append( + ScenarioViolation( + flag="--streaming", + current_value=False, + required_value=True, + message=( + f"scenario {spec.name!r} requires --streaming; the " + "per-token latency metrics (TTFT, ITL) need streaming" + ), + ) + ) + else: + endpoint.streaming = True + _logger.info("Scenario %r: forcing --streaming=true (was unset).", spec.name) + applied.append("streaming") + + +def _apply_ignore_eos( + run: BenchmarkRun, + spec: ScenarioSpec, + violations: list[ScenarioViolation], + applied: list[str], +) -> None: + """Inject ``ignore_eos=true`` into ``endpoint.extra`` (the wire body). + + ``EndpointConfig.extra`` is merged into every request body by the formatters. + Auto-injects when absent; raises when explicitly falsy. + """ + if not spec.require_ignore_eos: + return + extra = run.cfg.endpoint.extra + ignore_eos = extra.get("ignore_eos") + if ignore_eos is None: + extra["ignore_eos"] = True + _logger.info("Scenario %r: injecting ignore_eos=true (was absent).", spec.name) + applied.append("ignore_eos") + elif _is_falsy_extra_input(ignore_eos): + violations.append( + ScenarioViolation( + flag="extra_inputs.ignore_eos", + current_value=ignore_eos, + required_value=True, + message=f"scenario {spec.name!r} requires ignore_eos=true", + ) + ) + else: + applied.append("ignore_eos") + + +def _apply_forbid_input_truncation( + run: BenchmarkRun, + spec: ScenarioSpec, + violations: list[ScenarioViolation], + applied: list[str], +) -> None: + """Reject ``--synthesis-max-isl`` (the ``FileDataset.synthesis.max_isl`` ISL + filter). A weka graph dataset without a ``synthesis`` block has no such + field, so this degrades to a graceful no-op. + """ + if not spec.forbid_input_truncation: + return + dataset = run.cfg.get_default_dataset() + synthesis = getattr(dataset, "synthesis", None) + max_isl = getattr(synthesis, "max_isl", None) + if max_isl is not None: + violations.append( + ScenarioViolation( + flag="--synthesis-max-isl", + current_value=max_isl, + required_value=None, + message=( + f"scenario {spec.name!r} forbids client-side input " + "truncation; --synthesis-max-isl drops over-length traces, " + "falsifying the workload" + ), + ) + ) + else: + applied.append("forbid_input_truncation") + + +def _apply_require_loader( + run: BenchmarkRun, + spec: ScenarioSpec, + violations: list[ScenarioViolation], + applied: list[str], +) -> None: + """Require the active loader to be a canonical weka graph workload. + + The DatasetResolver does NOT populate ``dataset_types`` for weka inputs, so + detection is the graph-workload sniff: a weka graph workload reports + ``"weka_trace"``; else None. When the input is a HuggingFace dataset id + (rather than a local file) the resolved repo is additionally pinned to the + canonical SemiAnalysis weka org/repo prefix, so a submission cannot swap in + a foreign corpus. A foreign HF id that merely carries the "weka" + marker is flagged submission_invalid. + """ + if spec.require_loader is None: + return + allowed = ( + (spec.require_loader,) + if isinstance(spec.require_loader, str) + else tuple(spec.require_loader) + ) + detected = _WEKA_GRAPH_LOADER if _is_graph_workload(run) else None + if detected not in allowed: + display = allowed[0] if len(allowed) == 1 else f"any of {sorted(allowed)}" + violations.append( + ScenarioViolation( + flag="--input-file (loader)", + current_value=detected, + required_value=display, + message=f"scenario {spec.name!r} requires loader={display}", + ) + ) + return + if not pin_weka_hf_repo(run, spec, violations): + applied.append("require_loader") + + +def _apply_require_cache_bust( + run: BenchmarkRun, + spec: ScenarioSpec, + violations: list[ScenarioViolation], + applied: list[str], +) -> None: + """Auto-fill the bare ``endpoint.cache_bust`` to the required value when at + default ``NONE`` and not user-set; violation on an explicit different value. + The auto-filled value drives the already-wired worker stamping. + """ + if spec.require_cache_bust is None: + return + endpoint = run.cfg.endpoint + actual = endpoint.cache_bust + if actual == spec.require_cache_bust: + applied.append("cache_bust") + return + explicit = "cache_bust" in endpoint.model_fields_set + if explicit: + violations.append( + ScenarioViolation( + flag="--cache-bust", + current_value=str(actual), + required_value=str(spec.require_cache_bust), + message=( + f"scenario {spec.name!r} requires " + f"cache_bust={spec.require_cache_bust}; got {actual}" + ), + ) + ) + elif actual == CacheBustTarget.NONE: + endpoint.cache_bust = spec.require_cache_bust + _logger.info( + "Scenario %r: auto-set --cache-bust=%s (was default).", + spec.name, + spec.require_cache_bust, + ) + applied.append("cache_bust") + + +def _apply_duration( + run: BenchmarkRun, + spec: ScenarioSpec, + violations: list[ScenarioViolation], + applied: list[str], +) -> None: + """Auto-fill / enforce the profiling-phase duration floor. + + Graph-IR completion is trace/stop-condition driven, so this is a + submission-validity + run-length lock: auto-fill unset durations from the + default; violation when an explicit duration is below the floor. + """ + profiling_phases = run.cfg.get_profiling_phases() + if spec.default_benchmark_duration_seconds is not None: + for phase in profiling_phases: + if phase.duration is None: + phase.duration = float(spec.default_benchmark_duration_seconds) + _logger.info( + "Scenario %r: auto-set duration=%ss (was unset).", + spec.name, + spec.default_benchmark_duration_seconds, + ) + applied.append("default_benchmark_duration") + + below_floor = False + for phase in profiling_phases: + duration = phase.duration or 0.0 + if duration < spec.min_benchmark_duration_seconds: + below_floor = True + violations.append( + ScenarioViolation( + flag="--benchmark-duration", + current_value=duration, + required_value=f">={spec.min_benchmark_duration_seconds}", + message=( + f"scenario {spec.name!r} requires duration >= " + f"{spec.min_benchmark_duration_seconds}s to reach " + "steady state and trigger KV offloading" + ), + ) + ) + if not below_floor: + applied.append("min_benchmark_duration") diff --git a/src/aiperf/config/cli_parameter.py b/src/aiperf/config/cli_parameter.py index d16bd060ff..29802b1b73 100644 --- a/src/aiperf/config/cli_parameter.py +++ b/src/aiperf/config/cli_parameter.py @@ -58,3 +58,4 @@ class Groups: ZMQ_COMMUNICATION = Group.create_ordered("ZMQ Communication") ACCURACY = Group.create_ordered("Accuracy") MULTI_RUN = Group.create_ordered("Multi-Run") + SCENARIO = Group.create_ordered("Scenario") diff --git a/src/aiperf/config/comm/base.py b/src/aiperf/config/comm/base.py index 5cd5ec02d9..5897f6fe2e 100644 --- a/src/aiperf/config/comm/base.py +++ b/src/aiperf/config/comm/base.py @@ -121,25 +121,25 @@ def get_address(self, address_type: CommAddress) -> str: # CommAddress to a resolver callable keeps BaseZMQCommunicationConfig.get_address # flat (no match ladder) and makes the mapping easy to extend. _ADDRESS_RESOLVERS: dict[CommAddress, Callable[[BaseZMQCommunicationConfig], str]] = { - CommAddress.EVENT_BUS_PROXY_FRONTEND: lambda c: c.event_bus_proxy_config.resolve_frontend( - c._remote_host + CommAddress.EVENT_BUS_PROXY_FRONTEND: lambda c: ( + c.event_bus_proxy_config.resolve_frontend(c._remote_host) ), - CommAddress.EVENT_BUS_PROXY_BACKEND: lambda c: c.event_bus_proxy_config.resolve_backend( - c._remote_host + CommAddress.EVENT_BUS_PROXY_BACKEND: lambda c: ( + c.event_bus_proxy_config.resolve_backend(c._remote_host) ), - CommAddress.DATASET_MANAGER_PROXY_FRONTEND: lambda c: c.dataset_manager_proxy_config.resolve_frontend( - c._remote_host + CommAddress.DATASET_MANAGER_PROXY_FRONTEND: lambda c: ( + c.dataset_manager_proxy_config.resolve_frontend(c._remote_host) ), - CommAddress.DATASET_MANAGER_PROXY_BACKEND: lambda c: c.dataset_manager_proxy_config.resolve_backend( - c._remote_host + CommAddress.DATASET_MANAGER_PROXY_BACKEND: lambda c: ( + c.dataset_manager_proxy_config.resolve_backend(c._remote_host) ), # Raw inference proxy is always local (within-pod IPC). Workers and record # processors are co-located in the same pod, so remote_host is ignored. - CommAddress.RAW_INFERENCE_PROXY_FRONTEND: lambda c: c.raw_inference_proxy_config.resolve_frontend( - None + CommAddress.RAW_INFERENCE_PROXY_FRONTEND: lambda c: ( + c.raw_inference_proxy_config.resolve_frontend(None) ), - CommAddress.RAW_INFERENCE_PROXY_BACKEND: lambda c: c.raw_inference_proxy_config.resolve_backend( - None + CommAddress.RAW_INFERENCE_PROXY_BACKEND: lambda c: ( + c.raw_inference_proxy_config.resolve_backend(None) ), CommAddress.CREDIT_ROUTER: lambda c: c.credit_router_address, CommAddress.CREDIT_RETURN_ROUTER: lambda c: c.credit_return_router_address, diff --git a/src/aiperf/config/config.py b/src/aiperf/config/config.py index ddc532f78b..94a6c97d1c 100644 --- a/src/aiperf/config/config.py +++ b/src/aiperf/config/config.py @@ -388,6 +388,86 @@ class BenchmarkConfig(BaseConfig, BenchmarkHelpersMixin): ), ] + scenario: Annotated[ + str | None, + Field( + default=None, + description="Lock all benchmark invariants for a named scenario " + "(e.g. 'inferencex-agentx-mvp'). Plain data here; the lock is " + "applied by the ScenarioResolver step in the pre-bootstrap resolver " + "chain (auto-fills defaults, validates, raises ScenarioLockError on " + "conflict). Distinct from the sweep ``scenarios`` strategy " + "(SweepConfig), which expands hand-picked named runs -- this field " + "is a single invariant LOCK, not a sweep.", + ), + ] + + unsafe_override: Annotated[ + bool, + Field( + default=False, + description="Convert scenario lock errors to warnings; stamps " + "submission_valid=false in the resolved scenario outcome. No-op " + "without ``scenario``. Plain data; consumed by the ScenarioResolver.", + ), + ] + + trajectory_start_min_ratio: Annotated[ + float | None, + Field( + default=None, + ge=0.0, + le=1.0, + description="Lower bound (fraction of each trace's recorded " + "wall-clock duration) of the per-trace t* snapshot-instant sampling " + "window for recorded graph replay. Unset (None) resolves to 0.0. " + "With both bounds at 0.0 the window is OFF (full native replay). " + "The inferencex-agentx-mvp scenario auto-applies 0.0 when unset " + "and locks an explicit conflicting value. Must be <= " + "trajectory_start_max_ratio.", + ), + ] + + trajectory_start_max_ratio: Annotated[ + float | None, + Field( + default=None, + ge=0.0, + le=1.0, + description="Upper bound (fraction of each trace's recorded " + "wall-clock duration) of the per-trace t* snapshot-instant sampling " + "window for recorded graph replay. Unset (None) resolves to 0.0 " + "(window OFF, full native replay); any value > 0 engages the t* " + "snapshot machinery (per-trace t* sampling + auto boundary-priming " + "warmup). The inferencex-agentx-mvp scenario auto-applies 1.0 when " + "unset and locks an explicit conflicting value.", + ), + ] + + burst_phase_starts: Annotated[ + bool, + Field( + default=False, + description="Burst-vs-spread control for phase starts on the " + "recorded-trace graph-IR replay path (--burst-phase-starts). " + "False = SPREAD (recorded leading offsets honored); True = BURST " + "(both phase starts collapse into a synchronized burst).", + ), + ] + + agentic_cache_warmup_duration: Annotated[ + float | None, + Field( + default=None, + gt=0, + description="Extended (cache-pressure) warmup duration in seconds " + "for recorded graph replay (--agentic-cache-warmup-duration). " + "After the boundary-priming warmup drains, the live post-t* " + "replay continues with zero idle delay and one-token outputs for " + "this long before profiling. None = no pressure stage.", + ), + ] + # ========================================================================== # VALIDATORS # ========================================================================== @@ -398,7 +478,7 @@ def normalize_before_validation(cls, data: Any) -> Any: """Normalize input data before Pydantic validation. Handles singular/plural aliases and warmup/profiling-to-phases - shorthand. See `_benchmark_normalizers.normalize_benchmark_input`. + shorthand. See `aiperf.config.loader.normalizers.normalize_benchmark_input`. """ return normalize_benchmark_input(data) @@ -410,7 +490,7 @@ def parse_phases(cls, v: Any) -> list[Any]: The dict shape is rejected with a migration-pointing message; valid shorthand inputs (`warmup:` / `profiling:` top-level, or a single flat config under `phases:`) are converted to lists by the - pre-model normalizers in `_benchmark_normalizers`. + pre-model normalizers in `aiperf.config.loader.normalizers`. """ if isinstance(v, dict): raise ValueError( @@ -431,10 +511,22 @@ def parse_phases(cls, v: Any) -> list[Any]: def parse_datasets(cls, v: Any) -> list[Any]: """Parse dataset configurations into a list shape, validating each item has a name. - See `_benchmark_normalizers.parse_datasets_input`. + See `aiperf.config.loader.normalizers.parse_datasets_input`. """ return parse_datasets_input(v) + @model_validator(mode="after") + def validate_trajectory_window_ordered(self) -> Self: + """Reject an inverted t* window (min ratio above max ratio).""" + lo = self.trajectory_start_min_ratio or 0.0 + hi = self.trajectory_start_max_ratio or 0.0 + if lo > hi: + raise ValueError( + "trajectory_start_min_ratio must be <= trajectory_start_max_ratio; " + f"got [{lo}, {hi}]" + ) + return self + @model_validator(mode="after") def validate_phase_names_unique(self) -> Self: """Reject duplicate phase names — they must be unique within the list.""" diff --git a/src/aiperf/config/dataset/config.py b/src/aiperf/config/dataset/config.py index 09db0bef28..9bb87b08c8 100644 --- a/src/aiperf/config/dataset/config.py +++ b/src/aiperf/config/dataset/config.py @@ -47,7 +47,11 @@ ) from aiperf.config.loader.normalizers import _hoist_synthetic_prompt_fields from aiperf.config.types import SamplingDistribution -from aiperf.plugin.enums import DatasetSamplingStrategy, PublicDatasetType +from aiperf.plugin.enums import ( + DatasetSamplingStrategy, + GraphAdapterType, + PublicDatasetType, +) _logger = AIPerfLogger(__name__) @@ -302,7 +306,10 @@ class FileDataset(BaseConfig): "Can be absolute or relative. Mutually exclusive with `records:`. " "Supported formats depend on the format field: " "JSONL for single_turn/multi_turn, JSONL trace files for mooncake_trace/" - "bailian_trace, Parquet for baseten_trace, directories for random_pool.", + "bailian_trace, Parquet for baseten_trace, directories for random_pool." + " Graph trace inputs (`.jsonl` / `.jsonl.gz` / segmented `.gz` directories) " + "are recognized as graph workloads independently of `format` — auto-detected, " + "or forced with `graph_format`.", ), ] @@ -329,10 +336,28 @@ class FileDataset(BaseConfig): "mooncake_trace / bailian_trace / baseten_trace / burst_gpt_trace: " "timestamped trace files for replay. " "sagemaker_data_capture: JSONL captured by SageMaker DataCapture. " - "random_pool: directory of reusable prompts.", + "random_pool: directory of reusable prompts." + " This selects a *custom dataset loader* (conversation-style inputs) and is " + "ignored for graph workloads, which are auto-detected from `path` or forced " + "via `graph_format`.", ), ] + graph_format: Annotated[ + GraphAdapterType | None, + Field( + default=None, + description="Force a graph-adapter format for this `--input-file`, " + "overriding workload auto-detection. Accepts a registered graph " + "adapter name (e.g. `dynamo_trace`, `weka_trace`, `dag_jsonl`, " + "`native`). When set, the file is treated as a graph workload and " + "parsed by the named adapter; this is the only way to load a " + "hand-authored `native` graph file or run a `dag_jsonl` conversation " + "file on the graph runtime, both otherwise excluded from " + "auto-detection. None = auto-detect.", + ), + ] = None + sampling: Annotated[ DatasetSamplingStrategy, Field( @@ -352,7 +377,9 @@ class FileDataset(BaseConfig): "Allows scaling timestamps and token lengths before replay. " "Applies to trace formats such as mooncake_trace and baseten_trace, " "except speedup_ratio, which is rejected for baseten_trace " - "(use replay_speedup / --replay-speedup to scale replay pacing there).", + "(use replay_speedup / --replay-speedup to scale replay pacing there). " + "Also used by graph (weka) replay, whose `synthesis.max_osl` caps the " + "recorded output length (`--synthesis-max-osl`).", ), ] diff --git a/src/aiperf/config/dataset/resolver.py b/src/aiperf/config/dataset/resolver.py index 1c37287721..3c1d332b45 100644 --- a/src/aiperf/config/dataset/resolver.py +++ b/src/aiperf/config/dataset/resolver.py @@ -83,6 +83,7 @@ def resolve(self, run: BenchmarkRun) -> None: ``FileNotFoundError`` when a file dataset path does not exist. """ from aiperf.config.dataset import FileDataset, PublicDataset + from aiperf.dataset.graph.workload_detect import resolve_graph_workload from aiperf.plugin import plugins acc = _DatasetResolution() @@ -95,9 +96,17 @@ def resolve(self, run: BenchmarkRun) -> None: continue if not isinstance(ds, FileDataset): continue - self._resolve_one(name=ds.name, ds=ds, format_map=format_map, acc=acc) + self._resolve_one( + run=run, name=ds.name, ds=ds, format_map=format_map, acc=acc + ) self._publish(run, acc) + # Eagerly resolve-and-memoize graph-ness so single-run child processes + # inherit it via the pickled run and never re-walk the registry. A + # no-op when a `_resolve_one` site already populated it; for + # synthetic/public/inline runs this stores None WITH the resolved + # marker ("not a graph run", not "never checked"). + resolve_graph_workload(run) @staticmethod def _publish(run: BenchmarkRun, acc: _DatasetResolution) -> None: @@ -123,6 +132,7 @@ def _publish(run: BenchmarkRun, acc: _DatasetResolution) -> None: def _resolve_one( self, *, + run: BenchmarkRun, name: str, ds: object, format_map: dict[str, object], @@ -133,12 +143,53 @@ def _resolve_one( self._resolve_inline(name=name, ds=ds, format_map=format_map, acc=acc) return + # A weka HuggingFace ``org/name`` repo id is a graph workload reference, + # not a local path (resolved from the hub/cache by the weka adapter; the + # graph-store planes own its record counting), so skip file resolution. + # Populate the graph memo here too: this branch never reaches the main + # graph check below, and the accessor stores the RAW ``org/name`` id + # verbatim (never .resolve()d -- HF ids are not filesystem paths). + from pathlib import Path + + from aiperf.dataset.graph.adapters.weka.trace import ( + _looks_like_hf_dataset_id, + ) + from aiperf.dataset.graph.workload_detect import ( + is_graph_workload_path, + resolve_graph_workload, + ) + + is_default_dataset = ds is run.cfg.get_default_dataset() + raw_path = ds.path # type: ignore[attr-defined] + if raw_path is not None and _looks_like_hf_dataset_id(raw_path): + acc.paths[name] = Path(str(raw_path)) + if is_default_dataset: + resolve_graph_workload(run) + return + # 1. Resolve and validate path - resolved = ds.path.resolve() # type: ignore[attr-defined] + resolved = raw_path.resolve() # type: ignore[attr-defined] if not resolved.exists(): raise FileNotFoundError(f"Dataset '{name}' file not found: {resolved}") acc.paths[name] = resolved + # A local graph workload (dynamo/weka trace file, segmented dir) is + # owned by the graph-store build/schedule planes (parse_graph_workload), + # not the custom-dataset-loader path. Skip type detection + record + # counting, which text-mode read the file and crash on a .gz binary. + # Forced via `--graph-format`, OR auto-detected as a trace adapter. + # `graph_format` covers the `native`/ambiguous case that the path-level + # content sniff (which excludes `native`) cannot see. For the default + # dataset the check IS the accessor, so graph-ness is derived (and the + # file sniffed) exactly once and memoized onto ``run.resolved``. + if is_default_dataset: + if resolve_graph_workload(run) is not None: + return + elif getattr(ds, "graph_format", None) is not None or is_graph_workload_path( + resolved + ): + return + # 2. Detect dataset type from explicit format or via can_load. # Pydantic defaults ``format`` to SINGLE_TURN, so a falsy check isn't # enough — when the user didn't *explicitly* set format, fall back to diff --git a/src/aiperf/config/dataset/trace.py b/src/aiperf/config/dataset/trace.py index c96341aa6d..bbe55f4988 100644 --- a/src/aiperf/config/dataset/trace.py +++ b/src/aiperf/config/dataset/trace.py @@ -9,7 +9,7 @@ from __future__ import annotations -from typing import Annotated +from typing import Annotated, Literal from pydantic import ( ConfigDict, @@ -98,3 +98,62 @@ class SynthesisConfig(BaseConfig): "Traces with output_length > max_osl are capped to this value (not filtered).", ), ] + + max_context_length: Annotated[ + int | None, + Field( + ge=1, + default=None, + description="Maximum per-trace context length (tokens) for graph-plane " + "dataset selection (`--max-context-length`). Traces whose input+output " + "context would exceed this cap are excluded from selection. None (the " + "default) applies no context-length filter. Raw explicit value carried " + "verbatim from the CLI; the derived selection default is computed " + "elsewhere. Ignored by non-graph datasets.", + ), + ] + + allow_dataset_wrap: Annotated[ + bool | None, + Field( + default=None, + description="Whether graph-plane dataset selection may wrap (reuse the " + "finite trace pool) to satisfy `--request-count` (`--allow-dataset-wrap` / " + "`--no-allow-dataset-wrap`). None (the default) means unset -- the effective " + "value is derived downstream and surfaced on `run.resolved.allow_dataset_wrap`; " + "an explicit True/False here is the raw user intent carried verbatim so the " + "resolver can distinguish unset from explicit. Ignored by non-graph datasets.", + ), + ] + + idle_gap_cap_seconds: Annotated[ + float | None, + Field( + ge=0.0, + default=60.0, + description="Per-trace idle-gap cap (seconds) for recorded graph " + "(`weka_trace`, `dynamo_trace`) replay (`--synthesis-idle-gap-cap`). " + "Every recorded " + "inter-request idle gap longer than this is compressed to the cap, " + "shifting later requests left by the excess, so a recorded multi-hour " + "idle gap does not park the replay. Default 60.0. Set to null to " + "disable warping and replay the raw recorded timestamps. Ignored by " + "non-graph datasets.", + ), + ] + + corpus: Annotated[ + Literal["coding", "sonnet"] | None, + Field( + default=None, + description="Corpus backing recorded graph (`weka_trace`, `dynamo_trace`) " + "real-content " + "synthesis (`--prompt-corpus`). `coding` (the default when unset) uses " + "the procedural CodingContentGenerator pool -- the same corpus the " + "recorded weka workloads were captured against. `sonnet` uses the " + "Shakespeare PromptGenerator pool, which yields matching token counts " + "but different bytes (useful only to reproduce golden fixtures built " + "from that pool). Ignored by non-graph " + "datasets.", + ), + ] diff --git a/src/aiperf/config/endpoint.py b/src/aiperf/config/endpoint.py index a5a8fee1e0..21cb9f7498 100644 --- a/src/aiperf/config/endpoint.py +++ b/src/aiperf/config/endpoint.py @@ -22,6 +22,7 @@ ) from aiperf.common.enums import ( + CacheBustTarget, ConnectionReuseStrategy, ModelSelectionStrategy, RequestContentType, @@ -30,6 +31,7 @@ from aiperf.config.loader.parsing import normalize_http_urls from aiperf.plugin.enums import ( EndpointType, + SessionRoutingType, TransportType, URLSelectionStrategy, ) @@ -63,6 +65,7 @@ class EndpointDefaults: WAIT_FOR_MODEL_TIMEOUT = 0.0 WAIT_FOR_MODEL_INTERVAL = 5.0 WAIT_FOR_MODEL_MODE = "inference" + CACHE_BUST = CacheBustTarget.NONE class TemplateConfig(BaseConfig): @@ -239,6 +242,20 @@ def _redact_api_key(self, value: str | None) -> str | None: ), ] + cache_bust: Annotated[ + CacheBustTarget, + Field( + default=EndpointDefaults.CACHE_BUST, + description="Where to inject a per-trace-instance cache-bust marker on " + "the graph-IR replay path. 'none' (default): recorded bytes sent " + "verbatim. 'first_turn_prefix': prepend a '[rid:<12hex>]' marker to the " + "first user turn, shared across the trace instance's turns but distinct " + "per instance and reset on recycle, so the inference server's KV cache " + "sees a distinct prefix per trace instance, defeating cross-instance " + "prefix-cache hits.", + ), + ] + template: Annotated[ TemplateConfig | None, Field( @@ -316,6 +333,34 @@ def _redact_headers(self, value: dict[str, str]) -> dict[str, str]: ), ] + session_routing: Annotated[ + SessionRoutingType | None, + Field( + default=None, + description=( + "Session-routing transform stamping per-session identity on " + "every request so an external router (SGLang Model Gateway, Dynamo, " + "generic session-affinity LBs) can pin a session's turns to " + "one replica. One mode per run; parameterize with " + "--session-routing-opt key=value." + ), + ), + ] + + session_routing_opts: Annotated[ + dict[str, Any], + Field( + default_factory=dict, + description=( + "Mode-specific options for --session-routing, validated " + "against the selected plugin's Options model (unknown keys and " + "invalid values are rejected at startup). Values are " + "canonicalized/coerced to the plugin's Options model types at " + "validation, so downstream consumers see typed values." + ), + ), + ] + wait_for_model_timeout: Annotated[ float, Field( @@ -525,3 +570,29 @@ def _validate_request_content_type(self) -> Self: f"video_generation); endpoint type {self.type} does not." ) return self + + @model_validator(mode="after") + def validate_session_routing(self) -> Self: + """Fail fast: opts require a mode; opts must satisfy the plugin's Options. + + Canonicalizes the opts to the plugin's Options model types so downstream + consumers (including the pickled BenchmarkRun that reaches workers) carry + coerced values (e.g. ``{"timeout_seconds": 600}``, not ``"600"``). + """ + if self.session_routing is None: + if self.session_routing_opts: + raise ValueError( + "--session-routing-opt requires --session-routing to select a mode." + ) + return self + from aiperf.plugin import plugins + from aiperf.plugin.enums import PluginType + + routing_cls = plugins.get_class( + PluginType.SESSION_ROUTING, str(self.session_routing) + ) + options = routing_cls.Options(**self.session_routing_opts) + canonical = options.model_dump(mode="json", exclude_unset=True) + if canonical != self.session_routing_opts: + self.session_routing_opts = canonical + return self diff --git a/src/aiperf/config/flags/_converter_dataset.py b/src/aiperf/config/flags/_converter_dataset.py index 8cd397f628..56fb4265ec 100644 --- a/src/aiperf/config/flags/_converter_dataset.py +++ b/src/aiperf/config/flags/_converter_dataset.py @@ -272,6 +272,8 @@ def _flat_dataset_fields(cli: CLIConfig) -> dict[str, Any]: out[key] = value if _set(cli, "dataset_filters"): out["filters"] = _parse_dataset_filters(cli.dataset_filters) + if _set(cli, "graph_format") and cli.graph_format is not None: + out["graph_format"] = cli.graph_format return out @@ -303,11 +305,17 @@ def _resolve_entries(cli: CLIConfig) -> int | None: conversations (the runner recycles them to fill request_count). 3. ``cli.request_count`` (explicitly set) — fallback so a single ``--request-count N`` invocation produces ``N`` unique entries when - the user did not pin the conversation count separately. - - Returns None when none was explicitly set. The caller MUST omit the - ``entries`` key from the output dict in that case so the dataset class's - own Pydantic default applies (``SyntheticDataset.entries=100``; + the user did not pin the conversation count separately. **Synthetic + and public datasets only.** On a file dataset ``--request-count`` is + a recycle count, not a corpus size (a single trace can emit many + requests), so it must NOT cap ``FileDataset.entries`` — the file + corpus is sized by the file itself unless the user explicitly pins + it with ``--num-dataset-entries`` / ``--num-conversations`` (steps 1 + and 2, which apply to every dataset kind). See graph #1106. + + Returns None when no applicable source was explicitly set. The caller MUST + omit the ``entries`` key from the output dict in that case so the dataset + class's own Pydantic default applies (``SyntheticDataset.entries=100``; ``File/Public.entries=None``). Emitting ``entries=None`` into the dict would crash AIPerfConfig validation on synthetic (``int_type, got NoneType``). @@ -323,7 +331,12 @@ def _resolve_entries(cli: CLIConfig) -> int | None: if isinstance(v, list): return max(v) if v else None return v - if "request_count" in s: + # A file dataset draws its corpus from the input file; --request-count only + # recycles that corpus and must not cap it. Mirror _apply_dataset_type's + # kind precedence (public wins over file) so a public dataset still + # back-fills. Synthetic (neither field set) also back-fills. + is_file_dataset = bool(cli.input_file) and not cli.public_dataset + if "request_count" in s and not is_file_dataset: v = cli.request_count if isinstance(v, list): return max(v) if v else None @@ -422,7 +435,7 @@ def _apply_synthesis(d: dict[str, Any], cli: CLIConfig) -> None: Synthesis is only meaningful for trace-format file datasets (the Synthesizer is invoked from BaseTraceDatasetLoader). The synthesis - fields live flat on CLIConfig (post-Task-13), so we only emit a + fields live flat on CLIConfig, so we only emit a ``synthesis`` sub-dict when the resulting dataset is a FileDataset and at least one field was explicitly set or carries a non-default value. """ @@ -440,6 +453,10 @@ def _apply_synthesis(d: dict[str, Any], cli: CLIConfig) -> None: ("synthesis_output_len_multiplier", "output_len_multiplier"), ("synthesis_max_isl", "max_isl"), ("synthesis_max_osl", "max_osl"), + ("max_context_length", "max_context_length"), + ("allow_dataset_wrap", "allow_dataset_wrap"), + ("prompt_corpus", "corpus"), + ("synthesis_idle_gap_cap", "idle_gap_cap_seconds"), ): if cli_attr in set_fields: value = getattr(cli, cli_attr) diff --git a/src/aiperf/config/flags/_converter_endpoint.py b/src/aiperf/config/flags/_converter_endpoint.py index afe4773c61..e3e560bf99 100644 --- a/src/aiperf/config/flags/_converter_endpoint.py +++ b/src/aiperf/config/flags/_converter_endpoint.py @@ -38,6 +38,21 @@ def _endpoint_template_from_extra( } +def _parse_routing_opts(values: list[str]) -> dict[str, str]: + opts: dict[str, str] = {} + for item in values: + key, separator, value = item.partition("=") + key, value = key.strip(), value.strip() + if not separator or not key or not value: + raise ValueError( + f"Invalid --session-routing-opt {item!r}; expected non-empty key=value" + ) + if key in opts: + raise ValueError(f"Duplicate --session-routing-opt key {key!r}") + opts[key] = value + return opts + + def _endpoint_template_fallback(endpoint: dict[str, Any]) -> None: from aiperf.plugin.enums import EndpointType @@ -68,10 +83,12 @@ def _endpoint_template_fallback(endpoint: dict[str, Any]) -> None: "transport": "transport", "use_legacy_max_tokens": "use_legacy_max_tokens", "use_server_token_count": "use_server_token_count", + "cache_bust": "cache_bust", "connection_reuse_strategy": "connection_reuse", "download_video_content": "download_video_content", "request_content_type": "request_content_type", "session_header": "session_header", + "session_routing": "session_routing", } @@ -97,6 +114,8 @@ def build_endpoint(cli: CLIConfig) -> dict[str, Any]: extra = dict(cli.extra_inputs) _endpoint_template_from_extra(endpoint, extra) endpoint["extra"] = extra + if "session_routing_opt" in cli_set and cli.session_routing_opt: + endpoint["session_routing_opts"] = _parse_routing_opts(cli.session_routing_opt) _endpoint_template_fallback(endpoint) return endpoint diff --git a/src/aiperf/config/flags/_converter_profiling.py b/src/aiperf/config/flags/_converter_profiling.py index da57905542..ddc6d97e7c 100644 --- a/src/aiperf/config/flags/_converter_profiling.py +++ b/src/aiperf/config/flags/_converter_profiling.py @@ -281,11 +281,22 @@ def _validate_profiling(prof: dict[str, Any], cli: CLIConfig) -> None: if ( not any(k in prof for k in ("requests", "duration", "sessions")) and prof["type"] != PhaseType.FIXED_SCHEDULE + and not _cli_is_graph_workload(cli) ): - # Why: when no bound is given for an unbounded run, default to - # 10 requests so the run terminates in a reasonable time. - # Deliberate override of the PhaseConfig default (which would - # leave it unbounded). + # Why: when no bound is given for a non-graph unbounded run, default to + # 10 requests so the run terminates in a reasonable time. Deliberate + # override of the PhaseConfig default (which would leave it unbounded). + # + # A BARE graph run is deliberately EXCLUDED here so it stays UNBOUNDED + # (no auto-10): it does a SINGLE CORPUS PASS -- each loaded trace runs + # exactly once, then the lanes stop (the graph strategy's + # ``_recycle_has_stop_condition`` is False with no numeric bound). The + # corpus size is not knowable here (weka HF is streamed; the + # max-context-length filter runs at parse time), so the no-stop + # concurrency phase validates against the graph dataset in + # ``check_phase_dataset_compatibility`` -- mirroring FixedSchedulePhase + # inferring its stop from the trace. Injecting ``requests=10`` would + # truncate the benchmark to 10 dispatches. prof.setdefault("requests", 10) delay_set = "request_cancellation_delay" in cli.model_fields_set if cli.request_cancellation_rate: @@ -304,6 +315,35 @@ def _validate_profiling(prof: dict[str, Any], cli: CLIConfig) -> None: ) +def _cli_is_graph_workload(cli: CLIConfig) -> bool: + """True when the CLI invocation targets a graph workload. + + Mirrors the dataset resolver's graph check + (``config.dataset.resolver.DatasetResolver._resolve_one``): an explicit + ``--graph-format`` forces graph mode (including ``native``, which + auto-detection excludes), otherwise the ``--input-file`` path is sniffed + against the graph-adapter registry via + :func:`~aiperf.dataset.graph.workload_detect.is_graph_workload_path` + (which excludes ``native``/``dag_jsonl``). A graph corpus size is NOT + knowable at CLI-conversion time, so this only decides whether a bare run + stays unbounded (single corpus pass) instead of taking the auto-10 bound. + Any detection failure degrades to False (keep the non-graph auto-10). + """ + if getattr(cli, "graph_format", None) is not None: + return True + input_file = getattr(cli, "input_file", None) + if input_file is None: + return False + from pathlib import Path + + from aiperf.dataset.graph.workload_detect import is_graph_workload_path + + try: + return is_graph_workload_path(Path(input_file)) + except Exception: + return False + + def _maybe_auto_promote_trace( prof: dict[str, Any], cli: CLIConfig, file_path: Path | None ) -> None: @@ -366,7 +406,26 @@ def _maybe_set_dag_root_sessions( def _apply_dataset_aware_autodefaults(prof: dict[str, Any], cli: CLIConfig) -> None: - """Apply dataset-sensitive CLI defaults for trace/fixed/dag datasets.""" + """CLI-only counterpart of the v1 dataset-aware autodefaults. + + Three behaviors, all conditional on the user's CLI invocation supplying + a dataset file (no behavior change for YAML-only configs, which are + expected to be complete): + + 1. Trace auto-promotion: a trace ``--custom-dataset-type`` whose first + record carries a ``timestamp`` field flips the phase to + fixed_schedule unless the user passed ``--no-fixed-schedule``. + 2. fixed_schedule autodefault: when a fixed_schedule phase has no + stop condition, fill ``requests`` from the dataset record count + (single-pass). + 3. Forking-dataset autodefault: when the dataset is ``dag_jsonl`` and + no stop condition is set, fill ``sessions`` from the DAG root + count so the run executes each root once instead of truncating + mid-tree. + + Bare-string ``--custom-dataset-type`` (no ``--input-file``) is a no-op + for I/O-dependent steps. + """ from aiperf.config.phases import PhaseType @@ -454,7 +513,19 @@ def _count_dataset_records(file_path: object) -> int: def build_profiling(cli: CLIConfig) -> dict[str, Any]: - """Produce the canonical profiling-phase dict from ``cli``.""" + """Produce the canonical profiling-phase dict from ``cli``. + + Reads load-generator settings (concurrency, rate, ramps, cancellation), + schedule/replay flags, and session-turn count directly from ``cli`` + (all fields are top-level on the flat CLIConfig). Returns a dict whose ``type`` + is one of ``PhaseType.{CONCURRENCY, POISSON, GAMMA, CONSTANT, + USER_CENTRIC, FIXED_SCHEDULE}`` plus the keys mapped by + ``_PROF_FIELD_ROUTES`` and any ramp/cancellation sub-dicts. + + Raises: + ValueError: when USER_CENTRIC mode is selected but + ``conversation_turn_mean`` is < 2. + """ from aiperf.config.phases import PhaseType fields_set = cli.model_fields_set diff --git a/src/aiperf/config/flags/_section_fields.py b/src/aiperf/config/flags/_section_fields.py index deae2cce04..924191f2ed 100644 --- a/src/aiperf/config/flags/_section_fields.py +++ b/src/aiperf/config/flags/_section_fields.py @@ -17,6 +17,7 @@ ENDPOINT_FIELDS: frozenset[str] = frozenset( { "api_key", + "cache_bust", "connection_reuse_strategy", "custom_endpoint", "download_video_content", @@ -24,6 +25,8 @@ "model_selection_strategy", "request_content_type", "session_header", + "session_routing", + "session_routing_opt", "streaming", "timeout_seconds", "transport", @@ -42,8 +45,11 @@ { # ----- top-level input flat fields ----- "custom_dataset_type", + "graph_format", "dataset_filters", "dataset_sampling_strategy", + "max_context_length", + "allow_dataset_wrap", "extra_inputs", "input_file", "fixed_schedule", @@ -120,6 +126,8 @@ "synthesis_output_len_multiplier", "synthesis_max_isl", "synthesis_max_osl", + "prompt_corpus", + "synthesis_idle_gap_cap", } ) diff --git a/src/aiperf/config/flags/cli_config.py b/src/aiperf/config/flags/cli_config.py index bd469b6998..d677dd1ce6 100644 --- a/src/aiperf/config/flags/cli_config.py +++ b/src/aiperf/config/flags/cli_config.py @@ -17,7 +17,7 @@ AIPerfConfig. The converter at aiperf.config.flags.converter translates a populated CLIConfig into the canonical AIPerfConfig. -This file is intentionally large (~3200 LOC, ~200 fields) — every CLI flag +This file is intentionally large (~3900 LOC, ~230 fields) — every CLI flag is a top-level field by design, so size scales linearly with field count. The flat shape is the post-flatten architecture. Section dividers group fields by their CLIParameter ``Groups.X``. The pydantic-fields @@ -39,6 +39,7 @@ from aiperf.common.enums import ( AIPerfLogLevel, AudioFormat, + CacheBustTarget, ConnectionReuseStrategy, ConvergenceStat, ExportLevel, @@ -79,8 +80,10 @@ DatasetSamplingStrategy, EndpointType, GPUTelemetryCollectorType, + GraphAdapterType, PublicDatasetType, SearchPlannerType, + SessionRoutingType, TransportType, UIType, URLSelectionStrategy, @@ -327,6 +330,26 @@ class CLIConfig(BaseConfig): ), ] = EndpointDefaults.USE_SERVER_TOKEN_COUNT + cache_bust: Annotated[ + CacheBustTarget, + Field( + description=( + "Where to inject a per-trace-instance cache-bust marker on the " + "graph-IR replay path. 'none' (default): recorded bytes are sent " + "verbatim. 'first_turn_prefix': prepend a '[rid:<12hex>]' marker to " + "the first user turn. The marker is shared across all turns of one " + "trace instance (its own prefix stays cacheable) but differs between " + "distinct instances and resets on recycle, so the inference server's " + "KV cache sees a distinct prefix per trace instance, defeating " + "cross-instance prefix-cache hits that would inflate cache-hit metrics." + ), + ), + CLIParameter( + name=("--cache-bust",), + group=Groups.ENDPOINT, + ), + ] = EndpointDefaults.CACHE_BUST + connection_reuse_strategy: Annotated[ ConnectionReuseStrategy, Field( @@ -389,6 +412,42 @@ class CLIConfig(BaseConfig): ), ] = None + session_routing: Annotated[ + SessionRoutingType | None, + Field( + description=( + "Session-aware routing mode: stamps per-session identity on " + "every request for router affinity. Built-ins: dynamo_headers " + "(X-Dynamo-Session-ID + parent header), dynamo_nvext " + "(nvext.session_control bind/close request-body metadata), " + "smg_routing_key (X-SMG-Routing-Key for the SGLang Model Gateway " + "manual policy), session_id_header (custom additive header). " + "Parameterize with --session-routing-opt." + ), + ), + CLIParameter( + name=("--session-routing",), + group=Groups.ENDPOINT, + ), + ] = None + + session_routing_opt: Annotated[ + list[str], + Field( + default_factory=list, + description=( + "Repeatable key=value option for the selected --session-routing " + "mode (e.g. --session-routing-opt timeout_seconds=600), " + "validated against the plugin's Options model." + ), + ), + CLIParameter( + name=("--session-routing-opt",), + consume_multiple=True, + group=Groups.ENDPOINT, + ), + ] + @property def url(self) -> str: """Return the first URL for backward compatibility.""" @@ -567,6 +626,23 @@ def url(self) -> str: ), ] = None + graph_format: Annotated[ + GraphAdapterType | None, + Field( + description="Select the agentic-workload format for `--input-file`, " + "overriding graph-adapter auto-detection. Registered names: " + "`dynamo_trace`, `weka_trace`, `dag_jsonl`, `native`. Requires " + "`--input-file`. Independent of `--custom-dataset-type` (that selects " + "a custom dataset loader; this selects a graph adapter). Use `native` " + "to load a hand-authored Graph IR workload file, or `dag_jsonl` to " + "run a DAG conversation file on the graph runtime.", + ), + CLIParameter( + name=("--graph-format",), + group=Groups.INPUT, + ), + ] = None + dataset_sampling_strategy: Annotated[ DatasetSamplingStrategy | None, Field( @@ -582,6 +658,40 @@ def url(self) -> str: ), ] = None + max_context_length: Annotated[ + int | None, + Field( + ge=1, + description="Maximum per-trace context length (tokens) for graph-plane " + "dataset selection. Graph traces whose input+output context would exceed " + "this cap are excluded from the eligible set before selection. Applies to " + "graph-adapter workloads (`--input-file` graph traces); ignored by " + "synthetic/public datasets. Unset (the default) applies no context-length " + "filter.", + ), + CLIParameter( + name=("--max-context-length",), + group=Groups.INPUT, + ), + ] = None + + allow_dataset_wrap: Annotated[ + bool | None, + Field( + description="Allow graph-plane dataset selection to wrap (reuse the finite " + "trace pool) when `--request-count` exceeds the number of eligible traces. " + "`--allow-dataset-wrap` forces wrapping on; `--no-allow-dataset-wrap` forces " + "it off. Unset (the default, None) defers to a derived default computed at " + "resolution time and surfaced on `run.resolved.allow_dataset_wrap`. Applies " + "to graph-adapter workloads; ignored by synthetic/public datasets.", + ), + CLIParameter( + name=("--allow-dataset-wrap",), + group=Groups.INPUT, + negative="--no-allow-dataset-wrap", + ), + ] = None + random_seed: Annotated[ int | None, Field( @@ -757,7 +867,9 @@ def url(self) -> str: ge=1, description="Total number of unique entries to generate for the dataset. Each entry represents one user message that can be " "used as a turn in conversations. Entries are reused across conversations and turns according to `--dataset-sampling-strategy`. " - "Higher values provide more diversity.", + "Higher values provide more diversity. On the graph plane (graph-adapter workloads), this bounds the distinct traces selected: " + "unset means all eligible traces (after `--max-context-length` and other filters) are used; N means up to N distinct traces are " + "selected after filters. It does not count requests -- see `--request-count`.", ), CLIParameter( name=( @@ -1774,6 +1886,34 @@ def url(self) -> str: CLIParameter(name=("--synthesis-max-osl",), group=Groups.SYNTHESIS), ] = None + prompt_corpus: Annotated[ + Literal["coding", "sonnet"] | None, + Field( + default=None, + description="Corpus backing recorded graph (`weka_trace`, `dynamo_trace`) " + "real-content synthesis. " + "`coding` (default when unset) uses the procedural CodingContentGenerator pool " + "the recorded weka workloads were captured against; `sonnet` uses the Shakespeare " + "pool, which yields matching token counts but different bytes (useful only to " + "reproduce golden fixtures built from that pool). Ignored by " + "non-graph datasets.", + ), + CLIParameter(name=("--prompt-corpus",), group=Groups.SYNTHESIS), + ] = None + + synthesis_idle_gap_cap: Annotated[ + float | None, + Field( + default=None, + ge=0.0, + description="Per-trace idle-gap cap (seconds) for recorded graph " + "(`weka_trace`, `dynamo_trace`) replay. Recorded inter-request idle gaps longer " + "than this are compressed to the cap so a recorded multi-hour gap does not park " + "the replay. Defaults to 60s when unset. Ignored by non-graph datasets.", + ), + CLIParameter(name=("--synthesis-idle-gap-cap",), group=Groups.SYNTHESIS), + ] = None + ############################################################################## # Load Generator ############################################################################## @@ -1895,7 +2035,10 @@ def url(self) -> str: "on the timing mode and dataset size. For synthetic datasets, this will be `max(10, concurrency * 2)`. " "Pass a comma-separated list (e.g. `--request-count 100,500,1000`) to sweep over multiple " "request counts; the converter promotes the list to a sweep on phases.profiling.requests " - "before AIPerfConfig validation.", + "before AIPerfConfig validation. On the graph plane (graph-adapter workloads), this counts individual " + "requests, not traces -- a single trace can emit many requests, so `--request-count` does not bound the " + "number of distinct traces selected (use `--num-dataset-entries` for that). When it exceeds the eligible " + "trace pool, `--allow-dataset-wrap` controls whether selection wraps to reuse traces.", ), BeforeValidator(parse_int_or_int_list), CLIParameter( @@ -1912,7 +2055,9 @@ def url(self) -> str: Field( gt=0, description="Duration in seconds to ramp session concurrency from 1 to target. " - "Useful for gradual warm-up of the target system.", + "On graph-IR replay (weka/dynamo/dag traces) this ramps LANE admission: " + "concurrent trace-instance lanes are admitted 1 -> --concurrency over the " + "duration. Useful for gradual warm-up of the target system.", ), CLIParameter( name=("--concurrency-ramp-duration",), @@ -2090,6 +2235,86 @@ def url(self) -> str: ), ] = None + agentic_cache_warmup_duration: Annotated[ + float | None, + Field( + gt=0, + description="Extended (cache-pressure) warmup duration in seconds for " + "the recorded-trace graph-IR replay path. After the boundary-priming warmup " + "drains, AIPerf continues the live post-t* replay with zero idle " + "delay and one-token outputs for this long, recycling fresh " + "templates onto freed lanes, then drains and hands each lane's " + "execution frontier to profiling (with residual delays), driving " + "the server KV cache to steady-state pressure before benchmarking. " + "If not set, no extended warmup runs.", + ), + CLIParameter( + name=("--agentic-cache-warmup-duration",), + group=Groups.WARMUP, + ), + ] = None + + burst_phase_starts: Annotated[ + bool, + Field( + description="Burst-vs-spread control for phase starts on the recorded-trace graph-IR " + "replay path. Default False = SPREAD: each WARMUP / PROFILING trace's " + "first firing is offset by its recorded (idle-gap-warped, <=60s) lead " + "from the t* snapshot instant. True = BURST: both phase starts collapse " + "into a synchronized burst -- every warmup priming credit and every " + "trace's earliest profiling firing fire at once at phase-time 0 (the " + "leading per-firing dispatch offset is dropped); relative inter-turn " + "delays after the first are still honored. Governs ONLY the two " + "phase-start dispatch patterns.", + ), + CLIParameter( + name=("--burst-phase-starts",), + group=Groups.WARMUP, + ), + ] = False + + trajectory_start_min_ratio: Annotated[ + float | None, + Field( + ge=0.0, + le=1.0, + description="Lower bound (fraction of each trace's recorded " + "wall-clock duration) of the per-trace t* snapshot-instant sampling " + "window for recorded graph replay. Unset resolves to 0.0. With both " + "bounds at 0.0 the window is OFF (full native replay). Per trace, " + "t* is drawn uniformly from [min, max] x trace_duration with a " + "trace-salted RNG: nodes firing before t* become cache-priming " + "WARMUP history, nodes at/after t* are PROFILING. Set min == max " + "for a deterministic t*. Must be <= --trajectory-start-max-ratio. " + "The inferencex-agentx-mvp scenario auto-applies 0.0 when unset " + "and locks an explicit conflicting value.", + ), + CLIParameter( + name=("--trajectory-start-min-ratio",), + group=Groups.LOAD_GENERATOR, + ), + ] = None + + trajectory_start_max_ratio: Annotated[ + float | None, + Field( + ge=0.0, + le=1.0, + description="Upper bound (fraction of each trace's recorded " + "wall-clock duration) of the per-trace t* snapshot-instant sampling " + "window for recorded graph replay. Unset resolves to 0.0 (window " + "OFF, full native replay); any value > 0 engages the t* snapshot " + "machinery (per-trace t* sampling plus the auto boundary-priming " + "warmup phase). See --trajectory-start-min-ratio for the sampling " + "contract. The inferencex-agentx-mvp scenario auto-applies 1.0 " + "when unset and locks an explicit conflicting value.", + ), + CLIParameter( + name=("--trajectory-start-max-ratio",), + group=Groups.LOAD_GENERATOR, + ), + ] = None + warmup_num_sessions: Annotated[ int | None, Field( @@ -3918,6 +4143,28 @@ def url(self) -> str: ), ] = False + scenario: Annotated[ + str | None, + Field( + description="Lock all benchmark invariants for a named scenario " + "(e.g. 'inferencex-agentx-mvp'). Conflicts with the locked " + "invariants raise ScenarioLockError at startup unless " + "--unsafe-override is also passed. Distinct from the sweep " + "``scenarios`` strategy (hand-picked named runs).", + ), + CLIParameter(name=("--scenario",), group=Groups.SCENARIO), + ] = None + + unsafe_override: Annotated[ + bool, + Field( + description="Convert scenario lock errors to warnings; stamps " + "submission_valid=false in the aggregate output. No-op without " + "--scenario.", + ), + CLIParameter(name=("--unsafe-override",), group=Groups.SCENARIO), + ] = False + ############################################################################## # Miscellaneous / Internal ############################################################################## diff --git a/src/aiperf/config/flags/converter.py b/src/aiperf/config/flags/converter.py index 28152ada02..ef8a231b9e 100644 --- a/src/aiperf/config/flags/converter.py +++ b/src/aiperf/config/flags/converter.py @@ -365,7 +365,7 @@ def _promote_cli_dataset_magic_lists( as a scalar placeholder so AIPerfConfig validation passes; each sweep variation overrides per-cell at expand time. The dataset converter is responsible for unwrapping the list at base-build time (see - ``_build_synthetic_prompts_block``). + ``_converter_dataset._build_prompts``). """ additions: dict[str, list[Any]] = {} for attr, path in _CLI_DATASET_MAGIC_LIST_PATHS: @@ -531,11 +531,11 @@ def _assemble_envelope_dict(cli: CLIConfig) -> dict[str, Any]: "endpoint": endpoint, "models": models, "phases": phases, - # Dataset name "main": kept in sync with _V1_DEFAULT_DATASET_NAME in + # Dataset name "main": kept in sync with _DEFAULT_DATASET_NAME in # search_recipes.builtins (which can't import from this module without # creating a load-order cycle through aiperf.config/__init__.py). - # If renaming, update both call sites and the regression test in - # tests/unit/search_recipes/test_grid_recipe_converter.py. + # If renaming, update both call sites and the tests that assert + # "datasets.main.*" sweep paths (tests/unit/config/, tests/unit/search_recipes/). "datasets": [{"name": "main", **ds}], "artifacts": artifacts, "gpu_telemetry": gpu_telemetry, @@ -559,6 +559,7 @@ def _assemble_envelope_dict(cli: CLIConfig) -> dict[str, Any]: artifacts["plot_required"] = plot_required _assemble_optional(nested, cli, recipe_output=recipe_output) + _apply_scenario_fields(nested, cli) _apply_recipe_sweep_parameters(nested, recipe_output, cli) _apply_recipe_scenarios(nested, recipe_output, cli) sweep_type = getattr(cli, "sweep_type", "grid") @@ -569,6 +570,30 @@ def _assemble_envelope_dict(cli: CLIConfig) -> dict[str, Any]: return _wrap_under_envelope(nested) +def _apply_scenario_fields(nested: dict[str, Any], cli: CLIConfig) -> None: + """Write the scenario-lock fields onto the benchmark body. + + ``scenario`` / ``unsafe_override`` are plain data on ``BenchmarkConfig`` + (the lock is applied later by ``ScenarioResolver``). They are NOT envelope + keys, so ``_wrap_under_envelope`` moves them under ``benchmark:``. Only + written when the user explicitly set them, to keep ``model_fields_set`` + clean for downstream "was this set?" checks. + """ + set_fields = cli.model_fields_set + if "scenario" in set_fields: + nested["scenario"] = cli.scenario + if "unsafe_override" in set_fields: + nested["unsafe_override"] = cli.unsafe_override + if "trajectory_start_min_ratio" in set_fields: + nested["trajectory_start_min_ratio"] = cli.trajectory_start_min_ratio + if "trajectory_start_max_ratio" in set_fields: + nested["trajectory_start_max_ratio"] = cli.trajectory_start_max_ratio + if "burst_phase_starts" in set_fields: + nested["burst_phase_starts"] = cli.burst_phase_starts + if "agentic_cache_warmup_duration" in set_fields: + nested["agentic_cache_warmup_duration"] = cli.agentic_cache_warmup_duration + + def _apply_parameter_sweep_meta_to_sweep( nested: dict[str, Any], cli: CLIConfig ) -> None: @@ -576,8 +601,8 @@ def _apply_parameter_sweep_meta_to_sweep( Schema-2.0 moved these from ``MultiRunConfig`` to the per-sweep config (``GridSweep.iteration_order`` / ``GridSweep.same_seed`` / - ``GridSweep.cooldown_seconds``). They are CLI-set on - ``CLIConfig.sweeping``; lift them onto whatever sweep block + ``GridSweep.cooldown_seconds``). They are CLI-set as flat top-level + ``CLIConfig`` fields (the SWEEPING_FIELDS section); lift them onto whatever sweep block ``_assemble_optional`` / ``_promote_magic_lists_to_sweep_block`` already produced. No-op when no sweep is in flight. diff --git a/src/aiperf/config/flags/resolver.py b/src/aiperf/config/flags/resolver.py index e976e14267..ac82a5dc74 100644 --- a/src/aiperf/config/flags/resolver.py +++ b/src/aiperf/config/flags/resolver.py @@ -168,6 +168,9 @@ def build_cli_overrides( out, "wandb", build_wandb(cli, base_enabled=wandb_base_enabled) ) + if "random_seed" in cli.model_fields_set: + out["random_seed"] = cli.random_seed + if "no_sweep_table" in cli.model_fields_set: out["no_sweep_table"] = cli.no_sweep_table diff --git a/src/aiperf/config/loader/parsing.py b/src/aiperf/config/loader/parsing.py index 3f154fc70f..f71624751d 100644 --- a/src/aiperf/config/loader/parsing.py +++ b/src/aiperf/config/loader/parsing.py @@ -293,19 +293,27 @@ def parse_str_or_list_of_positive_values(input: Any) -> list[Any]: def parse_file(value: str | None) -> Path | None: - """Parse an existing file/directory path from a CLI value. + """Parse an input source from a CLI value: a local file/dir or a weka HF id. + + Accepts an existing local file/directory path, or a recognized weka + HuggingFace ``org/name`` repo id (e.g. ``semianalysisai/cc-traces-weka-062126``) + that is NOT a local path. The HF id is a graph workload reference resolved by + the weka adapter at load time, not a filesystem path, so it must bypass the + file/dir existence gate. It is returned as a ``Path`` whose ``str()`` round-trips + to the original id (single-slash ``org/name`` is path-stable), which + ``resolve_graph_workload`` / ``WekaTraceAdapter.can_load`` then re-detect. Args: - value: Path string from CLI/config input. ``None`` or an empty string - disables the path and returns ``None``. + value: Path string or weka HF id from CLI/config input. ``None`` or an + empty string disables the source and returns ``None``. Returns: - A ``Path`` pointing to an existing file or directory, or ``None`` when no - value was provided. + A ``Path`` pointing to an existing file/directory or carrying a weka HF + repo id, or ``None`` when no value was provided. Raises: - ValueError: If ``value`` is not a string, or if the path does not exist as - a file or directory. + ValueError: If ``value`` is not a string, or if it is neither an existing + file/directory nor a recognized weka HF dataset id. """ if not value: @@ -316,8 +324,15 @@ def parse_file(value: str | None) -> Path | None: path = Path(value) if path.is_file() or path.is_dir(): return path - else: - raise ValueError(f"'{value}' is not a valid file or directory") + # A weka HF repo id is a valid graph workload source even though no such + # local path exists; the adapter resolves it from the HuggingFace hub/cache. + from aiperf.dataset.graph.adapters.weka.trace import ( + _looks_like_hf_dataset_id, + ) + + if _looks_like_hf_dataset_id(value): + return path + raise ValueError(f"'{value}' is not a valid file or directory") def validate_sequence_distribution(v: str | None) -> str | None: diff --git a/src/aiperf/config/phases.py b/src/aiperf/config/phases.py index c8b7551e4f..123b06570d 100644 --- a/src/aiperf/config/phases.py +++ b/src/aiperf/config/phases.py @@ -220,11 +220,14 @@ class BasePhaseConfig(AdaptiveScalePhaseMixin, BaseConfig): ), ] - # Subclasses set False to opt out (e.g. FixedSchedulePhase, where the - # stop condition is inferred from the dataset). Otherwise CLI users - # get autodefaults applied in the CLI->YAML converter (see - # ``aiperf.config.flags._converter_profiling``); YAML users must be - # explicit. + # Subclasses set False to opt out (e.g. FixedSchedulePhase, where the stop + # condition is inferred from the trace dataset). The required-stop check + # itself is NOT enforced here -- a phase alone cannot see its dataset's + # graph-ness. It lives in + # ``aiperf.config.resolution.predicates.check_phase_dataset_compatibility`` + # (run from ``BenchmarkConfig``), which exempts graph workloads: a bare + # graph run infers a single-corpus-pass stop from the loaded corpus, just + # as FixedSchedulePhase infers its stop from the trace. _stop_condition_required: ClassVar[bool] = True # ========================================================================= @@ -247,16 +250,6 @@ def _validate_phase_constraints(self) -> Self: ) if self.exclude_from_results != required: self.exclude_from_results = required - if ( - self._stop_condition_required - and self.requests is None - and self.duration is None - and self.sessions is None - ): - raise ValueError( - f"Phase '{self.name}': at least one of " - "'requests', 'duration', or 'sessions' must be specified" - ) if ( self.prefill_concurrency is not None and self.concurrency is not None diff --git a/src/aiperf/config/resolution/__init__.py b/src/aiperf/config/resolution/__init__.py index 78567b428a..2d7baba14a 100644 --- a/src/aiperf/config/resolution/__init__.py +++ b/src/aiperf/config/resolution/__init__.py @@ -37,6 +37,7 @@ ConfigResolver, ConfigResolverChain, GpuMetricsResolver, + ScenarioResolver, TimingResolver, TokenizerResolver, build_default_resolver_chain, @@ -52,6 +53,7 @@ "FailurePolicy", "GpuMetricsResolver", "ResolvedConfig", + "ScenarioResolver", "TimingResolver", "TokenizerResolver", "build_default_resolver_chain", diff --git a/src/aiperf/config/resolution/graph_dispatch_resolver.py b/src/aiperf/config/resolution/graph_dispatch_resolver.py new file mode 100644 index 0000000000..794f92766b --- /dev/null +++ b/src/aiperf/config/resolution/graph_dispatch_resolver.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Graph-dispatch config resolver. + +Runs AFTER :class:`~aiperf.config.resolution.resolvers.ScenarioResolver` so any +scenario-driven ``endpoint.cache_bust`` auto-fill is already in place when this +step derives the wrap default from it. Reads ``run.cfg`` and populates two +graph-plane fields on ``run.resolved``: + +- ``allow_dataset_wrap``: honors the explicit raw user value stashed on + ``dataset.synthesis.allow_dataset_wrap`` (``--allow-dataset-wrap`` / + ``--no-allow-dataset-wrap``); when unset, derives the default from the + resolved cache-bust target (``cache_bust != NONE``). +- ``dataset_sampling_strategy``: surfaces the default dataset's ``sampling`` + strategy (the linear path lands per-dataset strategies on + ``dataset_sampling_strategies``, but graph workloads skip file resolution). + +No-op for non-graph runs -- detection is the single memoized +``resolve_graph_workload`` accessor every consumer shares. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from aiperf.common.aiperf_logger import AIPerfLogger + +if TYPE_CHECKING: + from aiperf.config.resolution.plan import BenchmarkRun + +_logger = AIPerfLogger(__name__) + + +class GraphDispatchResolver: + """Derive graph-plane dispatch defaults after the scenario lock is applied.""" + + def resolve(self, run: BenchmarkRun) -> None: + """Populate ``allow_dataset_wrap`` + ``dataset_sampling_strategy``. + + No-op unless the run is a graph workload. Reads the (possibly + scenario-auto-filled) ``endpoint.cache_bust`` to derive the wrap default + when the user left ``--allow-dataset-wrap`` unset. + """ + from aiperf.dataset.graph.workload_detect import resolve_graph_workload + + if resolve_graph_workload(run) is None: + return + + from aiperf.common.enums import CacheBustTarget + + dataset = run.cfg.get_default_dataset() + + raw_wrap = self._raw_allow_dataset_wrap(dataset) + if raw_wrap is not None: + run.resolved.allow_dataset_wrap = bool(raw_wrap) + else: + run.resolved.allow_dataset_wrap = ( + run.cfg.endpoint.cache_bust != CacheBustTarget.NONE + ) + _logger.debug( + lambda: f"Resolved allow_dataset_wrap={run.resolved.allow_dataset_wrap} " + f"(explicit={raw_wrap is not None}, cache_bust={run.cfg.endpoint.cache_bust})" + ) + + sampling = getattr(dataset, "sampling", None) + if sampling is not None: + run.resolved.dataset_sampling_strategy = sampling + + @staticmethod + def _raw_allow_dataset_wrap(dataset: object) -> bool | None: + """Return the raw explicit ``synthesis.allow_dataset_wrap`` (None if unset). + + The CLI flag's raw value lands on the default dataset's synthesis block; + an absent synthesis block or absent field both read as unset (None), so + the caller falls back to the cache-bust-derived default. + """ + synthesis = getattr(dataset, "synthesis", None) + return getattr(synthesis, "allow_dataset_wrap", None) if synthesis else None diff --git a/src/aiperf/config/resolution/plan.py b/src/aiperf/config/resolution/plan.py index 960b69c0e7..6fa9c7a504 100644 --- a/src/aiperf/config/resolution/plan.py +++ b/src/aiperf/config/resolution/plan.py @@ -14,6 +14,11 @@ GPUTelemetryMode, SweepMode, ) + +# Direct-module import (not the aiperf.common.models package facade): the +# package __init__ pulls the full model surface and works here only by +# fragile init ordering; dataset_models.py is the precedent for this form. +from aiperf.common.models.base_models import AIPerfBaseModel from aiperf.common.redact import build_cli_command from aiperf.config.base import BaseConfig from aiperf.config.comm import BaseZMQCommunicationConfig @@ -301,6 +306,24 @@ def _check_trials_matches_num_runs_cap(self) -> Self: return self +class GraphWorkloadRef(AIPerfBaseModel): + """Resolved reference to a run's graph workload input. + + The single source of graph-ness for a run: populated once (config + resolver chain, or first ``resolve_graph_workload`` call) and read by + every consumer thereafter -- no per-consumer re-sniffing. + """ + + path: Path = Field( + description="Graph input file/dir, or the HF org/name id verbatim " + "(never .resolve()d -- HF ids are not filesystem paths)." + ) + format: str = Field( + description="Adapter format name (weka_trace / dynamo_trace / native " + "/ dag_jsonl)." + ) + + class ResolvedConfig(BaseModel): """Runtime-computed state populated after construction. @@ -388,6 +411,42 @@ class ResolvedConfig(BaseModel): description="Pre-built ZMQ communication config. " "Avoids rebuilding in every service's CommunicationMixin.", ) + scenario_outcome: Any = Field( + default=None, + description="ScenarioOutcome from the ScenarioResolver step, or None " + "when no --scenario was set. Typed Any to avoid importing " + "aiperf.common.scenario at config-module load (the resolver writes a " + "concrete ScenarioOutcome). Carries scenario_name / submission_valid / " + "submission_invalid_reasons surfaced by the metrics-JSON exporter.", + ) + graph_workload: GraphWorkloadRef | None = Field( + default=None, + description="Graph workload resolved at config resolution, or memoized " + "on first access; None until resolved for non-graph runs -- consumers " + "use resolve_graph_workload(run), never this field directly.", + ) + graph_workload_resolved: bool = Field( + default=False, + description="True once detection ran (distinguishes 'not a graph run' " + "from 'never checked').", + ) + allow_dataset_wrap: bool | None = Field( + default=None, + description="Effective graph-plane dataset-wrap policy for this run: " + "whether selection may reuse the finite trace pool to satisfy " + "`--request-count`. None until resolution derives it (the resolver sets " + "the derived default when the user left `--allow-dataset-wrap` unset, or " + "carries the explicit user True/False otherwise). Read by downstream " + "graph-selection consumers via `run.resolved.allow_dataset_wrap`.", + ) + dataset_sampling_strategy: DatasetSamplingStrategy | None = Field( + default=None, + description="Effective dataset sampling strategy for a graph workload, " + "surfaced from the default dataset's `sampling` field by " + "GraphDispatchResolver. None for non-graph runs (whose per-dataset " + "strategies land on `dataset_sampling_strategies` instead). Read by " + "graph-selection consumers via `run.resolved.dataset_sampling_strategy`.", + ) class BenchmarkRun(BaseModel): diff --git a/src/aiperf/config/resolution/predicates.py b/src/aiperf/config/resolution/predicates.py index 7719460b31..59c3d13bf2 100644 --- a/src/aiperf/config/resolution/predicates.py +++ b/src/aiperf/config/resolution/predicates.py @@ -205,6 +205,33 @@ def requires_multi_turn(phase_type: PhaseType) -> bool: # ============================================================================= +def is_graph_dataset(dataset: DatasetConfig) -> bool: + """True when a dataset resolves to a graph workload. + + Mirrors the dataset resolver's graph check + (``config.dataset.resolver.DatasetResolver._resolve_one``): an explicit + ``graph_format`` forces graph mode (including ``native``, which + auto-detection excludes), otherwise the file ``path`` is sniffed against + the graph-adapter registry via + :func:`~aiperf.dataset.graph.workload_detect.is_graph_workload_path` + (which excludes ``native``/``dag_jsonl``). Non-file datasets (synthetic / + public / inline) and any detection failure degrade to False. + """ + if getattr(dataset, "graph_format", None) is not None: + return True + raw_path = getattr(dataset, "path", None) + if raw_path is None: + return False + from pathlib import Path + + from aiperf.dataset.graph.workload_detect import is_graph_workload_path + + try: + return is_graph_workload_path(Path(raw_path)) + except Exception: + return False + + def check_phase_dataset_compatibility( phase: BasePhaseConfig, dataset: DatasetConfig, @@ -226,6 +253,23 @@ def check_phase_dataset_compatibility( """ errors: list[str] = [] + # Required-stop rule. A phase that requires a numeric stop + # (``_stop_condition_required``, True everywhere except FixedSchedulePhase) + # with none of requests/duration/sessions set is only valid against a graph + # workload: the graph plane infers a single-corpus-pass stop from the loaded + # corpus (each loaded trace runs exactly once, then the lanes stop), just as + # FixedSchedulePhase infers its stop from the trace. Non-graph datasets must + # carry an explicit stop condition. + if ( + phase._stop_condition_required + and get_stop_condition(phase) == "none" + and not is_graph_dataset(dataset) + ): + errors.append( + f"Phase '{phase_name}': at least one of " + "'requests', 'duration', or 'sessions' must be specified" + ) + if requires_sequential_sampling(phase.type) and is_file_dataset(dataset): sampling = get_sampling_strategy(dataset) if sampling != DatasetSamplingStrategy.SEQUENTIAL: diff --git a/src/aiperf/config/resolution/resolvers.py b/src/aiperf/config/resolution/resolvers.py index 4b32c74104..f9a14ce101 100644 --- a/src/aiperf/config/resolution/resolvers.py +++ b/src/aiperf/config/resolution/resolvers.py @@ -17,6 +17,7 @@ from aiperf.common.aiperf_logger import AIPerfLogger from aiperf.config.dataset.resolver import DatasetResolver +from aiperf.config.resolution.graph_dispatch_resolver import GraphDispatchResolver if TYPE_CHECKING: from aiperf.config.resolution.plan import BenchmarkRun @@ -35,6 +36,8 @@ "ConfigResolverChain", "DatasetResolver", "GpuMetricsResolver", + "GraphDispatchResolver", + "ScenarioResolver", "TimingResolver", "TokenizerResolver", "build_default_resolver_chain", @@ -374,6 +377,24 @@ def _validate_fixed_schedule_timing( ) +class ScenarioResolver: + """Apply the locked ``--scenario`` invariants to ``run.cfg`` / ``run.resolved``. + + Runs between :class:`DatasetResolver` and :class:`TimingResolver` so the + scenario's auto-filled phase durations are in place before TimingResolver + sums them. No-op when ``run.cfg.scenario`` is None. Delegates to + ``aiperf.common.scenario.apply_scenario``, which stores the + :class:`~aiperf.common.scenario.ScenarioOutcome` on + ``run.resolved.scenario_outcome``. + """ + + def resolve(self, run: BenchmarkRun) -> None: + """Apply the scenario lock (auto-fill defaults, validate invariants).""" + from aiperf.common.scenario import apply_scenario + + apply_scenario(run) + + def build_default_resolver_chain() -> ConfigResolverChain: """Build the default resolver chain for pre-bootstrap resolution.""" return ConfigResolverChain( @@ -383,6 +404,8 @@ def build_default_resolver_chain() -> ConfigResolverChain: GpuMetricsResolver(), CommConfigResolver(), DatasetResolver(), + ScenarioResolver(), + GraphDispatchResolver(), TimingResolver(), ] ) diff --git a/src/aiperf/config/schema/aiperf-config.schema.json b/src/aiperf/config/schema/aiperf-config.schema.json index e0d1be58e1..e821b7dd82 100644 --- a/src/aiperf/config/schema/aiperf-config.schema.json +++ b/src/aiperf/config/schema/aiperf-config.schema.json @@ -2986,6 +2986,108 @@ "default": null, "description": "Accuracy benchmarking configuration. When set, enables accuracy evaluation alongside performance profiling." }, + "scenario": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Lock all benchmark invariants for a named scenario (e.g. 'inferencex-agentx-mvp'). Plain data here; the lock is applied by the ScenarioResolver step in the pre-bootstrap resolver chain (auto-fills defaults, validates, raises ScenarioLockError on conflict). Distinct from the sweep ``scenarios`` strategy (SweepConfig), which expands hand-picked named runs -- this field is a single invariant LOCK, not a sweep.", + "title": "Scenario" + }, + "unsafeOverride": { + "default": false, + "description": "Convert scenario lock errors to warnings; stamps submission_valid=false in the resolved scenario outcome. No-op without ``scenario``. Plain data; consumed by the ScenarioResolver.", + "title": "Unsafeoverride", + "type": "boolean" + }, + "trajectoryStartMinRatio": { + "anyOf": [ + { + "maximum": 1.0, + "minimum": 0.0, + "type": "number" + }, + { + "type": "null" + }, + { + "type": "string", + "pattern": ".*\\{\\{.*\\}\\}.*", + "description": "Jinja2 template (e.g., '{{ variable }}')." + }, + { + "type": "string", + "pattern": ".*\\$\\{[A-Za-z_][A-Za-z0-9_]*(?::[^}]*)?\\}.*", + "description": "Environment variable (e.g., '${VAR}' or '${VAR:default}')." + } + ], + "default": null, + "description": "Lower bound (fraction of each trace's recorded wall-clock duration) of the per-trace t* snapshot-instant sampling window for recorded graph replay. Unset (None) resolves to 0.0. With both bounds at 0.0 the window is OFF (full native replay). The inferencex-agentx-mvp scenario auto-applies 0.0 when unset and locks an explicit conflicting value. Must be <= trajectory_start_max_ratio.", + "title": "Trajectorystartminratio", + "x-jinja2-supported": true + }, + "trajectoryStartMaxRatio": { + "anyOf": [ + { + "maximum": 1.0, + "minimum": 0.0, + "type": "number" + }, + { + "type": "null" + }, + { + "type": "string", + "pattern": ".*\\{\\{.*\\}\\}.*", + "description": "Jinja2 template (e.g., '{{ variable }}')." + }, + { + "type": "string", + "pattern": ".*\\$\\{[A-Za-z_][A-Za-z0-9_]*(?::[^}]*)?\\}.*", + "description": "Environment variable (e.g., '${VAR}' or '${VAR:default}')." + } + ], + "default": null, + "description": "Upper bound (fraction of each trace's recorded wall-clock duration) of the per-trace t* snapshot-instant sampling window for recorded graph replay. Unset (None) resolves to 0.0 (window OFF, full native replay); any value > 0 engages the t* snapshot machinery (per-trace t* sampling + auto boundary-priming warmup). The inferencex-agentx-mvp scenario auto-applies 1.0 when unset and locks an explicit conflicting value.", + "title": "Trajectorystartmaxratio", + "x-jinja2-supported": true + }, + "burstPhaseStarts": { + "default": false, + "description": "Burst-vs-spread control for phase starts on the recorded-trace graph-IR replay path (--burst-phase-starts). False = SPREAD (recorded leading offsets honored); True = BURST (both phase starts collapse into a synchronized burst).", + "title": "Burstphasestarts", + "type": "boolean" + }, + "agenticCacheWarmupDuration": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "number" + }, + { + "type": "null" + }, + { + "type": "string", + "pattern": ".*\\{\\{.*\\}\\}.*", + "description": "Jinja2 template (e.g., '{{ variable }}')." + }, + { + "type": "string", + "pattern": ".*\\$\\{[A-Za-z_][A-Za-z0-9_]*(?::[^}]*)?\\}.*", + "description": "Environment variable (e.g., '${VAR}' or '${VAR:default}')." + } + ], + "default": null, + "description": "Extended (cache-pressure) warmup duration in seconds for recorded graph replay (--agentic-cache-warmup-duration). After the boundary-priming warmup drains, the live post-t* replay continues with zero idle delay and one-token outputs for this long before profiling. None = no pressure stage.", + "title": "Agenticcachewarmupduration", + "x-jinja2-supported": true + }, "model": { "description": "Shorthand for 'models'. Accepts a single value which will be converted to a list/dict. Model configuration. Accepts a single model name string, a list of model names, or an advanced configuration with strategy and weighted items. All forms are normalized to ModelsAdvanced.", "oneOf": [ @@ -3278,7 +3380,7 @@ } ], "default": null, - "description": "Path to file or directory containing benchmark dataset. Can be absolute or relative. Mutually exclusive with `records:`. Supported formats depend on the format field: JSONL for single_turn/multi_turn, JSONL trace files for mooncake_trace/bailian_trace, Parquet for baseten_trace, directories for random_pool.", + "description": "Path to file or directory containing benchmark dataset. Can be absolute or relative. Mutually exclusive with `records:`. Supported formats depend on the format field: JSONL for single_turn/multi_turn, JSONL trace files for mooncake_trace/bailian_trace, Parquet for baseten_trace, directories for random_pool. Graph trace inputs (`.jsonl` / `.jsonl.gz` / segmented `.gz` directories) are recognized as graph workloads independently of `format` — auto-detected, or forced with `graph_format`.", "title": "Path" }, "records": { @@ -3311,7 +3413,19 @@ "format": { "$ref": "#/$defs/DatasetFormat", "default": "single_turn", - "description": "Dataset file format determining parsing logic and expected file structure. single_turn: JSONL with single prompt-response exchanges. multi_turn: JSONL with conversation history. mooncake_trace / bailian_trace / baseten_trace / burst_gpt_trace: timestamped trace files for replay. sagemaker_data_capture: JSONL captured by SageMaker DataCapture. random_pool: directory of reusable prompts." + "description": "Dataset file format determining parsing logic and expected file structure. single_turn: JSONL with single prompt-response exchanges. multi_turn: JSONL with conversation history. mooncake_trace / bailian_trace / baseten_trace / burst_gpt_trace: timestamped trace files for replay. sagemaker_data_capture: JSONL captured by SageMaker DataCapture. random_pool: directory of reusable prompts. This selects a *custom dataset loader* (conversation-style inputs) and is ignored for graph workloads, which are auto-detected from `path` or forced via `graph_format`." + }, + "graphFormat": { + "anyOf": [ + { + "$ref": "#/$defs/GraphAdapterType" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Force a graph-adapter format for this `--input-file`, overriding workload auto-detection. Accepts a registered graph adapter name (e.g. `dynamo_trace`, `weka_trace`, `dag_jsonl`, `native`). When set, the file is treated as a graph workload and parsed by the named adapter; this is the only way to load a hand-authored `native` graph file or run a `dag_jsonl` conversation file on the graph runtime, both otherwise excluded from auto-detection. None = auto-detect." }, "sampling": { "$ref": "#/$defs/DatasetSamplingStrategy", @@ -3328,7 +3442,7 @@ } ], "default": null, - "description": "Trace synthesis/transformation configuration. Allows scaling timestamps and token lengths before replay. Applies to trace formats such as mooncake_trace and baseten_trace, except speedup_ratio, which is rejected for baseten_trace (use replay_speedup / --replay-speedup to scale replay pacing there)." + "description": "Trace synthesis/transformation configuration. Allows scaling timestamps and token lengths before replay. Applies to trace formats such as mooncake_trace and baseten_trace, except speedup_ratio, which is rejected for baseten_trace (use replay_speedup / --replay-speedup to scale replay pacing there). Also used by graph (weka) replay, whose `synthesis.max_osl` caps the recorded output length (`--synthesis-max-osl`)." }, "entries": { "anyOf": [ @@ -7550,6 +7664,15 @@ } ] }, + "CacheBustTarget": { + "description": "Where (and how) to inject a per-session cache-bust marker.\n\nA unique marker injected into a session's first-turn prompt PREFIX makes the\ninference server's KV cache see a distinct prefix per session, defeating\ncross-session prefix-cache hits that would otherwise inflate cache-hit\nmetrics. The graph snapshot-replay path\nneeds only the ``FIRST_TURN_PREFIX`` and ``NONE`` variants.", + "enum": [ + "none", + "first_turn_prefix" + ], + "title": "CacheBustTarget", + "type": "string" + }, "CancellationConfig": { "additionalProperties": false, "description": "Configuration for request cancellation testing.\n\nEnables testing server behavior when clients cancel requests mid-flight.", @@ -9229,6 +9352,11 @@ "title": "Useservertokencount", "type": "boolean" }, + "cacheBust": { + "$ref": "#/$defs/CacheBustTarget", + "default": "none", + "description": "Where to inject a per-trace-instance cache-bust marker on the graph-IR replay path. 'none' (default): recorded bytes sent verbatim. 'first_turn_prefix': prepend a '[rid:<12hex>]' marker to the first user turn, shared across the trace instance's turns but distinct per instance and reset on recycle, so the inference server's KV cache sees a distinct prefix per trace instance, defeating cross-instance prefix-cache hits." + }, "template": { "anyOf": [ { @@ -9286,6 +9414,24 @@ "description": "HTTP header name used to carry the per-session affinity identifier. When set, replaces the default `X-Correlation-ID` header. Useful when the inference server expects a custom session-affinity header (e.g. `--session-header X-Session-ID`).", "title": "Sessionheader" }, + "sessionRouting": { + "anyOf": [ + { + "$ref": "#/$defs/SessionRoutingType" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Session-routing transform stamping per-session identity on every request so an external router (SGLang Model Gateway, Dynamo, generic session-affinity LBs) can pin a session's turns to one replica. One mode per run; parameterize with --session-routing-opt key=value." + }, + "sessionRoutingOpts": { + "additionalProperties": true, + "description": "Mode-specific options for --session-routing, validated against the selected plugin's Options model (unknown keys and invalid values are rejected at startup). Values are canonicalized/coerced to the plugin's Options model types at validation, so downstream consumers see typed values.", + "title": "Sessionroutingopts", + "type": "object" + }, "waitForModelTimeout": { "default": 0.0, "description": "Enable a pre-flight readiness probe by setting this to a positive value (seconds). aiperf applies this timeout to each URL/model probe before starting the benchmark, aborting with a non-zero exit if any probe times out. For multiple URLs or models, worst-case wall-clock time can be roughly this timeout multiplied by the number of URL/model probes. The probe strategy is controlled by `--wait-for-model-mode`, which defaults to sending a 1-token inference request. 0 (default) disables the probe. Eliminates the need for external shell-based readiness loops in containers and Kubernetes recipes.", @@ -9437,7 +9583,7 @@ } ], "default": null, - "description": "Path to file or directory containing benchmark dataset. Can be absolute or relative. Mutually exclusive with `records:`. Supported formats depend on the format field: JSONL for single_turn/multi_turn, JSONL trace files for mooncake_trace/bailian_trace, Parquet for baseten_trace, directories for random_pool.", + "description": "Path to file or directory containing benchmark dataset. Can be absolute or relative. Mutually exclusive with `records:`. Supported formats depend on the format field: JSONL for single_turn/multi_turn, JSONL trace files for mooncake_trace/bailian_trace, Parquet for baseten_trace, directories for random_pool. Graph trace inputs (`.jsonl` / `.jsonl.gz` / segmented `.gz` directories) are recognized as graph workloads independently of `format` — auto-detected, or forced with `graph_format`.", "title": "Path" }, "records": { @@ -9470,7 +9616,19 @@ "format": { "$ref": "#/$defs/DatasetFormat", "default": "single_turn", - "description": "Dataset file format determining parsing logic and expected file structure. single_turn: JSONL with single prompt-response exchanges. multi_turn: JSONL with conversation history. mooncake_trace / bailian_trace / baseten_trace / burst_gpt_trace: timestamped trace files for replay. sagemaker_data_capture: JSONL captured by SageMaker DataCapture. random_pool: directory of reusable prompts." + "description": "Dataset file format determining parsing logic and expected file structure. single_turn: JSONL with single prompt-response exchanges. multi_turn: JSONL with conversation history. mooncake_trace / bailian_trace / baseten_trace / burst_gpt_trace: timestamped trace files for replay. sagemaker_data_capture: JSONL captured by SageMaker DataCapture. random_pool: directory of reusable prompts. This selects a *custom dataset loader* (conversation-style inputs) and is ignored for graph workloads, which are auto-detected from `path` or forced via `graph_format`." + }, + "graphFormat": { + "anyOf": [ + { + "$ref": "#/$defs/GraphAdapterType" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Force a graph-adapter format for this `--input-file`, overriding workload auto-detection. Accepts a registered graph adapter name (e.g. `dynamo_trace`, `weka_trace`, `dag_jsonl`, `native`). When set, the file is treated as a graph workload and parsed by the named adapter; this is the only way to load a hand-authored `native` graph file or run a `dag_jsonl` conversation file on the graph runtime, both otherwise excluded from auto-detection. None = auto-detect." }, "sampling": { "$ref": "#/$defs/DatasetSamplingStrategy", @@ -9487,7 +9645,7 @@ } ], "default": null, - "description": "Trace synthesis/transformation configuration. Allows scaling timestamps and token lengths before replay. Applies to trace formats such as mooncake_trace and baseten_trace, except speedup_ratio, which is rejected for baseten_trace (use replay_speedup / --replay-speedup to scale replay pacing there)." + "description": "Trace synthesis/transformation configuration. Allows scaling timestamps and token lengths before replay. Applies to trace formats such as mooncake_trace and baseten_trace, except speedup_ratio, which is rejected for baseten_trace (use replay_speedup / --replay-speedup to scale replay pacing there). Also used by graph (weka) replay, whose `synthesis.max_osl` caps the recorded output length (`--synthesis-max-osl`)." }, "entries": { "anyOf": [ @@ -11022,6 +11180,16 @@ } ] }, + "GraphAdapterType": { + "enum": [ + "dynamo_trace", + "weka_trace", + "dag_jsonl", + "native" + ], + "title": "GraphAdapterType", + "type": "string" + }, "GridSweep": { "additionalProperties": false, "description": "Grid sweep - all combinations of parameters (Cartesian product).", @@ -14784,6 +14952,16 @@ "title": "ServiceRunType", "type": "string" }, + "SessionRoutingType": { + "enum": [ + "dynamo_headers", + "dynamo_nvext", + "smg_routing_key", + "session_id_header" + ], + "title": "SessionRoutingType", + "type": "string" + }, "SobolSweep": { "additionalProperties": false, "description": "Sobol quasi-Monte-Carlo sweep — low-discrepancy joint coverage.\n\nNote on ``scramble=False``: scipy's raw Sobol sequence starts at the\norigin (the first row is exactly all-zeros), so the first variant\nwill land on ``(lo, lo, ...)`` for every dimension. With small\n``samples`` this produces degenerate corner-only coverage. Prefer\nthe default ``scramble=True`` (Owen-scrambled) unless you have a\nspecific reason to want raw Sobol points.", @@ -15133,6 +15311,86 @@ "description": "Maximum output sequence length cap. Traces with output_length > max_osl are capped to this value (not filtered).", "title": "Maxosl", "x-jinja2-supported": true + }, + "maxContextLength": { + "anyOf": [ + { + "minimum": 1, + "type": "integer" + }, + { + "type": "null" + }, + { + "type": "string", + "pattern": ".*\\{\\{.*\\}\\}.*", + "description": "Jinja2 template (e.g., '{{ variable }}')." + }, + { + "type": "string", + "pattern": ".*\\$\\{[A-Za-z_][A-Za-z0-9_]*(?::[^}]*)?\\}.*", + "description": "Environment variable (e.g., '${VAR}' or '${VAR:default}')." + } + ], + "default": null, + "description": "Maximum per-trace context length (tokens) for graph-plane dataset selection (`--max-context-length`). Traces whose input+output context would exceed this cap are excluded from selection. None (the default) applies no context-length filter. Raw explicit value carried verbatim from the CLI; the derived selection default is computed elsewhere. Ignored by non-graph datasets.", + "title": "Maxcontextlength", + "x-jinja2-supported": true + }, + "allowDatasetWrap": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Whether graph-plane dataset selection may wrap (reuse the finite trace pool) to satisfy `--request-count` (`--allow-dataset-wrap` / `--no-allow-dataset-wrap`). None (the default) means unset -- the effective value is derived downstream and surfaced on `run.resolved.allow_dataset_wrap`; an explicit True/False here is the raw user intent carried verbatim so the resolver can distinguish unset from explicit. Ignored by non-graph datasets.", + "title": "Allowdatasetwrap" + }, + "idleGapCapSeconds": { + "anyOf": [ + { + "minimum": 0.0, + "type": "number" + }, + { + "type": "null" + }, + { + "type": "string", + "pattern": ".*\\{\\{.*\\}\\}.*", + "description": "Jinja2 template (e.g., '{{ variable }}')." + }, + { + "type": "string", + "pattern": ".*\\$\\{[A-Za-z_][A-Za-z0-9_]*(?::[^}]*)?\\}.*", + "description": "Environment variable (e.g., '${VAR}' or '${VAR:default}')." + } + ], + "default": 60.0, + "description": "Per-trace idle-gap cap (seconds) for recorded graph (`weka_trace`, `dynamo_trace`) replay (`--synthesis-idle-gap-cap`). Every recorded inter-request idle gap longer than this is compressed to the cap, shifting later requests left by the excess, so a recorded multi-hour idle gap does not park the replay. Default 60.0. Set to null to disable warping and replay the raw recorded timestamps. Ignored by non-graph datasets.", + "title": "Idlegapcapseconds", + "x-jinja2-supported": true + }, + "corpus": { + "anyOf": [ + { + "enum": [ + "coding", + "sonnet" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Corpus backing recorded graph (`weka_trace`, `dynamo_trace`) real-content synthesis (`--prompt-corpus`). `coding` (the default when unset) uses the procedural CodingContentGenerator pool -- the same corpus the recorded weka workloads were captured against. `sonnet` uses the Shakespeare PromptGenerator pool, which yields matching token counts but different bytes (useful only to reproduce golden fixtures built from that pool). Ignored by non-graph datasets.", + "title": "Corpus" } }, "title": "SynthesisConfig", diff --git a/src/aiperf/controller/multiprocess_service_manager.py b/src/aiperf/controller/multiprocess_service_manager.py index df6f3f972b..e98a44c1e6 100644 --- a/src/aiperf/controller/multiprocess_service_manager.py +++ b/src/aiperf/controller/multiprocess_service_manager.py @@ -95,8 +95,9 @@ async def run_service( process.start() self.debug( - lambda pid=process.pid, - type=service_type: f"Service {type} started as process (pid: {pid})" + lambda pid=process.pid, type=service_type: ( + f"Service {type} started as process (pid: {pid})" + ) ) self.multi_process_info.append( @@ -269,7 +270,9 @@ async def _wait_for_process(self, info: MultiProcessRunInfo) -> None: info.process.kill() else: self.debug( - lambda: f"Service {info.service_type} process stopped (pid: {info.process.pid})" + lambda: ( + f"Service {info.service_type} process stopped (pid: {info.process.pid})" + ) ) async def wait_for_all_services_start( diff --git a/src/aiperf/credit/callback_handler.py b/src/aiperf/credit/callback_handler.py index 19da73c6a2..547b278464 100644 --- a/src/aiperf/credit/callback_handler.py +++ b/src/aiperf/credit/callback_handler.py @@ -31,6 +31,19 @@ from aiperf.timing.phase.stop_conditions import StopConditionChecker from aiperf.timing.strategies.core import TimingStrategyProtocol +# A graph-return observer receives ``(credit, error, cancelled)`` for every +# graph-IR credit return, fired UNCONDITIONALLY (independent of phase-send +# gating). ``GraphIRReplayStrategy._on_graph_return`` is the production +# implementation (routing to the owning ``CreditDispatchAdapter.resolve``). +GraphReturnObserver = Callable[["Credit", "str | None", bool], None] + +# A graph-first-token observer receives the ``FirstToken`` message for every +# graph-IR credit's TTFT event (post-TTFT first-token anchoring), fired +# independent of phase-handler gating. +# ``GraphIRReplayStrategy._on_graph_first_token`` is the production +# implementation (routing to the owning adapter's ``on_first_token``). +GraphFirstTokenObserver = Callable[["FirstToken"], None] + _logger = AIPerfLogger(__name__) @@ -88,6 +101,69 @@ def __init__(self, concurrency_manager: ConcurrencyManager) -> None: self._concurrency_manager = concurrency_manager self._phase_handlers: dict[CreditPhase, PhaseCallbackContext] = {} self._branch_orchestrator: BranchOrchestrator | None = None + self._graph_return_observer: GraphReturnObserver | None = None + self._graph_first_token_observer: GraphFirstTokenObserver | None = None + + def set_graph_return_observer(self, observer: GraphReturnObserver | None) -> None: + """Register (or detach) the unconditional graph-IR return observer. + + Called by ``GraphIRReplayStrategy.setup_phase`` (through the + runner-injected registrar) before phase start; teardown detaches via + :meth:`clear_graph_return_observer` (compare-and-clear), falling back + to ``None`` here only when no unregister channel is wired (unit + harness). The observer fires + for every credit return carrying a ``trace_id`` (a graph credit), + BEFORE the gated ``strategy.handle_credit_return`` path -- which is + skipped once ``can_send_any_turn()`` is False and would otherwise + strand in-flight dispatch Futures for already-issued graph credits. + """ + self._graph_return_observer = observer + + def set_graph_first_token_observer( + self, observer: GraphFirstTokenObserver | None + ) -> None: + """Register (or detach) the graph-IR first-token observer. + + Called by ``GraphIRReplayStrategy.setup_phase`` (through the + runner-injected registrar) before phase start; teardown detaches via + :meth:`clear_graph_first_token_observer` (compare-and-clear), falling + back to ``None`` here only when no unregister channel is wired (unit + harness). The observer fires + in ``on_first_token`` for every ``FirstToken`` carrying a ``trace_id`` + (a graph credit), AFTER the prefill-slot release block and independent + of phase-handler gating -- so a graph node's TTFT anchor is delivered + even when prefill limiting is off or the phase can no longer send. + """ + self._graph_first_token_observer = observer + + def clear_graph_return_observer(self, expected: GraphReturnObserver) -> None: + """Detach the graph-return observer ONLY if ``expected`` is installed. + + Compare-and-clear guard for deferred teardown: a seamless non-final + graph phase defers ``teardown_phase`` to its background return-wait + completion, which can fire AFTER the next phase's ``setup_phase`` + already installed ITS observer on this single shared slot. An + unconditional detach here would null the live phase's observer and + silently drop every subsequent graph return (parked dispatch Futures + would strand until the dispatch timeout). Equality, not identity: + bound methods are re-created per attribute access, so the teardown's + ``strategy._on_graph_return`` is a different object than the one + installed by ``setup_phase`` but compares equal. + """ + if self._graph_return_observer == expected: + self._graph_return_observer = None + + def clear_graph_first_token_observer( + self, expected: GraphFirstTokenObserver + ) -> None: + """Detach the first-token observer ONLY if ``expected`` is installed. + + Same compare-and-clear guard as :meth:`clear_graph_return_observer`, + for the graph first-token slot (post-TTFT anchoring): a stale phase's + deferred teardown must never null the next phase's live observer. + """ + if self._graph_first_token_observer == expected: + self._graph_first_token_observer = None def set_branch_orchestrator(self, orchestrator: BranchOrchestrator | None) -> None: """Inject (or detach) the DAG branch orchestrator. @@ -214,6 +290,18 @@ async def on_credit_return( credit_return: Return details including credit and status. """ credit = credit_return.credit + + # Graph-IR returns resolve a parked dispatch Future via a DEDICATED, + # UNCONDITIONAL observer -- BEFORE the phase-handler lookup and the + # gated strategy path. The gated path no-ops once the phase can no + # longer send (or when no handler is registered), which would strand + # the executor coroutine awaiting this node's dispatch. Firing here + # guarantees every graph credit's Future is resolved or rejected. + if credit.trace_id is not None and self._graph_return_observer is not None: + self._graph_return_observer( + credit, credit_return.error, credit_return.cancelled + ) + handler = self._lookup_active_phase_handler(credit, worker_id) if handler is None: return @@ -433,7 +521,10 @@ def _release_slots_for_return( # whether completed or cancelled). DAG children (agent_depth > 0) # inherit the root's session slot via the dispatch path that # bypasses ``acquire_session_slot``; releasing here would underflow. - if credit.is_final_turn and credit.agent_depth == 0: + # Graph-IR credits (``trace_id is not None``) likewise NEVER acquire a + # session slot (``CreditIssuer.issue_graph_credit`` bypasses it), so + # releasing one here would underflow the counter -- skip them. + if credit.is_final_turn and credit.agent_depth == 0 and credit.trace_id is None: concurrency.release_session_slot(phase) # On phase end, release slots for sessions still in flight. @@ -457,29 +548,43 @@ def _release_slots_for_return( async def on_first_token(self, first_token: FirstToken) -> None: """Handle first token event (TTFT) from worker. - Releases prefill concurrency slot, allowing another request - to start prefilling. + Releases the prefill concurrency slot (when limiting is active) and + fires the graph first-token observer (post-TTFT first-token anchoring). + + A ``FirstToken`` now arrives whenever the emitting credit set + ``first_token_event`` -- which includes graph credits with prefill + limiting OFF. The prefill-released counter must therefore only advance + when limiting is actually in force, or the stat inflates; the slot + release itself is a no-op when disabled. Args: - first_token: TTFT event details including credit_id and phase. + first_token: TTFT event details including credit_id, phase, and the + graph-identity fields (trace_id / x_correlation_id / turn_index). """ phase = first_token.phase handler = self._phase_handlers.get(phase) - if not handler: + if handler is None: _logger.debug( lambda: ( f"TTFT for unregistered phase {phase}, " f"credit_id={first_token.credit_id}" ) ) - return - - if handler.handle_first_token is not None: - await handler.handle_first_token(first_token) - - # Track the release - handler.progress.increment_prefill_released() - - # Release the prefill slot - handler.concurrency_manager.release_prefill_slot(phase) + else: + # First-token-aware strategies (e.g. graph replay) observe TTFT here. + if handler.handle_first_token is not None: + await handler.handle_first_token(first_token) + # Only count the release when prefill limiting is active; a + # first-token-anchoring event arrives even without limiting. + if handler.concurrency_manager.prefill_limiting_enabled: + handler.progress.increment_prefill_released() + handler.concurrency_manager.release_prefill_slot(phase) + + # Graph first-token observer fires AFTER the prefill block, independent + # of phase-handler gating, for every FirstToken carrying a trace_id. + if ( + first_token.trace_id is not None + and self._graph_first_token_observer is not None + ): + self._graph_first_token_observer(first_token) diff --git a/src/aiperf/credit/issuer.py b/src/aiperf/credit/issuer.py index 41a14da326..50415c2c91 100644 --- a/src/aiperf/credit/issuer.py +++ b/src/aiperf/credit/issuer.py @@ -14,6 +14,7 @@ from __future__ import annotations +import hashlib import time from typing import TYPE_CHECKING @@ -30,6 +31,7 @@ from aiperf.timing.phase.progress_tracker import PhaseProgressTracker from aiperf.timing.phase.stop_conditions import StopConditionChecker from aiperf.timing.request_cancellation import RequestCancellationSimulator + from aiperf.timing.session_tree import SessionTreeRegistry class CreditIssuer: @@ -62,6 +64,8 @@ def __init__( cancellation_policy: RequestCancellationSimulator, lifecycle: PhaseLifecycle, url_selection_strategy: URLSelectionStrategyProtocol | None = None, + session_tree_registry: SessionTreeRegistry | None = None, + session_tree_registry_enabled: bool | None = None, ) -> None: """Initialize credit issuer. @@ -75,6 +79,14 @@ def __init__( lifecycle: Phase lifecycle for timestamp data. url_selection_strategy: Optional URL selection strategy for multi-URL load balancing. If None, url_index will be None in credits. + session_tree_registry: Optional per-tree finality ledger (DAG + datasets only). When engaged, a root session start opens a tree + and the issuer stamps ``is_parent_final`` / ``is_tree_final`` on + every emitted credit from live tree state. None on non-DAG paths + (finality stays conservative ``(None, False)``). + session_tree_registry_enabled: Explicit engage override. When None, + the registry is engaged iff one was supplied (DAG runs engage it + in every phase). Set True/False to force it on/off (tests). """ self._phase = phase self._stop_checker = stop_checker @@ -84,6 +96,63 @@ def __init__( self._cancellation_policy = cancellation_policy self._lifecycle = lifecycle self._url_selection_strategy = url_selection_strategy + # Per-TEMPLATE URL affinity for sticky graph credits, keyed by the + # nonce-stripped template id and minted DETERMINISTICALLY from it + # (`_stable_graph_url_index`) so the WARMUP priming issuer and the + # PROFILING replay issuer -- separate per-phase objects -- land every + # instance/recycle of one template on the same backend. Entries + # persist for the issuer's lifetime BY DESIGN (a recycle after + # end_graph_trace must reuse the primed backend); the issuer is + # per-phase and the map is bounded by the corpus's template count. + self._graph_url_affinity: dict[str, int] = {} + self._session_tree_registry = ( + session_tree_registry + if ( + session_tree_registry_enabled + if session_tree_registry_enabled is not None + else session_tree_registry is not None + ) + else None + ) + + def _open_session_tree(self, turn: TurnToSend) -> None: + """Open a session tree for a root session-start credit just admitted. + + No-op when tree accounting is not engaged (non-DAG). A root session + start is always depth 0, so the tree root id is the root's own + ``x_correlation_id``. + """ + if self._session_tree_registry is not None: + self._session_tree_registry.open_tree( + turn.effective_root_correlation_id, self._phase, root_pending=True + ) + + def _finality_for_issue(self, turn: TurnToSend) -> tuple[bool | None, bool]: + """Issue-time lineage finality from ``SessionTreeRegistry`` state. + + Conservative by spec: returns ``None``/``False`` whenever indeterminate + (including the non-DAG path where no registry is engaged). + """ + registry = self._session_tree_registry + if registry is None: + return None, False + root_id = turn.effective_root_correlation_id + is_root = turn.parent_correlation_id is None + is_parent_final: bool | None = None + if not is_root and turn.parent_correlation_id == root_id: + # v1: parent finality is determinable only when the parent IS the + # root (the registry tracks per-tree, not per-intermediate-node). + is_parent_final = registry.root_terminal(root_id) + is_tree_final = registry.is_last_tree_request( + root_id, + is_final_turn=turn.is_final_turn, + is_root_credit=is_root, + # Any-mode branch flag, NOT the FORK-only has_forks: a final turn + # declaring SPAWN branches spawns descendants at return-intercept, + # after this stamp, so it must never read as tree-final. + has_branches=turn.has_branches, + ) + return is_parent_final, is_tree_final def can_acquire_and_start_new_session(self) -> bool: """Check if a session slot can be acquired and a new session can be started.""" @@ -134,6 +203,7 @@ async def issue_credit(self, turn: TurnToSend) -> bool: ) if not acquired: return False + self._open_session_tree(turn) # Prefill concurrency: one slot per request, released when TTFT arrives. # Limits concurrent prompt processing which is the GPU-intensive phase. @@ -182,6 +252,7 @@ async def try_issue_credit(self, turn: TurnToSend) -> bool | None: ) if not acquired: return None # No slot - credit not issued + self._open_session_tree(turn) acquired = self._concurrency_manager.try_acquire_prefill_slot( self._phase, can_proceed_fn @@ -209,15 +280,33 @@ async def _issue_credit_internal(self, turn: TurnToSend) -> bool: time.perf_counter_ns() - self._lifecycle.started_at_perf_ns ) - # Get URL index from strategy (for multi-URL load balancing) - # Only advance the round-robin on the first turn of a conversation. - # Subsequent turns will use the url_index stored in the worker's UserSession. - is_first_turn = turn.turn_index == 0 - url_index = ( - self._url_selection_strategy.next_url_index() - if self._url_selection_strategy and is_first_turn - else None - ) + # Get URL index from strategy (for multi-URL load balancing). + # Sticky graph credits get per-template URL affinity minted + # deterministically from the nonce-stripped template identity (NOT the + # per-trajectory ``x_correlation_id``, which embeds a fresh uuid and + # would let warmup priming and profiling replay land on + # different backends); linear turns only advance the round-robin on + # the first turn of a conversation -- subsequent turns use the url_index + # stored in the worker's UserSession. + if turn.trace_id is not None and self._url_selection_strategy: + # Key on the nonce-stripped TEMPLATE id: instance ids carry fresh + # nonces ({template}::{nonce}), and a nonce-bearing key would + # silently re-shuffle URLs between warmup and profiling. + affinity_key = turn.trace_id.split("::", 1)[0] + if affinity_key in self._graph_url_affinity: + url_index = self._graph_url_affinity[affinity_key] + else: + url_index = self._stable_graph_url_index(affinity_key) + self._graph_url_affinity[affinity_key] = url_index + else: + is_first_turn = turn.turn_index == 0 + url_index = ( + self._url_selection_strategy.next_url_index() + if self._url_selection_strategy and is_first_turn + else None + ) + + is_parent_final, is_tree_final = self._finality_for_issue(turn) credit = Credit( id=credit_index, @@ -231,8 +320,15 @@ async def _issue_credit_internal(self, turn: TurnToSend) -> bool: url_index=url_index, agent_depth=turn.agent_depth, parent_correlation_id=turn.parent_correlation_id, + root_correlation_id=turn.root_correlation_id, has_forks=turn.has_forks, + is_parent_final=is_parent_final, + is_tree_final=is_tree_final, branch_mode=turn.branch_mode, + trace_id=turn.trace_id, + node_ordinal=turn.node_ordinal, + phase_variant=turn.phase_variant, + first_token_event=turn.first_token_event, ) await self._credit_router.send_credit(credit=credit) @@ -242,6 +338,115 @@ async def _issue_credit_internal(self, turn: TurnToSend) -> bool: return not is_final_credit + async def issue_graph_credit(self, turn: TurnToSend) -> bool: + """Issue a graph-IR credit, BYPASSING the linear session-slot lifecycle. + + Graph traces are a fan-out DAG: the executor fires nodes in dataflow- + readiness order, not the linear ``turn_index==0`` -> ``num_turns-1`` + sequence the session-slot acquire/release arithmetic assumes. Engaging + ``acquire_session_slot`` (turn0) here would either fail to acquire (the + first-fired node is rarely ordinal 0) or leak the slot (the terminal + node's ordinal rarely equals ``num_turns-1``), deadlocking the phase + once every slot leaks. Trace-admission concurrency is owned by the + ``GraphIRReplayStrategy`` instead. + + A prefill slot is STILL acquired per request (per-request prompt- + processing back-pressure is orthogonal to session accounting). The + ``can_send_any_turn`` stop gate is honored so cancellation / + duration / request-count caps still apply. + + The matching release-side bypass lives in + ``CreditCallbackHandler._release_slots_for_return`` (gated on + ``credit.trace_id is not None``), so a graph credit never touches the + session-slot counter on either side. + + Returns: + True iff the credit was placed on the wire. + """ + # Graph credits use the DAG-child stop gate, NOT ``can_send_any_turn``. + # ``GraphIRReplayStrategy`` owns completion (a trace is DONE when its + # executor drains), so the linear ``SessionCountStopCondition`` + # (``--num-conversations``) must NOT truncate a fan-out DAG mid-trace -- + # it opts out of DAG-child gating for exactly this reason. Cancellation, + # duration, and ``--request-count`` still apply. + if not self._stop_checker.can_send_dag_child_turn(): + return False + acquired = await self._concurrency_manager.acquire_prefill_slot( + self._phase, self._stop_checker.can_send_dag_child_turn + ) + if not acquired: + return False + await self._issue_credit_internal(turn) + return True + + def _stable_graph_url_index(self, affinity_key: str) -> int: + """Deterministically map a graph template onto a URL index. + + The mapping must be a pure function of TEMPLATE identity, NOT of mint + order: warmup priming and profiling replay run in SEPARATE per-phase + issuers, and every instance/recycle of one template must land on the + backend that primed its KV -- so the key (the nonce-stripped template + id) hashes deterministically onto the URL list. Distribution across a corpus's templates is statistically + uniform under the hash; a strategy exposing no ``urls`` falls back to + its round-robin mint. + """ + urls = getattr(self._url_selection_strategy, "urls", None) + if not urls: + return self._url_selection_strategy.next_url_index() + digest = hashlib.sha256(affinity_key.encode()).digest() + return int.from_bytes(digest[:8], "big") % len(urls) + + async def end_graph_trace(self, trace_id: str, phase_variant: str) -> None: + """Close a graph instance's sticky lifecycle (router session only). + + Called ONCE per instance by ``GraphIRReplayStrategy`` at the + adapter-reap points (all in-flight dispatches for the instance + drained, or phase teardown for retained adapters). Forwards to the + router, which closes the instance's sticky session and notifies the + sticky worker (``GraphTraceEnd``). URL affinity is deliberately NOT + evicted: it keys on the nonce-stripped template, and a recycle of the + template must land on the backend that primed its KV. Idempotent end + to end. + """ + await self._credit_router.end_graph_trace(trace_id, phase_variant) + + def mark_graph_sending_complete(self) -> None: + """Signal that a graph phase will issue no further credits. + + ``GraphIRReplayStrategy`` owns completion: a graph credit NEVER trips + ``is_final_credit`` (``CreditCounter.increment_sent`` returns False for + every ``trace_id``-carrying credit), so the issuer never auto-freezes + the sent counts or sets ``all_credits_sent_event`` the way it does for + linear phases (``_issue_credit_internal``). The strategy calls this once + its executors have all drained -- i.e. every node-dispatch that will + ever be issued has been issued -- to freeze the authoritative + ``final_requests_sent`` (the per-node record count) and unblock the + ``PhaseRunner``'s ``_wait_for_sending_complete``. + + Idempotent: setting an already-set ``asyncio.Event`` is a no-op, and + re-freezing simply re-snapshots the (now stable) sent counters. + """ + self._progress.freeze_sent_counts() + self._progress.all_credits_sent_event.set() + + def graph_all_returned(self) -> bool: + """True iff every issued graph credit has returned or cancelled. + + Reads the same ``check_all_returned_or_cancelled`` predicate the + callback handler uses; meaningful only AFTER + ``mark_graph_sending_complete`` has frozen ``final_requests_sent``. + """ + return self._progress.check_all_returned_or_cancelled() + + def set_graph_all_returned_event(self) -> None: + """Set the phase's all-credits-returned event for a graph phase. + + Used by ``GraphIRReplayStrategy`` only in the degenerate + no-credits-issued case (an empty trace set), where no credit return + will ever fire the event via the callback handler. + """ + self._progress.all_credits_returned_event.set() + # ========================================================================= # DAG dispatch helpers (used by BranchOrchestrator) # ========================================================================= @@ -304,7 +509,9 @@ async def dispatch_join_turn(self, pending: PendingBranchJoin) -> bool: num_turns=pending.parent_num_turns, agent_depth=pending.parent_agent_depth, parent_correlation_id=pending.parent_parent_correlation_id, + root_correlation_id=pending.parent_root_correlation_id, has_forks=pending.parent_has_forks_on_gated_turn, + has_branches=pending.parent_has_branches_on_gated_turn, branch_mode=pending.parent_branch_mode, ) return await self._dispatch_dag_turn(turn) diff --git a/src/aiperf/credit/messages.py b/src/aiperf/credit/messages.py index 00a092b5f5..a0a3675b03 100644 --- a/src/aiperf/credit/messages.py +++ b/src/aiperf/credit/messages.py @@ -133,21 +133,33 @@ class CreditReturn( worker_id: str | None = None -class FirstToken(Struct, frozen=True, kw_only=True, tag_field="t", tag="ft"): +class FirstToken( + Struct, omit_defaults=True, frozen=True, kw_only=True, tag_field="t", tag="ft" +): """Worker reports first token received (TTFT event). Sent by worker to router when first valid token is received from inference server. - Router forwards to timing manager to release prefill concurrency slot. + Router forwards to timing manager to release prefill concurrency slot AND to the + graph first-token observer (post-TTFT first-token anchoring). Attributes: credit_id: ID of the credit this TTFT is for. phase: Credit phase for routing to correct phase tracker. ttft_ns: Time to first token in nanoseconds (duration from request start). + trace_id: Graph-IR trace instance this event addresses (e.g. ``t-1#0``); + None for non-graph (template/DAG) dispatch. The graph first-token + observer keys off this to identify the emitting node's trace. + x_correlation_id: Conversation instance ID of the emitting credit; None + when the worker did not stamp it (non-graph fast path). + turn_index: 0-based turn index of the emitting credit; None when unset. """ credit_id: int phase: CreditPhase ttft_ns: int + trace_id: str | None = None + x_correlation_id: str | None = None + turn_index: int | None = None # Union type for decoding worker -> router messages @@ -172,6 +184,32 @@ class CancelCredits(Struct, frozen=True, kw_only=True, tag_field="t", tag="cc"): credit_ids: set[int] -# Union type for decoding router -> worker messages -# Credit is sent directly (no wrapper), CancelCredits for cancellation -RouterToWorkerMessage: TypeAlias = Credit | CancelCredits +class GraphTraceEnd( + Struct, + frozen=True, + kw_only=True, + tag_field="t", + tag="te", # codespell:ignore te +): + """Router notifies the sticky worker that a graph trace execution ended. + + Sent ONCE per instance when the strategy reaps a trace instance's + dispatch adapter (all in-flight dispatches drained) or at phase teardown + for retained adapters. The router closes the instance's sticky session + (graph sessions key on ``trace_id``) before forwarding; the worker uses + it to evict the trace's dynamic-content pool entries. Ids only -- no + content. + + Attributes: + trace_id: Graph-IR trace instance id (``{template}::{nonce}``). + phase_variant: Graph-IR phase variant label (``"profiling"``/``"warmup"``). + """ + + trace_id: str + phase_variant: str + + +# Union type for decoding router -> worker messages: Credit is sent directly +# (no wrapper); CancelCredits for cancellation; GraphTraceEnd for graph trace +# lifecycle. +RouterToWorkerMessage: TypeAlias = Credit | CancelCredits | GraphTraceEnd diff --git a/src/aiperf/credit/sticky_router.py b/src/aiperf/credit/sticky_router.py index b041fe7f4c..28dc0b30e2 100644 --- a/src/aiperf/credit/sticky_router.py +++ b/src/aiperf/credit/sticky_router.py @@ -10,6 +10,12 @@ x_correlation_id (UUID). All turns in a session route to the same worker. conversation_id: Template ID from the dataset (can be reused across sessions). +Graph credits (trace_id set) key their session on the trace INSTANCE id +instead of the per-trajectory corr: every trajectory of one graph trace +instance shares one session -- and therefore one worker -- because +dynamic-slot capture/splice pools and spawn state are worker-local and keyed +by the instance. + Includes: - WorkerLoad: Worker load tracking for fair load balancing - StickyCreditRouter: Main router class @@ -33,6 +39,7 @@ CancelCredits, CreditReturn, FirstToken, + GraphTraceEnd, WorkerReady, WorkerShutdown, WorkerToRouterMessage, @@ -59,8 +66,10 @@ class WorkerLoad: Note on active_sessions: - active_sessions and active_session_ids only represent the number of sticky sessions assigned - to the worker, which inherently means that it only tracks sessions with MORE turns left. This is - because sticky sessions are only created when more than 1 turn exists, and are removed when SENDING the final turn. + to the worker. For linear turns this inherently means only sessions with MORE turns left are + tracked: sticky sessions are created only when more than 1 turn exists, and are removed when + SENDING the final turn. Graph credits (trace_id set) create their session on first + sight and are removed only by end_graph_trace. """ worker_id: str @@ -119,6 +128,15 @@ async def cancel_all_credits(self) -> None: """ ... + async def end_graph_trace(self, trace_id: str, phase_variant: str) -> None: + """Close a graph instance's sticky session and notify its worker. + + Graph session lifecycle is owned by this explicit call, never by + ``is_final_turn`` (graph credits mint ``turn_index`` per node key, so + turn counting cannot express trace completion). Idempotent. + """ + ... + def mark_credits_complete(self) -> None: """Mark that all credits have been issued and returned. @@ -264,9 +282,10 @@ def __init__( Callable[[FirstToken], Awaitable[None]] | None ) = None - # Sticky sessions: x_correlation_id -> worker_id - # Routes all turns of a conversation to the same worker. Required because - # workers cache UserSession state by x_correlation_id. + # Sticky sessions: session key -> worker_id. Linear sessions key on + # x_correlation_id (workers cache UserSession state by it); graph + # sessions key on the trace INSTANCE id so every trajectory of one + # instance shares one worker (worker-local dynamic-slot/spawn state). self._sticky_sessions: dict[str, str] = {} self._cancellation_pending: bool = False @@ -341,8 +360,16 @@ async def send_credit(self, credit: Credit) -> None: if not credit.x_correlation_id: raise RuntimeError("x_correlation_id must be set in Credit") - x_correlation_id = credit.x_correlation_id - sticky_worker_id = self._sticky_sessions.get(x_correlation_id) + # ONE routing key per plane: linear sessions key on the trajectory + # corr (legacy contract); graph credits key on the trace INSTANCE id, + # so EVERY trajectory of one instance shares one session/worker -- + # dynamic-slot capture/splice pools and spawn state are worker-local + # and keyed by the instance, so a child trajectory placed elsewhere + # could never see its producer's content (pool_missing). + session_key = ( + credit.trace_id if credit.trace_id is not None else credit.x_correlation_id + ) + sticky_worker_id = self._sticky_sessions.get(session_key) # Use existing sticky session if worker still valid if sticky_worker_id and sticky_worker_id in self._workers: @@ -381,25 +408,70 @@ async def send_credit(self, credit: Credit) -> None: worker_id = best_worker_id - # Only create sticky session if there are more turns coming. Single-turn - # conversations don't need routing state since there's no next turn. - if not credit.is_final_turn: - self._sticky_sessions[x_correlation_id] = worker_id + # Graph credits (trace_id set) ALWAYS create the session on first + # sight regardless of is_final_turn: a t*-chopped / errored + # trajectory may never dispatch its recorded final turn, so + # turn-based lifecycle would leak; graph sessions are closed + # deterministically by ``end_graph_trace`` at adapter reap. Linear + # turns keep the existing rule: only create when more turns are + # coming. + if credit.trace_id is not None or not credit.is_final_turn: + self._sticky_sessions[session_key] = worker_id load = self._workers[worker_id] load.active_sessions += 1 - load.active_session_ids.add(x_correlation_id) + load.active_session_ids.add(session_key) # Cleanup on final turn - only decrement if session was actually tracked - # (single-turn sessions never get added to _sticky_sessions) - if credit.is_final_turn and self._sticky_sessions.pop(x_correlation_id, None): + # (single-turn sessions never get added to _sticky_sessions). Gated OFF + # for graph credits (trace_id set): a trajectory's recorded final turn + # may never dispatch (t*-chop, errors), so this cleanup would be + # unreliable there. Graph session lifecycle is owned exclusively by + # `end_graph_trace`. + if ( + credit.trace_id is None + and credit.is_final_turn + and self._sticky_sessions.pop(session_key, None) + ): load = self._workers[worker_id] load.active_sessions -= 1 - load.active_session_ids.discard(x_correlation_id) + load.active_session_ids.discard(session_key) self._track_credit_sent(worker_id, credit.id) await self._router_client.send_to(worker_id, credit) + async def end_graph_trace(self, trace_id: str, phase_variant: str) -> None: + """Close a graph instance's sticky session and notify the sticky worker. + + Graph sessions key on the trace INSTANCE id, so ONE call closes the + whole instance (every trajectory) and the worker receives exactly one + ``GraphTraceEnd`` to evict its instance-keyed pool entry. + + State mutation is SYNCHRONOUS before any await (a cancelled caller must + not leave the session half-closed); the worker-forward is best-effort + (the worker-side pool has an LRU backstop for lost notifications). + Idempotent: a missing session (zero-credit trace, repeated call, or a + session already reaped by worker unregistration) is a no-op. + """ + worker_id = self._sticky_sessions.pop(trace_id, None) + if worker_id is None: + return + load = self._workers.get(worker_id) + if load is None: + # Worker unregistered mid-trace; its load entry (and pool) are gone. + return + load.active_sessions -= 1 + load.active_session_ids.discard(trace_id) + try: + await self._router_client.send_to( + worker_id, + GraphTraceEnd(trace_id=trace_id, phase_variant=phase_variant), + ) + except Exception as e: + self.debug( + lambda e=e: f"GraphTraceEnd forward to {worker_id} failed: {e!r}" + ) + async def cancel_all_credits(self) -> None: """Send cancellation requests to all workers with in-flight credits.""" # Mark cancellation first, so we suppress warnings for workers that unregister with in-flight credits. @@ -547,8 +619,8 @@ def _unregister_worker(self, worker_id: str) -> None: self.warning( f"Worker {worker_id} unregistered with {len(orphaned_session_ids)} active sessions, will reassign" ) - for x_correlation_id in orphaned_session_ids: - self._sticky_sessions.pop(x_correlation_id, None) + for session_key in orphaned_session_ids: + self._sticky_sessions.pop(session_key, None) if not worker_load: # Warn but continue - may happen if shutdown message arrives before ready message. diff --git a/src/aiperf/credit/structs.py b/src/aiperf/credit/structs.py index 0a80926ff3..1469ab3558 100644 --- a/src/aiperf/credit/structs.py +++ b/src/aiperf/credit/structs.py @@ -29,7 +29,7 @@ class Credit( Attributes: id: Sequential number of the credit in the credit phase. - phase: Type of credit phase (e.g., "warmup", "profile"). + phase: Type of credit phase (e.g., "warmup", "profiling"). conversation_id: Template ID from the dataset. x_correlation_id: Conversation instance ID for sticky routing (X-Correlation-ID header). turn_index: Index of the turn in the conversation (0-based). @@ -63,16 +63,57 @@ class Credit( url_index: int | None = None agent_depth: int = 0 parent_correlation_id: str | None = None + root_correlation_id: str | None = None + """x_correlation_id of the depth-0 root of this credit's session TREE. + + Stable across the whole tree: the root carries its own x_correlation_id + (left None on the wire when it equals x_correlation_id to keep the struct + small), and every descendant (child, subchild) inherits the root's id. + This is the key used for per-tree finality accounting + (``SessionTreeRegistry``) and is persisted in the export so analysis groups + a tree under one lane. Effective value is + ``root_correlation_id or x_correlation_id``.""" has_forks: bool = False + is_parent_final: bool | None = None + """Parent conversation had already returned its final turn at issue time. + None for roots / when not determinable. Issue-time stamp, never copied.""" + is_tree_final: bool = False + """Provably the last request the whole session tree will send (conservative + False when indeterminate). Issue-time stamp from SessionTreeRegistry.""" branch_mode: ConversationBranchMode = ConversationBranchMode.FORK """DAG branch mode for this credit. Ignored when parent_correlation_id is None (i.e. for root sessions). FORK = inherit parent turn_list; SPAWN = fresh context. Default FORK keeps wire footprint small via msgspec omit_defaults.""" + trace_id: str | None = None + """Graph-IR trace instance identifier this credit addresses + (``{template}::{nonce}``, e.g. ``t-1::3f2a...``); + None for non-graph (template/DAG) dispatch.""" + + node_ordinal: int | None = None + """Ordinal of the graph-IR node this credit addresses; the worker + materializes the node's request from the graph mmap store by this ordinal. + None for non-graph dispatch.""" + + phase_variant: str = "profiling" + """Graph-IR phase variant label for this credit (e.g. ``"profiling"``, + ``"warmup"``). Distinct from the credit-router ``phase`` enum.""" + + first_token_event: bool = False + """When True the worker emits a ``FirstToken`` event on TTFT for this credit + even with prefill-concurrency limiting off (post-TTFT first-token anchoring). + The graph first-token observer keys off the event's ``trace_id``. Default + False keeps the wire footprint small via msgspec ``omit_defaults``.""" + @property def is_final_turn(self) -> bool: return self.turn_index == self.num_turns - 1 + @property + def effective_root_correlation_id(self) -> str: + """Tree root id, defaulting to this credit's own ``x_correlation_id``.""" + return self.root_correlation_id or self.x_correlation_id + class CreditContext( Struct, omit_defaults=True, kw_only=True, tag_field="t", tag="cctx" @@ -132,13 +173,52 @@ class TurnToSend(Struct, frozen=True): num_turns: int agent_depth: int = 0 parent_correlation_id: str | None = None + root_correlation_id: str | None = None + """x_correlation_id of the depth-0 root of this turn's session TREE. + + None for a root turn (the root IS its own tree root); set on every + descendant to the root's id. Propagated onto the issued ``Credit`` and used + for per-tree finality accounting. Effective value is + ``root_correlation_id or x_correlation_id``.""" has_forks: bool = False + has_branches: bool = False + """True iff the originating turn declares ANY branch (FORK or SPAWN) in its + metadata ``branch_ids``. Superset of ``has_forks``, which is FORK-only and + owned by the sticky router's deferred-eviction logic — do not conflate the + two. Consumed by finality stamping: a turn that will spawn descendants on + its return can never be the tree's provably-last request, even when the + registry shows nothing outstanding yet (SPAWN children register only at + return-intercept, AFTER issue-time stamping).""" branch_mode: ConversationBranchMode = ConversationBranchMode.FORK + trace_id: str | None = None + """Graph-IR trace instance identifier this turn addresses + (``{template}::{nonce}``, e.g. ``t-1::3f2a...``); None for non-graph + (template/DAG) dispatch. Copied into the issued Credit.""" + + node_ordinal: int | None = None + """Ordinal of the graph-IR node this turn addresses; the worker materializes + the node's request from the graph mmap store by this ordinal. None for + non-graph dispatch. Copied into the issued Credit.""" + + phase_variant: str = "profiling" + """Graph-IR phase variant label for this turn (e.g. ``"profiling"``, + ``"warmup"``). Copied into the issued Credit.""" + + first_token_event: bool = False + """When True the issued Credit requests a per-credit ``FirstToken`` event on + TTFT regardless of prefill-concurrency limiting (post-TTFT first-token + anchoring). Copied into the issued Credit.""" + @property def is_final_turn(self) -> bool: return self.turn_index == self.num_turns - 1 + @property + def effective_root_correlation_id(self) -> str: + """Tree root id, defaulting to this turn's own ``x_correlation_id``.""" + return self.root_correlation_id or self.x_correlation_id + @classmethod def from_previous_credit( cls, credit: Credit, next_meta: "TurnMetadata | None" = None @@ -149,7 +229,10 @@ def from_previous_credit( credit: The previous turn's credit. next_meta: Metadata for the NEW turn being built. When provided, the ``has_forks`` flag is derived from it so the sticky - router can defer parent-entry eviction until DAG children drain. + router can defer parent-entry eviction until DAG children drain, + and ``has_branches`` (any-mode) is derived from its + ``branch_ids`` so finality stamping stays conservative on + spawning turns. """ return cls( conversation_id=credit.conversation_id, @@ -158,6 +241,12 @@ def from_previous_credit( num_turns=credit.num_turns, agent_depth=credit.agent_depth, parent_correlation_id=credit.parent_correlation_id, + root_correlation_id=credit.root_correlation_id, has_forks=next_meta.has_forks if next_meta is not None else False, + has_branches=bool(next_meta.branch_ids) if next_meta is not None else False, branch_mode=credit.branch_mode, + trace_id=credit.trace_id, + node_ordinal=credit.node_ordinal, + phase_variant=credit.phase_variant, + first_token_event=credit.first_token_event, ) diff --git a/src/aiperf/dataset/_mp_context.py b/src/aiperf/dataset/_mp_context.py new file mode 100644 index 0000000000..9c4626780c --- /dev/null +++ b/src/aiperf/dataset/_mp_context.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Dedicated multiprocessing context for trace-loader worker pools. + +Trace-loader pools (the Weka graph adapter's parallel directory/row parsers +in :mod:`aiperf.dataset.graph.adapters.weka.trace_parallel`) fork +worker processes after the parent has loaded HF tokenizers and exercised +their Rust thread pool. Under the default ``fork`` start method, that +inherits broken rayon state and ``transformers`` whose offline-mode flag was +cached at parent-import time -- the combination DEADLOCKS the workers (live +symptom: workers sit idle at ~2% CPU each holding 64 inherited rayon threads +and many GB RSS while the parse hangs indefinitely). + +Forking from a long-lived ``forkserver`` helper instead bypasses parent state +entirely: the helper is a fresh Python interpreter that imports only the +modules in ``_LOADER_PRELOAD``. The helper additionally instantiates the +benchmark's tokenizer (driven by env vars set in +:func:`get_loader_mp_context`) so every worker fork CoW-shares the in-memory +copy instead of re-loading from disk. + +This context is intentionally *separate* from any future service-spawning +context -- its sole consumer is trace-loader worker pools, so its preload +list and lifecycle are scoped to that use case. +""" + +from __future__ import annotations + +import contextlib +import multiprocessing +import os + +from aiperf.common.constants import IS_LINUX + +_LOADER_PRELOAD = [ + # Module imports happen once in the forkserver helper so workers don't + # pay the transformers/HF import cost on every spawn. Order matters + # only loosely -- the tokenizer-preload module is last so its instance + # creation finds Tokenizer already imported. + "aiperf.common.tokenizer", + "aiperf.common.hash_id_random_generator", + "numpy", + # Side-effecting: instantiates the tokenizer named by + # AIPERF_LOADER_PRELOAD_TOKENIZER into the helper's heap. + "aiperf.dataset._tokenizer_preload", +] + +_ENV_PRELOAD_NAME = "AIPERF_LOADER_PRELOAD_TOKENIZER" +_ENV_PRELOAD_TRUST = "AIPERF_LOADER_PRELOAD_TRUST_REMOTE_CODE" +_ENV_PRELOAD_REVISION = "AIPERF_LOADER_PRELOAD_REVISION" + +_loader_ctx: multiprocessing.context.BaseContext | None = None + + +def configure_loader_tokenizer_env( + *, + trust_remote_code: bool, + revision: str | None, +) -> None: + """Publish the run's tokenizer ``trust_remote_code``/``revision`` for the preload. + + The parallel loader threads only the tokenizer NAME into + :func:`get_loader_mp_context`; trust/revision travel via env because the + forkserver helper snapshots the env once at spawn and every worker (preload + hit or on-demand fallback) reads the SAME inherited triple. Must run in the + parent before the first loader pool is opened -- the DatasetManager calls + it at graph-configure time, before any parse. + """ + os.environ[_ENV_PRELOAD_TRUST] = "true" if trust_remote_code else "false" + os.environ[_ENV_PRELOAD_REVISION] = revision or "main" + + +def get_loader_mp_context( + *, + preload_tokenizer: str | None = None, + trust_remote_code: bool | None = None, + revision: str | None = None, +) -> multiprocessing.context.BaseContext: + """Return the trace-loader-specific multiprocessing context. + + On Linux this is a ``forkserver`` context whose helper is started eagerly + with stdio redirected to ``/dev/null`` and (optionally) the named + tokenizer pre-instantiated in its heap so workers CoW-share it. On + macOS this is a ``spawn`` context (no helper; each worker is a fresh + interpreter, and ``preload_tokenizer`` is a no-op). + + ``trust_remote_code``/``revision`` default to ``None`` meaning "keep the + values :func:`configure_loader_tokenizer_env` published" (falling back to + ``False``/``"main"`` when nothing was published); explicit arguments + override. Clobbering the published env here would silently preload with + the wrong tokenizer revision for runs pinned to a non-``main`` revision. + + The context is built once and cached; later calls with a different + ``preload_tokenizer`` reuse the original helper. Callers are expected + to share a single tokenizer per process lifetime (the typical AIPerf + flow). Workers receiving a different name fall back to on-demand load. + """ + global _loader_ctx + if _loader_ctx is not None: + return _loader_ctx + + # Env must be set BEFORE the forkserver helper is spawned: it reads + # these at module-import time and instantiates the tokenizer once in + # its own heap, where every forked worker CoW-shares it. + if preload_tokenizer: + os.environ[_ENV_PRELOAD_NAME] = preload_tokenizer + if trust_remote_code is not None: + os.environ[_ENV_PRELOAD_TRUST] = "true" if trust_remote_code else "false" + else: + os.environ.setdefault(_ENV_PRELOAD_TRUST, "false") + if revision is not None: + os.environ[_ENV_PRELOAD_REVISION] = revision + else: + os.environ.setdefault(_ENV_PRELOAD_REVISION, "main") + + method = "forkserver" if IS_LINUX else "spawn" + ctx = multiprocessing.get_context(method) + if method == "forkserver": + ctx.set_forkserver_preload(_LOADER_PRELOAD) + _eagerly_start_forkserver() + _loader_ctx = ctx + return _loader_ctx + + +def _eagerly_start_forkserver() -> None: + """Boot the forkserver helper with stdio pointing at ``/dev/null``. + + Must run before any fork through the context so the helper inherits + ``/dev/null`` rather than the parent's possibly-captured stdio (pytest, + Textual dashboard, etc.). If the helper is already running, we're too + late to redirect -- bail out. + """ + from multiprocessing import forkserver as _fs + + if getattr(_fs, "_forkserver", None) and getattr( + _fs._forkserver, "_forkserver_pid", None + ): + return + + devnull_fd = os.open(os.devnull, os.O_RDWR) + saved = [os.dup(fd) for fd in (0, 1, 2)] + try: + for fd in (0, 1, 2): + os.dup2(devnull_fd, fd) + with contextlib.suppress(Exception): + _fs.ensure_running() + finally: + for fd, original in zip((0, 1, 2), saved, strict=False): + os.dup2(original, fd) + os.close(original) + os.close(devnull_fd) diff --git a/src/aiperf/dataset/_tokenizer_preload.py b/src/aiperf/dataset/_tokenizer_preload.py new file mode 100644 index 0000000000..f24a839247 --- /dev/null +++ b/src/aiperf/dataset/_tokenizer_preload.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Forkserver preload module: load the trace-loader tokenizer once in the helper. + +Listed in :data:`aiperf.dataset._mp_context._LOADER_PRELOAD` so Python's +forkserver helper imports it at startup. Any tokenizer instantiated here +lives in the helper's anonymous heap; every worker child forked from it +CoW-shares those pages instead of re-loading the tokenizer from disk. + +For a 700 MiB Qwen tokenizer with 16 workers, this takes the per-spawn +cost from ~700 ms x 16 (sequential disk reads under file-lock contention) +down to a single shared resident copy. + +Configuration is via environment variables, populated by +:func:`aiperf.dataset._mp_context.get_loader_mp_context` **before** it +calls :func:`multiprocessing.forkserver.ensure_running`. The env is +inherited into the helper when Python spawns it, and into every worker +forked from the helper: + + AIPERF_LOADER_PRELOAD_TOKENIZER tokenizer name to preload + AIPERF_LOADER_PRELOAD_TRUST_REMOTE_CODE "true" or "false" (default false) + AIPERF_LOADER_PRELOAD_REVISION HF revision (default "main") + +Fail-soft: any failure is logged to stderr and silently skipped. The +worker's :func:`Tokenizer.from_pretrained` fallback covers misses, so a +preload failure never blocks the run -- it just means workers re-load +from disk individually. + +Fork-safety: we deliberately **do not** call ``tokenizer.encode`` or +``tokenizer.decode`` here. HF fast tokenizers spawn rayon threads at +first parallel encode; a forkserver that has triggered parallel state +would propagate stale thread references into every forked child. Loading +the tokenizer object alone does not trigger parallel execution. We also +set ``TOKENIZERS_PARALLELISM=false`` so HF does not emit its post-fork +"disabling parallelism to avoid deadlocks" warning in every worker. +""" + +from __future__ import annotations + +import os +import sys +from typing import Any + +_LOADED: dict[tuple[str, bool, str], Any] = {} + +_ENV_NAME = "AIPERF_LOADER_PRELOAD_TOKENIZER" +_ENV_TRUST = "AIPERF_LOADER_PRELOAD_TRUST_REMOTE_CODE" +_ENV_REVISION = "AIPERF_LOADER_PRELOAD_REVISION" + + +def _env_name() -> str: + return os.environ.get(_ENV_NAME, "").strip() + + +def _env_trust_remote_code() -> bool: + return os.environ.get(_ENV_TRUST, "false").strip().lower() in ("1", "true", "yes") + + +def _env_revision() -> str: + return os.environ.get(_ENV_REVISION, "main").strip() or "main" + + +def _preload() -> None: + name = _env_name() + if not name: + return + + os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") + + try: + from aiperf.common.tokenizer import Tokenizer + except ImportError as e: + print( + f"[aiperf.loader_tokenizer_preload] tokenizer module unavailable; " + f"skipping preload: {e!r}", + file=sys.stderr, + flush=True, + ) + return + + trust = _env_trust_remote_code() + revision = _env_revision() + + try: + # resolve_alias=True matches the worker/parent fallback in + # aiperf.dataset.graph.adapters.shared.content._build_generator, so a + # preload hit and a fallback miss load the SAME tokenizer for an + # aliased name. + tok = Tokenizer.from_pretrained( + name, + trust_remote_code=trust, + revision=revision, + resolve_alias=True, + ) + _LOADED[(name, trust, revision)] = tok + print( + f"[aiperf.loader_tokenizer_preload] preloaded '{name}' " + f"(trust_remote_code={trust}, revision={revision}) into forkserver heap", + file=sys.stderr, + flush=True, + ) + except Exception as e: # preload must never crash the forkserver helper + print( + f"[aiperf.loader_tokenizer_preload] failed to preload '{name}': {e!r}; " + "workers will load on demand", + file=sys.stderr, + flush=True, + ) + + +def get_preloaded( + name: str, + *, + trust_remote_code: bool = False, + revision: str = "main", +) -> Any | None: + """Return the preloaded tokenizer for ``(name, trust_remote_code, revision)``. + + Returns ``None`` when nothing was preloaded with this exact triple, so the + caller can fall through to :meth:`Tokenizer.from_pretrained`. + """ + return _LOADED.get((name, trust_remote_code, revision)) + + +def clear_preloaded() -> None: + """Drop every preloaded tokenizer. + + Primarily for test isolation: a test that triggers ``_preload`` (or + otherwise populates ``_LOADED``) would leak its tokenizer into every later + ``get_preloaded`` caller in the same process. + """ + _LOADED.clear() + + +_preload() diff --git a/src/aiperf/dataset/composer/base.py b/src/aiperf/dataset/composer/base.py index a7b08f0640..3b6aa81545 100644 --- a/src/aiperf/dataset/composer/base.py +++ b/src/aiperf/dataset/composer/base.py @@ -330,7 +330,9 @@ def _inject_context_prompts(self, conversations: list[Conversation]) -> None: if shared_system_prompt: conversation.system_message = shared_system_prompt self.trace( - lambda conv=conversation: f"Set system_message on conversation {conv.session_id}" + lambda conv=conversation: ( + f"Set system_message on conversation {conv.session_id}" + ) ) # Set user context prompt (unique per session) @@ -340,9 +342,10 @@ def _inject_context_prompts(self, conversations: list[Conversation]) -> None: ) conversation.user_context_message = user_context self.trace( - lambda idx=session_index, - conv=conversation: f"Set user_context_message for session {idx} " - f"(conversation {conv.session_id})" + lambda idx=session_index, conv=conversation: ( + f"Set user_context_message for session {idx} " + f"(conversation {conv.session_id})" + ) ) @staticmethod diff --git a/src/aiperf/dataset/composer/synthetic_rankings.py b/src/aiperf/dataset/composer/synthetic_rankings.py index 459ce54dce..0cf5a02ec4 100644 --- a/src/aiperf/dataset/composer/synthetic_rankings.py +++ b/src/aiperf/dataset/composer/synthetic_rankings.py @@ -109,6 +109,8 @@ def _create_turn(self, num_passages: int) -> Turn: self._finalize_turn(turn) self.debug( - lambda: f"[rankings] query_len={len(query_text)} chars, passages={num_passages}" + lambda: ( + f"[rankings] query_len={len(query_text)} chars, passages={num_passages}" + ) ) return turn diff --git a/src/aiperf/dataset/dataset_manager.py b/src/aiperf/dataset/dataset_manager.py index 38688ae792..f656d55cce 100644 --- a/src/aiperf/dataset/dataset_manager.py +++ b/src/aiperf/dataset/dataset_manager.py @@ -6,6 +6,7 @@ import gc import time from io import BytesIO +from pathlib import Path from typing import TYPE_CHECKING from urllib.parse import urlparse @@ -43,6 +44,10 @@ RequestInfo, SessionPayloads, ) +from aiperf.common.models.dataset_models import ( + GraphDatasetMetadata, + GraphSegmentClientMetadata, +) from aiperf.common.tokenizer import Tokenizer from aiperf.config.artifacts import OutputDefaults from aiperf.config.dataset import FileDataset, PublicDataset @@ -60,6 +65,7 @@ if TYPE_CHECKING: from aiperf.config.resolution.plan import BenchmarkRun + from aiperf.dataset.graph.store_build import GraphStoreBuildResult from aiperf.dataset.protocols import ( DatasetBackingStoreProtocol, DatasetClientStoreProtocol, @@ -149,7 +155,11 @@ async def _profile_configure_command( self.info(lambda: f"Configuring dataset for {self.service_id}") begin = time.perf_counter() await self._configure_dataset() - await self._generate_inputs_json_file() + # Skip inputs.json for graph runs: the worker materializes real + # per-node payloads from the unified store, so there is no + # conversation payload to dump -- entries would be empty placeholders. + if self._graph_workload_path() is None: + await self._generate_inputs_json_file() await self._configure_dataset_client_and_free_memory() duration = time.perf_counter() - begin @@ -159,7 +169,10 @@ async def _configure_dataset_client_and_free_memory(self) -> None: """Configure the dataset client for serving fallback requests, then free memory.""" conversation_count = len(self.dataset) - if not self._compress_only: + # Graph runs never initialize the conversation backing store (they serve + # no conversation requests -- workers read the unified segment store), so + # building the fallback client from it would fail. + if not self._compress_only and self._graph_workload_path() is None: client_metadata = self._backing_store.get_client_metadata() ClientStoreClass = plugins.get_class( PluginType.DATASET_CLIENT_STORE, client_metadata.client_type @@ -293,8 +306,10 @@ def _generate_input_payloads( ) endpoint: EndpointProtocol = EndpointClass(model_endpoint=model_endpoint) self.debug( - lambda: f"Created endpoint protocol for {model_endpoint.endpoint.type}, " - f"class: {endpoint.__class__.__name__}", + lambda: ( + f"Created endpoint protocol for {model_endpoint.endpoint.type}, " + f"class: {endpoint.__class__.__name__}" + ), ) session_payloads_map: dict[str, list] = {} for conversation in self.dataset.values(): @@ -406,6 +421,44 @@ def _load_synthetic_dataset(self) -> list[Conversation]: self._default_context_mode = composer.get_default_context_mode() return conversations + def _graph_workload_path(self) -> Path | None: + """Return the input path iff this run is a graph IR workload, else None.""" + from aiperf.dataset.graph.workload_detect import resolve_graph_workload + + ref = resolve_graph_workload(self.run) + return ref.path if ref is not None else None + + async def _build_graph_store(self, graph_path: Path) -> GraphStoreBuildResult: + """Gate the endpoint type, then run the ONE graph store build. + + The shared seam for :meth:`_configure_graph_dataset` (production, which + broadcasts the result's store/sidecar locations) and + :meth:`_configure_graph_workload` (facet-only callers): validate the + run's endpoint up front + (:func:`~aiperf.dataset.graph.workload_detect.validate_graph_endpoint_type` + -- the graph dispatch path emits a chat-completions body verbatim, so a + non-chat endpoint would 422 on every request), then hand the build to + :class:`~aiperf.dataset.graph.store_build.GraphStoreBuilder`. + """ + from aiperf.dataset.graph.store_build import GraphStoreBuilder + from aiperf.dataset.graph.workload_detect import validate_graph_endpoint_type + + validate_graph_endpoint_type(self.run) + return await GraphStoreBuilder(self.run).build(graph_path) + + async def _configure_graph_workload(self, graph_path: Path) -> GraphDatasetMetadata: + """Build the graph store for a graph workload and return the graph facet. + + The build itself lives in + :class:`~aiperf.dataset.graph.store_build.GraphStoreBuilder` (the ONE + streaming store build for every graph workload -- weka / dynamo / + native / dag_jsonl); this wrapper keeps the facet-returning contract + (:class:`GraphDatasetMetadata`: the trace universe plus the per-node + theoretical prefix-cache map) for direct callers and runs the same + endpoint gate as production via :meth:`_build_graph_store`. + """ + return (await self._build_graph_store(graph_path)).facet + async def _load_accuracy_dataset(self) -> list[Conversation]: from aiperf.dataset.loader.accuracy_dataset_loader import AccuracyDatasetLoader @@ -433,27 +486,42 @@ async def _load_accuracy_dataset(self) -> list[Conversation]: loader = AccuracyDatasetLoader(run=self.run) return await loader.load() - async def _configure_dataset(self) -> None: - self.dataset_configured.clear() - self._default_context_mode = None + async def _select_conversations(self) -> list[Conversation]: + """Load the run's conversations via the appropriate dataset path. + Precedence: accuracy > public > file-backed custom composer > synthetic. + + Graph IR workloads never reach this method: they are configured natively + by :meth:`_configure_graph_dataset` (no conversations), which + :meth:`_configure_dataset` routes to before any conversation work. + """ accuracy_cfg = self.run.cfg.accuracy - accuracy_enabled = accuracy_cfg.enabled if accuracy_cfg else False - default_dataset = self.run.cfg.get_default_dataset() + if accuracy_cfg and accuracy_cfg.enabled: + return await self._load_accuracy_dataset() - if accuracy_enabled: - conversations = await self._load_accuracy_dataset() - elif isinstance(default_dataset, PublicDataset): - conversations = await self._load_public_dataset() - elif isinstance(default_dataset, FileDataset) or ( + default_dataset = self.run.cfg.get_default_dataset() + if isinstance(default_dataset, PublicDataset): + return await self._load_public_dataset() + if isinstance(default_dataset, FileDataset) or ( getattr(default_dataset, "path", None) is not None ): # Use CUSTOM composer if the dataset is file-backed (FileDataset # or any composed/file-source variant exposing a `path`). The # composer auto-infers the format. - conversations = self._load_custom_dataset() - else: - conversations = self._load_synthetic_dataset() + return self._load_custom_dataset() + return self._load_synthetic_dataset() + + async def _configure_dataset(self) -> None: + self.dataset_configured.clear() + self._default_context_mode = None + + graph_path = self._graph_workload_path() + if graph_path is not None: + await self._configure_graph_dataset(graph_path) + return + + default_dataset = self.run.cfg.get_default_dataset() + conversations = await self._select_conversations() self.dataset = {conv.session_id: conv for conv in conversations} self._conversation_ids_cache = [ @@ -510,6 +578,44 @@ async def _configure_dataset(self) -> None: ) ) + async def _configure_graph_dataset(self, graph_path: Path) -> None: + """Graph-native configure: unified store + sidecar + graph broadcast. + + No stub conversations, no conversation mmap store, no inputs.json: + graph payloads live in the unified segment store keyed by + ``(trace_id, node_ordinal)`` and the schedule plane plans from the + graph_meta sidecar, so the conversation-shaped plumbing has nothing + to carry for a graph run. The broadcast's graph-typed client metadata + is the single source of truth for the store and sidecar locations -- + both come off the build's :class:`GraphStoreBuildResult` (the builder + hard-fails a build that never landed the mandatory sidecar). + """ + result = await self._build_graph_store(graph_path) + graph_meta = result.facet + self.dataset = {} + self._conversation_ids_cache = [] + default_dataset = self.run.cfg.get_default_dataset() + self.dataset_metadata = DatasetMetadata( + conversations=[], + sampling_strategy=getattr(default_dataset, "sampling", None), + graph=graph_meta, + ) + self.info( + f"graph dataset configured: {len(graph_meta.trace_ids)} traces, " + f"sidecar {result.sidecar_path}" + ) + await self.publish( + DatasetConfiguredNotification( + service_id=self.service_id, + metadata=self.dataset_metadata, + client_metadata=GraphSegmentClientMetadata( + store_base_path=result.base_path, + benchmark_id=self.run.benchmark_id, + sidecar_path=result.sidecar_path, + ), + ) + ) + @on_request(MessageType.CONVERSATION_REQUEST) async def _handle_conversation_request( self, message: ConversationRequestMessage diff --git a/src/aiperf/dataset/generator/_coding_cicd_docs.py b/src/aiperf/dataset/generator/_coding_cicd_docs.py new file mode 100644 index 0000000000..332cacb5c4 --- /dev/null +++ b/src/aiperf/dataset/generator/_coding_cicd_docs.py @@ -0,0 +1,439 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CI/CD output, config-file, markdown-doc, and test-output generators. + +Extracted from ``coding_content.py`` to keep that module under the +ergonomics file-size cap. Methods read ``self._template_rng`` and the +shared vocabulary tuples; behavior is unchanged. +""" + +from __future__ import annotations + +from aiperf.dataset.generator._coding_vocab import ( + _CLASSES, + _ERROR_MESSAGES, + _METHODS, + _MODULES, + _VARS, +) + + +class _CicdDocsMixin: + def _gen_cicd_output(self, language: str | None = None) -> str: + r = self._template_rng + mod = r.choice(_MODULES) + n_pass = r.randint(20, 200) + n_fail = r.randint(0, 5) + n_skip = r.randint(0, 10) + n_pkgs = r.randint(50, 300) + install_time = r.uniform(0.5, 10) + n_lint_files = r.randint(10, 100) + n_type_mods = r.randint(100, 500) + coverage = r.uniform(70, 99) + ver = f"{r.randint(1, 9)}.{r.randint(0, 99)}.{r.randint(0, 99)}" + artifact_size = r.uniform(0.1, 50) + status = "PASSED" if n_fail == 0 else "FAILED" + elapsed = r.randint(30, 600) + + lang_toolchain = { + "python": { + "install": f"pip install -r requirements.txt\n Resolved {n_pkgs} packages in {install_time:.1f}s", + "lint": f"ruff check . && ruff format --check .\n All checks passed ({n_lint_files} files)", + "typecheck": f"mypy src/\n Success: {n_type_mods} modules checked", + "test": f"pytest tests/ -v\n {n_pass} passed, {n_fail} failed, {n_skip} skipped\n Coverage: {coverage:.1f}%", + "build": f"python -m build\n Built {mod}-{ver}.tar.gz ({artifact_size:.1f} MB)", + }, + "go": { + "install": f"go mod download\n Resolved {n_pkgs} packages in {install_time:.1f}s", + "lint": f"golangci-lint run ./...\n All checks passed ({n_lint_files} files)", + "typecheck": f"go vet ./...\n Success: {n_type_mods} packages checked", + "test": f"go test -v -race -coverprofile=coverage.out ./...\n {n_pass} passed, {n_fail} failed, {n_skip} skipped\n Coverage: {coverage:.1f}%", + "build": f"go build -o bin/{mod} ./cmd/{mod}\n Built bin/{mod} ({artifact_size:.1f} MB)", + }, + "rust": { + "install": f"cargo fetch\n Resolved {n_pkgs} crates in {install_time:.1f}s", + "lint": f"cargo clippy -- -D warnings\n All checks passed ({n_lint_files} files)", + "typecheck": f"cargo check\n Checked {n_type_mods} crates", + "test": f"cargo test\n {n_pass} passed, {n_fail} failed, {n_skip} ignored\n Coverage: {coverage:.1f}%", + "build": f"cargo build --release\n Built target/release/{mod} ({artifact_size:.1f} MB)", + }, + "typescript": { + "install": f"npm ci\n Resolved {n_pkgs} packages in {install_time:.1f}s", + "lint": f"eslint src/ && prettier --check src/\n All checks passed ({n_lint_files} files)", + "typecheck": f"tsc --noEmit\n Success: {n_type_mods} modules checked", + "test": f"vitest run\n {n_pass} passed, {n_fail} failed, {n_skip} skipped\n Coverage: {coverage:.1f}%", + "build": f"npm run build\n Built dist/{mod}-{ver}.tgz ({artifact_size:.1f} MB)", + }, + } + toolchain = lang_toolchain.get( + language, r.choice(list(lang_toolchain.values())) + ) + + return f"""\ +=== CI Pipeline: {mod} === +Step 1/5: Installing dependencies... + {toolchain["install"]} +Step 2/5: Linting... + {toolchain["lint"]} +Step 3/5: Type checking... + {toolchain["typecheck"]} +Step 4/5: Running tests... + {toolchain["test"]} +Step 5/5: Building artifacts... + {toolchain["build"]} +Pipeline {status} in {elapsed}s +""" + + def _gen_config_file(self, language: str | None = None) -> str: + r = self._template_rng + mod = r.choice(_MODULES) + v1, v2, v3 = r.sample(_VARS, 3) + + lang_to_kinds: dict[str, list[str]] = { + "python": ["yaml", "toml", "dockerfile"], + "go": ["yaml", "makefile"], + "rust": ["toml"], + "typescript": ["yaml", "dockerfile"], + } + choices = ( + lang_to_kinds.get(language, ["yaml", "toml", "dockerfile", "makefile"]) + if language + else ["yaml", "toml", "dockerfile", "makefile"] + ) + kind = r.choice(choices) + if kind == "yaml": + return self._gen_config_yaml(mod, v1, v2, v3) + if kind == "toml": + return self._gen_config_toml(mod, v1, v2, v3) + if kind == "dockerfile": + return self._gen_config_dockerfile(mod, v1, v2, language) + return self._gen_config_makefile(mod) + + def _gen_config_yaml(self, mod: str, v1: str, v2: str, v3: str) -> str: + r = self._template_rng + port = r.randint(3000, 9999) + workers = r.randint(1, 16) + v2_val = r.randint(1, 1000) + v3_val = r.choice(_MODULES) + db_port = r.choice([5432, 3306, 27017, 6379]) + pool = r.randint(5, 50) + return f"""\ +# {mod} configuration +service: + name: {mod} + port: {port} + workers: {workers} + {v1}: + enabled: true + {v2}: {v2_val} + {v3}: "{v3_val}" + logging: + level: info + format: json + database: + host: localhost + port: {db_port} + pool_size: {pool} +""" + + def _gen_config_toml(self, mod: str, v1: str, v2: str, v3: str) -> str: + r = self._template_rng + ver = f"{r.randint(0, 9)}.{r.randint(0, 99)}.{r.randint(0, 99)}" + desc_cls = r.choice(_CLASSES) + desc_method = r.choice(_METHODS) + dep1, dep2 = r.choice(_MODULES), r.choice(_MODULES) + dep1_ver = f"{r.randint(1, 5)}.{r.randint(0, 20)}" + dep2_ver = f"{r.randint(0, 3)}.{r.randint(0, 40)}" + tool_mod = r.choice(_MODULES) + v1_val = r.randint(1, 100) + v2_val = r.choice(_MODULES) + return f"""\ +[project] +name = "{mod}" +version = "{ver}" +description = "{desc_cls} {desc_method} service" + +[dependencies] +{dep1} = "{dep1_ver}" +{dep2} = "{dep2_ver}" + +[tool.{tool_mod}] +{v1} = {v1_val} +{v2} = "{v2_val}" +{v3} = true +""" + + def _gen_config_dockerfile( + self, mod: str, v1: str, v2: str, language: str | None + ) -> str: + r = self._template_rng + env1_val = r.randint(1, 100) + env2_val = r.choice(_MODULES) + port = r.randint(3000, 9999) + docker_lang = language or "python" + if docker_lang == "python": + py_ver = r.randint(10, 13) + base_image = f"python:3.{py_ver}-slim" + install_cmd = "COPY requirements.txt .\nRUN pip install --no-cache-dir -r requirements.txt" + run_cmd = f'CMD ["python", "-m", "{mod}"]' + elif docker_lang == "go": + go_ver = f"1.{r.randint(21, 23)}" + base_image = f"golang:{go_ver}-alpine" + install_cmd = "COPY go.mod go.sum ./\nRUN go mod download" + run_cmd = f'CMD ["./bin/{mod}"]' + elif docker_lang == "rust": + base_image = "rust:1-slim" + install_cmd = "COPY Cargo.toml Cargo.lock ./\nRUN cargo fetch" + run_cmd = f'CMD ["./target/release/{mod}"]' + else: + node_ver = r.randint(18, 22) + base_image = f"node:{node_ver}-alpine" + install_cmd = "COPY package.json package-lock.json ./\nRUN npm ci" + run_cmd = f'CMD ["node", "dist/{mod}/index.js"]' + return f"""\ +FROM {base_image} + +WORKDIR /app + +{install_cmd} + +COPY src/ ./src/ + +ENV {v1.upper()}={env1_val} +ENV {v2.upper()}={env2_val} + +EXPOSE {port} + +{run_cmd} +""" + + def _gen_config_makefile(self, mod: str) -> str: + return f"""\ +.PHONY: build test lint clean + +build: +\t@echo "Building {mod}..." +\tgo build -o bin/{mod} ./cmd/{mod} + +test: +\t@echo "Testing {mod}..." +\tgo test -v -race ./... + +lint: +\tgolangci-lint run ./... + +clean: +\trm -rf bin/ dist/ *.egg-info +""" + + def _gen_markdown_doc(self, language: str | None = None) -> str: + r = self._template_rng + cls = r.choice(_CLASSES) + m1, m2 = r.sample(_METHODS, 2) + mod = r.choice(_MODULES) + v1 = r.choice(_VARS) + err = r.choice(_ERROR_MESSAGES) + + lang_examples = self._build_markdown_lang_examples(cls, m1, mod, v1) + example = lang_examples.get(language, r.choice(list(lang_examples.values()))) + + v2, v3 = r.sample(_VARS, 2) + err2 = r.choice(_ERROR_MESSAGES) + env_var = f"AIPERF_{mod.upper()}_{v1.upper()}" + + return f"""\ +# {cls} + +## Overview + +The `{cls}` class provides {m1} and {m2} operations for the `{mod}` module. + +## Usage + +```{example["fence"]} +{example["code"]} +``` + +## Configuration + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `{v1}` | {example["param_type"]} | required | Primary {v1} identifier | +| `{v2}` | {example["param_type"]} | `None` | Optional {v2} override | +| `{v3}` | int | `10` | Maximum {v3} per batch | +| `timeout` | float | `30.0` | Operation timeout in seconds | + +Environment variable override: `{env_var}` + +## API Reference + +### `{m1}({v1})` + +Performs the {m1} operation. + +**Parameters:** +- `{v1}` ({example["param_type"]}): The input {v1}. + +**Returns:** {example["return_type"]} + +### `{m2}()` + +Performs the {m2} operation. + +**Raises:** `ValueError` if {err}. + +## Errors + +| Error | Condition | Recovery | +|-------|-----------|----------| +| `ValueError` | {err} | Check {v1} parameter | +| `RuntimeError` | {err2} | Retry with backoff | +| `TimeoutError` | Operation exceeds timeout | Increase timeout or reduce {v3} | +""" + + def _build_markdown_lang_examples( + self, cls: str, m1: str, mod: str, v1: str + ) -> dict[str, dict[str, str]]: + r = self._template_rng + return { + "python": { + "fence": "python", + "code": f'from {mod} import {cls}\n\ninstance = {cls}({v1}="value")\nresult = await instance.{m1}()', + "param_type": r.choice( + ("str", "int", "float", "bool", "dict", "list", "Any", "Optional") + ), + "return_type": r.choice( + ("str", "int", "bool", "dict", "list", "None", "Any") + ), + }, + "go": { + "fence": "go", + "code": f'import "{mod}"\n\nc := {mod}.New{cls}("{v1}")\nerr := c.{m1.title()}(ctx)', + "param_type": r.choice( + ( + "string", + "int", + "int64", + "bool", + "[]byte", + "error", + "context.Context", + ) + ), + "return_type": r.choice(("string", "int", "bool", "error", f"*{cls}")), + }, + "rust": { + "fence": "rust", + "code": f'use {mod}::{cls};\n\nlet mut c = {cls}::new("{v1}");\nc.{m1}().await?;', + "param_type": r.choice( + ( + "&str", + "String", + "i64", + "bool", + "Vec", + "&[u8]", + "Option", + ) + ), + "return_type": r.choice( + ("Result<()>", "Result", "bool", "Option", "&str") + ), + }, + "typescript": { + "fence": "typescript", + "code": f"import {{ {cls} }} from './{mod}';\n\nconst c = new {cls}({{ {v1}: 'value' }});\nawait c.{m1}();", + "param_type": r.choice( + ( + "string", + "number", + "boolean", + "Record", + "unknown[]", + ) + ), + "return_type": r.choice( + ("string", "number", "boolean", "void", "Promise") + ), + }, + } + + def _gen_test_output(self, language: str | None = None) -> str: + r = self._template_rng + mod = r.choice(_MODULES) + cls = r.choice(_CLASSES) + methods = r.sample(list(_METHODS), 5) + + lang_to_kind = { + "python": "pytest", + "go": "go", + "rust": "cargo", + "typescript": "jest", + } + kind = ( + lang_to_kind[language] + if language in lang_to_kind + else r.choice(["pytest", "go", "cargo"]) + ) + if kind == "pytest": + lines = [ + "============================= test session starts =============================" + ] + lines.append(f"collected {r.randint(10, 100)} items\n") + for m in methods: + status = r.choice(["PASSED", "PASSED", "PASSED", "FAILED"]) + lines.append(f"tests/test_{mod}.py::Test{cls}::test_{m} {status}") + n_pass = sum(1 for line in lines if "PASSED" in line) + n_fail = len(methods) - n_pass + dur = r.uniform(0.5, 30.0) + lines.append(f"\n{'=' * 70}") + lines.append(f"{n_pass} passed, {n_fail} failed in {dur:.2f}s") + return "\n".join(lines) + "\n" + elif kind == "jest": + runner = r.choice(["JEST", "VITEST"]) + lines = [ + f" {runner} v{r.randint(28, 30)}.{r.randint(0, 9)}.{r.randint(0, 9)}" + ] + lines.append("") + results: list[str] = [] + for m in methods: + passed = r.choice([True, True, True, False]) + mark = "\u2713" if passed else "\u2717" + dur_ms = r.randint(1, 500) + results.append(f" {mark} {cls} > {m} ({dur_ms} ms)") + lines.append(results[-1]) + n_pass = sum(1 for res in results if "\u2713" in res) + n_fail = len(methods) - n_pass + dur = r.uniform(0.5, 15.0) + lines.append("") + lines.append( + f"Tests: {n_fail} failed, {n_pass} passed, {len(methods)} total" + ) + lines.append(f"Time: {dur:.3f} s") + lines.append(f"Ran all test suites matching /src/{mod}.test.ts/i.") + return "\n".join(lines) + "\n" + elif kind == "go": + lines = [] + for m in methods: + status = r.choice(["ok", "ok", "ok", "FAIL"]) + dur = r.uniform(0.001, 2.0) + lines.append(f"--- {status}: Test{m.title()} ({dur:.3f}s)") + lines.append( + f"{status} \t{mod}/{r.choice(_MODULES)}\t{r.uniform(0.1, 5.0):.3f}s" + ) + return "\n".join(lines) + "\n" + else: + lines = [f" Compiling {mod} v0.{r.randint(1, 99)}.{r.randint(0, 9)}"] + lines.append(f" Finished test target(s) in {r.uniform(1, 30):.2f}s") + lines.append(" Running unittests src/lib.rs\n") + for m in methods: + status = r.choice(["ok", "ok", "ok", "FAILED"]) + lines.append(f"test {mod}::{cls.lower()}::test_{m} ... {status}") + n_pass = sum(1 for line in lines if "... ok" in line) + n_fail = len(methods) - n_pass + lines.append( + f"\ntest result: {'ok' if n_fail == 0 else 'FAILED'}. " + f"{n_pass} passed; {n_fail} failed; 0 ignored" + ) + return "\n".join(lines) + "\n" diff --git a/src/aiperf/dataset/generator/_coding_conversations.py b/src/aiperf/dataset/generator/_coding_conversations.py new file mode 100644 index 0000000000..c33427067b --- /dev/null +++ b/src/aiperf/dataset/generator/_coding_conversations.py @@ -0,0 +1,253 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Conversation-template generators (mixin for CodingContentGenerator). + +Extracted from ``coding_content.py`` to keep that module under the +ergonomics file-size cap. Methods read ``self._template_rng`` and the +shared vocabulary tuples; behavior is unchanged. +""" + +from __future__ import annotations + +from aiperf.dataset.generator._coding_text import ( + _BRIDGE_ANALYZE, + _BRIDGE_EXPLAIN, + _BRIDGE_FIX, + _BRIDGE_PERF, + _BRIDGE_REFACTOR, + _BRIDGE_SUMMARY, + _BRIDGE_TEST, + _BRIDGE_WRITE_TEST, + _FOLLOWUP_QUESTIONS, + _LANGUAGES, +) + + +class _ConversationsMixin: + def _gen_conv_bugfix(self) -> str: + r = self._template_rng + lang = r.choice(_LANGUAGES) + ids = self._conv_ids() + + turns = [ + f"[User]\n{self._conv_user_msg(ids)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_read_long(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_FIX, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_TEST, ids)}\n\n" + f"{self._gen_tool_bash(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_SUMMARY, ids)}", + ] + return "\n\n".join(turns) + + def _gen_conv_review(self) -> str: + r = self._template_rng + lang = r.choice(_LANGUAGES) + ids = self._conv_ids() + + turns = [ + f"[User]\n{self._conv_user_msg(ids)}\n\n" + f"{self._gen_git_diff(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_read_long(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_FIX, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[User]\n{self._conv_bridge(_FOLLOWUP_QUESTIONS, ids)}", + ] + return "\n\n".join(turns) + + def _gen_conv_feature(self) -> str: + r = self._template_rng + lang = r.choice(_LANGUAGES) + ids = self._conv_ids() + + turns = [ + f"[User]\n{self._conv_user_msg(ids)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_search_verbose(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_read_long(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_FIX, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_WRITE_TEST, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_TEST, ids)}\n\n" + f"{self._gen_tool_bash_verbose(language=lang)}", + ] + return "\n\n".join(turns) + + def _gen_conv_debug(self) -> str: + r = self._template_rng + lang = r.choice(_LANGUAGES) + ids = self._conv_ids() + error_block = r.choice( + [ + lambda: self._gen_error_traceback(language=lang), + self._gen_cuda_error, + ] + )() + + turns = [ + f"[User]\n{self._conv_user_msg(ids)}\n\n{error_block}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_read_long(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_search_verbose(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_FIX, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_SUMMARY, ids)}", + ] + return "\n\n".join(turns) + + def _gen_conv_qa(self) -> str: + r = self._template_rng + lang = r.choice(_LANGUAGES) + ids = self._conv_ids() + + turns = [ + f"[User]\n{self._conv_user_msg(ids)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_read_long(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_EXPLAIN, ids)}", + f"[User]\n{self._conv_bridge(_FOLLOWUP_QUESTIONS, ids)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_FIX, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + ] + return "\n\n".join(turns) + + def _gen_conv_refactor(self) -> str: + """Multi-file refactoring: search callers, read multiple files, edit each.""" + r = self._template_rng + lang = r.choice(_LANGUAGES) + ids = self._conv_ids() + + turns = [ + f"[User]\n{self._conv_user_msg(ids)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_search_verbose(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_read_long(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_read(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_REFACTOR, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\nNow let me update the callers.\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_REFACTOR, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_TEST, ids)}\n\n" + f"{self._gen_tool_bash_verbose(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_SUMMARY, ids)}", + ] + return "\n\n".join(turns) + + def _gen_conv_perf(self) -> str: + """Performance investigation: profile, read hot path, optimize, benchmark.""" + r = self._template_rng + lang = r.choice(_LANGUAGES) + ids = self._conv_ids() + + turns = [ + f"[User]\n{self._conv_user_msg(ids)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_PERF, ids)}\n\n" + f"{self._gen_tool_bash(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_read_long(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_PERF, ids)}\n\n" + f"{self._gen_tool_search(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_FIX, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_TEST, ids)}\n\n" + f"{self._gen_tool_bash_verbose(language=lang)}", + f"[User]\n{self._conv_bridge(_FOLLOWUP_QUESTIONS, ids)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_EXPLAIN, ids)}\n\n" + f"{self._conv_bridge(_BRIDGE_SUMMARY, ids)}", + ] + return "\n\n".join(turns) + + def _gen_conv_cicd(self) -> str: + """CI/CD debugging: failing pipeline, read logs, fix config, re-run.""" + r = self._template_rng + lang = r.choice(_LANGUAGES) + ids = self._conv_ids() + + ci_output = self._gen_cicd_output(language=lang) + + turns = [ + f"[User]\nThe CI pipeline is failing on the {ids['module']} service. " + f"Can you take a look?\n\n{ci_output}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_read(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_read(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_FIX, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_TEST, ids)}\n\n" + f"{self._gen_tool_bash_verbose(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_SUMMARY, ids)}", + f"[User]\n{self._conv_bridge(_FOLLOWUP_QUESTIONS, ids)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_EXPLAIN, ids)}", + ] + return "\n\n".join(turns) + + def _gen_conv_ml_debug(self) -> str: + """ML/GPU debugging: CUDA error, read training code, fix, re-run.""" + ids = self._conv_ids() + + cuda_err = self._gen_cuda_error() + training_code = self._gen_ml_training_code() + training_log = self._gen_ml_training_log() + inference_code = self._gen_ml_inference_code() + + turns = [ + f"[User]\nI'm getting a CUDA error during training. " + f"Here's the error:\n\n{cuda_err}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"read\n" + f'train.py\n' + f"\n{training_code}\n", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"read\n" + f'inference.py\n' + f"\n{inference_code}\n", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_FIX, ids)}\n\n" + f"{self._gen_tool_edit(language='python')}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_TEST, ids)}\n\n" + f"bash\n" + f'python train.py --max-steps 10\n' + f"\n{training_log}\n", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_SUMMARY, ids)}", + "[User]\nCan you also check if the inference script has the same issue?", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._conv_bridge(_BRIDGE_EXPLAIN, ids)}\n\n" + f"{self._conv_bridge(_BRIDGE_SUMMARY, ids)}", + ] + return "\n\n".join(turns) + + def _gen_conv_test_write(self) -> str: + """Test writing session: read code, write tests, iterate on failures.""" + r = self._template_rng + lang = r.choice(_LANGUAGES) + ids = self._conv_ids() + + turns = [ + f"[User]\nWrite comprehensive tests for {ids['cls']}.{ids['method']}(). " + f"Cover the happy path, edge cases, and error handling.", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_read_long(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_search(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_WRITE_TEST, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_TEST, ids)}\n\n" + f"{self._gen_tool_bash_verbose(language=lang)}", + f"[User]\n{self._conv_bridge(_FOLLOWUP_QUESTIONS, ids)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_WRITE_TEST, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_TEST, ids)}\n\n" + f"{self._gen_tool_bash(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_SUMMARY, ids)}", + ] + return "\n\n".join(turns) diff --git a/src/aiperf/dataset/generator/_coding_conversations_advanced.py b/src/aiperf/dataset/generator/_coding_conversations_advanced.py new file mode 100644 index 0000000000..d36e68e3ed --- /dev/null +++ b/src/aiperf/dataset/generator/_coding_conversations_advanced.py @@ -0,0 +1,337 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Higher-level conversation generators: migration, deploy, security, distributed, etc. + +Extracted from ``coding_content.py`` to keep that module under the +ergonomics file-size cap. Methods read ``self._template_rng`` and the +shared vocabulary tuples; behavior is unchanged. +""" + +from __future__ import annotations + +from aiperf.dataset.generator._coding_text import ( + _BRIDGE_ANALYZE, + _BRIDGE_ARCHITECTURE_TRADEOFF, + _BRIDGE_DATA_ARCHITECTURE, + _BRIDGE_DEPLOY, + _BRIDGE_DISTRIBUTED, + _BRIDGE_EXPLAIN, + _BRIDGE_FIX, + _BRIDGE_OBSERVABILITY, + _BRIDGE_PERF, + _BRIDGE_REFACTOR, + _BRIDGE_SECURITY, + _BRIDGE_SUMMARY, + _BRIDGE_TEST, + _BRIDGE_WRITE_TEST, + _FOLLOWUP_QUESTIONS, + _LANGUAGES, +) +from aiperf.dataset.generator._coding_vocab import ( + _DB_TABLES, +) + + +class _ConversationsAdvancedMixin: + def _gen_conv_migration(self) -> str: + """Multi-file migration: search all usages, update each file, run tests.""" + r = self._template_rng + lang = r.choice(_LANGUAGES) + ids = self._conv_ids() + + turns = [ + f"[User]\nMigrate {ids['cls']}.{ids['method']}() from " + f"sync to async. It's called across multiple files in {ids['module']}. " + f"Update all callers and add backward compat.", + f"[Assistant]\nLet me find all the callers first.\n\n" + f"{self._gen_tool_search_verbose(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_read_long(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_read(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_read(language=lang)}", + f"[Assistant]\nI'll start with the core change to {ids['cls']}, " + f"then update each caller.\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\nNow updating the first caller.\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\nUpdating the second caller.\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\nUpdating the third caller and adding the " + f"backward-compat wrapper.\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_TEST, ids)}\n\n" + f"{self._gen_tool_bash_verbose(language=lang)}", + f"[User]\n{self._conv_bridge(_FOLLOWUP_QUESTIONS, ids)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_EXPLAIN, ids)}\n\n" + f"{self._conv_bridge(_BRIDGE_SUMMARY, ids)}", + ] + return "\n\n".join(turns) + + def _gen_conv_deploy(self) -> str: + """Deployment troubleshooting: check config, logs, fix, verify.""" + r = self._template_rng + lang = r.choice(_LANGUAGES) + ids = self._conv_ids() + + config_block = self._gen_config_file(language=lang) + json_resp = self._gen_json_response(language=lang) + + turns = [ + f"[User]\nThe {ids['module']} service keeps crashing after deploy. " + f"The health check is failing and pods are in CrashLoopBackOff.", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_DEPLOY, ids)}\n\n" + f"bash\n" + f'kubectl describe pod {ids["module"]}-' + f"{r.randint(1000, 9999)}-{r.choice('abcdef')}" + f"{r.choice('abcdef')}{r.choice('0123456789')}" + f"{r.choice('abcdef')}{r.choice('0123456789')}\n" + f"\n" + f"Name: {ids['module']}-deployment-{r.randint(1000, 9999)}\n" + f"Namespace: default\n" + f"Status: Running\n" + f"Containers:\n" + f" {ids['module']}:\n" + f" Image: registry.internal/{ids['module']}:latest\n" + f" State: Waiting (CrashLoopBackOff)\n" + f" Last State: Terminated (Error, exit code 1)\n" + f" Ready: False\n" + f" Restart Count: 7\n" + f" Limits:\n" + f" cpu: 2\n" + f" memory: 512Mi\n" + f" Requests:\n" + f" cpu: 500m\n" + f" memory: 256Mi\n" + f" Liveness: http-get http://:8080/health delay=10s timeout=3s period=5s\n" + f" Readiness: http-get http://:8080/ready delay=5s timeout=3s period=5s\n" + f"Events:\n" + f" Warning BackOff 2m (x7 over 10m) kubelet " + f"Back-off restarting failed container\n" + f"", + f"[Assistant]\nThe memory limit looks too low. Let me check the config.\n\n" + f"read\n" + f'kubernetes/deployment.yaml\n' + f"\n{config_block}\n", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_DEPLOY, ids)}\n\n" + f"bash\n" + f'kubectl logs deploy/{ids["module"]} ' + f"--tail=30\n" + f"\n" + f"{self._gen_error_traceback(language=lang)}\n" + f"", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_FIX, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\nLet me also increase the memory limits.\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_TEST, ids)}\n\n" + f"bash\n" + f'kubectl apply -f kubernetes/deployment.yaml ' + f"&& kubectl rollout status deploy/{ids['module']} --timeout=120s\n" + f"\n" + f"deployment.apps/{ids['module']} configured\n" + f'Waiting for deployment "{ids["module"]}" rollout to finish: ' + f"1 old replicas are pending termination...\n" + f'deployment "{ids["module"]}" successfully rolled out\n' + f"", + f"[Assistant]\nLet me verify the health check is passing now.\n\n" + f"bash\n" + f'curl -s http://localhost:8080/health ' + f"| python -m json.tool\n" + f"\n{json_resp}\n", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_SUMMARY, ids)}", + ] + return "\n\n".join(turns) + + def _gen_conv_security(self) -> str: + """Security vulnerability investigation: find vuln, analyze attack vectors, fix, test.""" + r = self._template_rng + lang = r.choice(_LANGUAGES) + ids = self._conv_ids() + + turns = [ + f"[User]\nI think there's a security vulnerability in the {ids['module']} " + f"service. The {ids['method']}() endpoint accepts user input for {ids['var']} " + f"without proper validation.", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_SECURITY, ids)}\n\n" + f"{self._gen_tool_read_long(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_search_verbose(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ARCHITECTURE_TRADEOFF, ids)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_SECURITY, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_WRITE_TEST, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_TEST, ids)}\n\n" + f"{self._gen_tool_bash_verbose(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_SUMMARY, ids)}", + ] + return "\n\n".join(turns) + + def _gen_conv_distributed(self) -> str: + """Distributed systems debugging: inconsistency, analyze replication, fix consensus.""" + r = self._template_rng + lang = r.choice(_LANGUAGES) + ids = self._conv_ids() + + config_block = self._gen_config_file(language=lang) + + turns = [ + f"[User]\nThere are inconsistent reads across replicas in the " + f"{ids['module']} service. After writing to {ids['var']} via " + f"{ids['cls']}.{ids['method']}(), some replicas return stale data.", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_DISTRIBUTED, ids)}\n\n" + f"read\n" + f'config/replication.yaml\n' + f"\n{config_block}\n", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_search_verbose(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ARCHITECTURE_TRADEOFF, ids)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_DISTRIBUTED, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_TEST, ids)}\n\n" + f"{self._gen_tool_bash_verbose(language=lang)}", + f"[User]\n{self._conv_bridge(_FOLLOWUP_QUESTIONS, ids)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_DISTRIBUTED, ids)}\n\n" + f"{self._conv_bridge(_BRIDGE_SUMMARY, ids)}", + ] + return "\n\n".join(turns) + + def _gen_conv_observability(self) -> str: + """Observability gap: add tracing, metrics, structured logging.""" + r = self._template_rng + lang = r.choice(_LANGUAGES) + ids = self._conv_ids() + + config_block = self._gen_config_file(language=lang) + + turns = [ + f"[User]\nCan't debug a production latency spike in {ids['module']}. " + f"There's no tracing or metrics on {ids['cls']}.{ids['method']}().", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_read_long(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_OBSERVABILITY, ids)}\n\n" + f"{self._gen_tool_search_verbose(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_OBSERVABILITY, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_OBSERVABILITY, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\nLet me also add the telemetry configuration.\n\n" + f"read\n" + f'config/telemetry.yaml\n' + f"\n{config_block}\n", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_TEST, ids)}\n\n" + f"bash\n" + f'curl -s http://localhost:8080/metrics ' + f"| head -20\n" + f"\n{self._gen_json_response(language=lang)}\n", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_SUMMARY, ids)}", + ] + return "\n\n".join(turns) + + def _gen_conv_db_optimize(self) -> str: + """Database optimization: EXPLAIN, read ORM code, add index, benchmark.""" + r = self._template_rng + lang = r.choice(_LANGUAGES) + ids = self._conv_ids() + + table = r.choice(_DB_TABLES) + sql_block = self._gen_sql_query() + + turns = [ + f"[User]\nThe {ids['method']}() query on the {table} table is taking " + f"over 5 seconds in production. Can you optimize it?", + f"[Assistant]\nLet me run EXPLAIN ANALYZE to see the query plan.\n\n" + f"bash\n" + f'psql -d mydb -c "EXPLAIN ANALYZE ' + f"SELECT * FROM {table} WHERE {ids['var']} = 'test'\"\n" + f"\n{sql_block}\n", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_DATA_ARCHITECTURE, ids)}\n\n" + f"{self._gen_tool_read_long(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ARCHITECTURE_TRADEOFF, ids)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_DATA_ARCHITECTURE, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_TEST, ids)}\n\n" + f"{self._gen_tool_bash_verbose(language=lang)}", + f"[User]\nShould we also partition the {table} table?", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ARCHITECTURE_TRADEOFF, ids)}\n\n" + f"{self._conv_bridge(_BRIDGE_SUMMARY, ids)}", + ] + return "\n\n".join(turns) + + def _gen_conv_architecture_review(self) -> str: + """Architecture review: read multiple files, deep multi-paragraph analysis, refactor.""" + r = self._template_rng + lang = r.choice(_LANGUAGES) + ids = self._conv_ids() + + turns = [ + f"[User]\nCan you do an architecture review of the {ids['module']} " + f"service? I'm concerned about coupling and scalability.", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_read_long(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_read_long(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ANALYZE, ids)}\n\n" + f"{self._gen_tool_search_verbose(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ARCHITECTURE_TRADEOFF, ids)}\n\n" + f"{self._conv_bridge(_BRIDGE_ARCHITECTURE_TRADEOFF, ids)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_REFACTOR, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[User]\nWhat about the scalability of {ids['cls']}? Will this " + f"approach hold up under 10x traffic?", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ARCHITECTURE_TRADEOFF, ids)}\n\n" + f"{self._conv_bridge(_BRIDGE_PERF, ids)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_FIX, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_TEST, ids)}\n\n" + f"{self._gen_tool_bash_verbose(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_SUMMARY, ids)}", + ] + return "\n\n".join(turns) + + def _gen_conv_incident_response(self) -> str: + """Production incident: cascading failure, diagnose, fix, add monitoring, post-mortem.""" + r = self._template_rng + lang = r.choice(_LANGUAGES) + ids = self._conv_ids() + + config_block = self._gen_config_file(language=lang) + error_block = self._gen_error_traceback(language=lang) + + turns = [ + f"[User]\nProduction incident: the {ids['module']} service is down " + f"and it's causing cascading failures in downstream services.", + "[Assistant]\nLet me check the service health immediately.\n\n" + "bash\n" + 'curl -s http://localhost:8080/health ' + "|| echo 'Connection refused'\n" + "\nConnection refused\n", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_DEPLOY, ids)}\n\n" + f"read\n" + f'kubernetes/deployment.yaml\n' + f"\n{config_block}\n", + f"[Assistant]\nLet me check the logs for the root cause.\n\n" + f"bash\n" + f'kubectl logs deploy/{ids["module"]} ' + f"--tail=50\n" + f"\n{error_block}\n", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_ARCHITECTURE_TRADEOFF, ids)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_FIX, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\nNow let me add a circuit breaker to prevent cascading failures.\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_OBSERVABILITY, ids)}\n\n" + f"{self._gen_tool_edit(language=lang)}", + f"[Assistant]\n{self._conv_bridge(_BRIDGE_TEST, ids)}\n\n" + f"{self._gen_tool_bash_verbose(language=lang)}", + f"[Assistant]\nPost-mortem summary: The {ids['module']} service experienced " + f"a cascading failure triggered by {ids['error']}. The root cause was " + f"{ids['cls']}.{ids['method']}() not handling the error gracefully, which " + f"caused the health check to fail and pods to restart in a loop. " + f"Fixes applied: error handling in {ids['method']}(), circuit breaker " + f"pattern for downstream calls, and Prometheus alerts for early detection.", + ] + return "\n\n".join(turns) diff --git a/src/aiperf/dataset/generator/_coding_errors_diff.py b/src/aiperf/dataset/generator/_coding_errors_diff.py new file mode 100644 index 0000000000..de7a55c4af --- /dev/null +++ b/src/aiperf/dataset/generator/_coding_errors_diff.py @@ -0,0 +1,382 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Error-traceback and git-diff generators (mixin for CodingContentGenerator). + +Extracted from ``coding_content.py`` to keep that module under the +ergonomics file-size cap. Methods read ``self._template_rng`` and the +shared vocabulary tuples; behavior is unchanged. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from aiperf.dataset.generator._coding_vocab import ( + _CLASSES, + _ERROR_MESSAGES, + _METHODS, + _MODULES, + _VARS, +) + + +@dataclass(frozen=True, slots=True) +class _DiffCtx: + """Pre-sampled randomness shared across per-language diff hunk builders.""" + + cls: str + m1: str + m2: str + m3: str + v1: str + v2: str + v3: str + ln: int + ln2: int + err: str + mod: str + hunk_old: int + hunk_new: int + + +def _diff_hunks_python(c: _DiffCtx) -> tuple[str, str, str]: + return ( + f"""\ +@@ -{c.ln},8 +{c.ln},14 @@ class {c.cls}: + def {c.m1}(self): +- {c.v1} = self._{c.m2}() +- return {c.v1} ++ try: ++ {c.v1} = await self._{c.m2}() ++ if {c.v1} is None: ++ raise ValueError("{c.err}") ++ return {c.v1} ++ except Exception as e: ++ logger.error(f"{c.cls}.{c.m1} failed: {{{{e}}}}") ++ raise""", + f"""\ +@@ -{c.ln2},5 +{c.ln2},9 @@ def {c.m2}({c.v1}): + {c.v2} = {c.mod}.{c.m3}({c.v1}) +- return {c.v2} ++ if not {c.v2}: ++ raise RuntimeError("{c.err}") ++ logger.info("{c.m2} completed: %s", {c.v2}) ++ return {{{{"{c.v1}": {c.v2}, "status": "ok"}}}}""", + f"""\ +@@ -{c.hunk_old},3 +{c.hunk_new},7 @@ ++import logging ++from {c.mod} import {c.cls} ++ ++logger = logging.getLogger(__name__)""", + ) + + +def _diff_hunks_go(c: _DiffCtx) -> tuple[str, str, str]: + return ( + f"""\ +@@ -{c.ln},6 +{c.ln},12 @@ func (s *{c.cls}) {c.m1.title()}() error {{{{ +- return nil ++ {c.v1}, err := s.{c.m2.title()}(ctx) ++ if err != nil {{{{ ++ return fmt.Errorf("{c.err}: %w", err) ++ }}}} ++ s.{c.v2} = {c.v1} ++ return nil""", + f"""\ +@@ -{c.ln2},4 +{c.ln2},8 @@ func (s *{c.cls}) {c.m2.title()}() (string, error) {{{{ + s.mu.RLock() + defer s.mu.RUnlock() +- return s.{c.v1}, nil ++ if s.{c.v1} == "" {{{{ ++ return "", fmt.Errorf("{c.err}") ++ }}}} ++ return fmt.Sprintf("%s:%d", s.{c.v1}, s.{c.v2}), nil""", + f"""\ +@@ -{c.hunk_old},3 +{c.hunk_new},7 @@ ++import ( ++ "fmt" ++ "log/slog" ++)""", + ) + + +def _diff_hunks_rust(c: _DiffCtx) -> tuple[str, str, str]: + return ( + f"""\ +@@ -{c.ln},5 +{c.ln},11 @@ impl {c.cls} {{{{ + pub fn {c.m1}(&self) -> Result<()> {{{{ +- Ok(()) ++ let {c.v1} = self.{c.m2}()?; ++ if {c.v1}.is_empty() {{{{ ++ anyhow::bail!("{c.err}"); ++ }}}} ++ tracing::info!("{c.m1} completed: {{}}", {c.v1}); ++ Ok(())""", + f"""\ +@@ -{c.ln2},4 +{c.ln2},7 @@ impl {c.cls} {{{{ + fn {c.m2}(&self) -> Result {{{{ +- Ok(self.{c.v1}.clone()) ++ let {c.v2} = &self.{c.v1}; ++ anyhow::ensure!(!{c.v2}.is_empty(), "{c.err}"); ++ Ok({c.v2}.clone())""", + f"""\ +@@ -{c.hunk_old},3 +{c.hunk_new},6 @@ ++use anyhow::Result; ++use tracing; ++use {c.mod}::{c.cls};""", + ) + + +def _diff_hunks_typescript(c: _DiffCtx) -> tuple[str, str, str]: + return ( + f"""\ +@@ -{c.ln},6 +{c.ln},12 @@ export class {c.cls} {{{{ + {c.m1}({c.v1}: string) {{{{ +- return this.{c.m2}({c.v1}); ++ try {{{{ ++ const {c.v2} = await this.{c.m2}({c.v1}); ++ if (!{c.v2}) throw new Error('{c.err}'); ++ return {{ status: 'ok', data: {c.v2} }}; ++ }}}} catch (err) {{{{ ++ console.error(`{c.cls}.{c.m1} failed: ${{{{err}}}}`); ++ throw err; ++ }}}}""", + f"""\ +@@ -{c.ln2},4 +{c.ln2},7 @@ export class {c.cls} {{{{ + private {c.m2}({c.v1}: string): {c.v2} {{{{ +- return this.#{c.v1}; ++ if (!this.#{c.v1}) {{{{ ++ throw new Error('{c.err}'); ++ }}}} ++ return this.#{c.v1};""", + f"""\ +@@ -{c.hunk_old},3 +{c.hunk_new},6 @@ ++import {{ {c.cls} }} from './{c.mod}'; ++import type {{ {c.v3.title()} }} from './types'; ++""", + ) + + +class _ErrorsDiffMixin: + def _gen_error_traceback(self, language: str | None = None) -> str: + r = self._template_rng + err = r.choice(_ERROR_MESSAGES) + cls = r.choice(_CLASSES) + m1, m2, m3, m4 = r.sample(_METHODS, 4) + + lang_to_kind = { + "python": "python", + "go": "go", + "rust": "rust", + "typescript": "node", + } + kind = ( + lang_to_kind[language] + if language in lang_to_kind + else r.choice(["python", "go", "rust", "node"]) + ) + file_pool = self._file_pool(language) + f1, f2, f3, f4 = r.sample(list(file_pool), 4) + ms = (m1, m2, m3, m4) + fs = (f1, f2, f3, f4) + if kind == "python": + return self._error_traceback_python(r, err=err, cls=cls, ms=ms, fs=fs) + elif kind == "go": + return self._error_traceback_go(r, err=err, cls=cls, ms=ms, fs=fs) + elif kind == "rust": + return self._error_traceback_rust(r, err=err, cls=cls, ms=ms, fs=fs) + else: + return self._error_traceback_node(r, err=err, cls=cls, ms=ms, fs=fs) + + def _error_traceback_python( + self, + r, + *, + err: str, + cls: str, + ms: tuple[str, str, str, str], + fs: tuple[str, str, str, str], + ) -> str: + m1, m2, m3, m4 = ms + f1, f2, f3, f4 = fs + v = r.choice(_VARS) + mod = r.choice(_MODULES) + err2 = r.choice(_ERROR_MESSAGES) + cls2 = r.choice(_CLASSES) + return f"""\ +Traceback (most recent call last): + File "{f1}", line {r.randint(10, 500)}, in {m1} + result = self.{m2}(data) + File "{f2}", line {r.randint(10, 300)}, in {m2} + {v} = await self._{m3}() + File "{f3}", line {r.randint(10, 200)}, in _{m3} + return {mod}.{m4}({v}) + File "{f4}", line {r.randint(1, 200)}, in {m4} + raise ValueError("{err}") +ValueError: {err} + +During handling of the above exception, another exception occurred: + +Traceback (most recent call last): + File "{f1}", line {r.randint(10, 500)}, in {m1} + self._{v} = {mod}.{m1}() + File "{f2}", line {r.randint(10, 300)}, in __init__ + raise RuntimeError("{err2}") +RuntimeError: {cls}.{m1}() failed: {err2} + +The above exception was the direct cause of the following exception: + +{cls2}Error: {cls}.{m1}() aborted after {err}: {err2} +""" + + def _error_traceback_go( + self, + r, + *, + err: str, + cls: str, + ms: tuple[str, str, str, str], + fs: tuple[str, str, str, str], + ) -> str: + m1, m2, m3, m4 = ms + f1, f2, f3, f4 = fs + g1 = r.randint(1, 100) + g2 = r.randint(101, 200) + cls2 = r.choice(_CLASSES) + return f"""\ +goroutine {g1} [running]: +runtime/debug.Stack() + /usr/local/go/src/runtime/debug/stack.go:{r.randint(10, 50)} +main.{cls}.{m1.title()}(...) + {f1}:{r.randint(10, 300)} +main.{cls}.{m2.title()}(0xc000{r.randint(10000, 99999):05x}) + {f2}:{r.randint(10, 300)} +main.{cls}.{m3.title()}(0xc000{r.randint(10000, 99999):05x}, 0x{r.randint(100, 999):x}) + {f3}:{r.randint(10, 300)} +panic: {err} + +goroutine {g2} [select]: +main.{cls2}.{m4.title()}(0xc000{r.randint(10000, 99999):05x}) + {f4}:{r.randint(10, 300)} +0x{r.randint(100, 999):x} +created by main.New{cls2} + {f4}:{r.randint(10, 100)} +""" + + def _error_traceback_rust( + self, + r, + *, + err: str, + cls: str, + ms: tuple[str, str, str, str], + fs: tuple[str, str, str, str], + ) -> str: + m1, m2, m3, _m4 = ms + f1, f2, f3, f4 = fs + mod1, mod2, mod3 = r.sample(list(_MODULES), 3) + return f"""\ +thread 'main' panicked at '{err}', {f1}:{r.randint(10, 300)} +stack backtrace: + 0: std::panicking::begin_panic + 1: {mod1}::{cls}::{m1} + at {f1}:{r.randint(10, 300)} + 2: {mod2}::{cls}::{m2} + at {f2}:{r.randint(10, 300)} + 3: {mod3}::{cls}::{m3} + at {f3}:{r.randint(10, 300)} + 4: {mod1}::main + at {f4}:{r.randint(10, 300)} + 5: std::rt::lang_start::{{{{closure}}}} + at /rustc/src/rt.rs:{r.randint(50, 200)} + 6: std::rt::lang_start + at /rustc/src/rt.rs:{r.randint(50, 200)} +note: run with `RUST_BACKTRACE=1` for a full backtrace +""" + + def _error_traceback_node( + self, + r, + *, + err: str, + cls: str, + ms: tuple[str, str, str, str], + fs: tuple[str, str, str, str], + ) -> str: + m1, m2, m3, m4 = ms + f1, f2, f3, f4 = fs + async_cls = r.choice(_CLASSES) + async_method = r.choice(_METHODS) + cls2 = r.choice(_CLASSES) + return f"""\ +Error: {err} + at {cls}.{m1} ({f1}:{r.randint(10, 300)}:{r.randint(1, 40)}) + at {cls}.{m2} ({f2}:{r.randint(10, 300)}:{r.randint(1, 40)}) + at {cls2}.{m3} ({f3}:{r.randint(10, 300)}:{r.randint(1, 40)}) + at processTicksAndRejections (node:internal/process/task_queues:{r.randint(50, 100)}) + at async {async_cls}.{async_method} ({f4}:{r.randint(10, 300)}) +Caused by: {r.choice(_ERROR_MESSAGES)} + at {cls2}.{m4} ({f3}:{r.randint(10, 300)}:{r.randint(1, 40)}) +""" + + def _gen_git_diff(self, language: str | None = None) -> str: + r = self._template_rng + file_pool = self._file_pool(language) + f1, f2, f3 = r.sample(list(file_pool), 3) + m1, m2, m3 = r.sample(_METHODS, 3) + v1, v2, v3 = r.sample(_VARS, 3) + cls = r.choice(_CLASSES) + ln = r.randint(10, 200) + ln2 = r.randint(50, 300) + err = r.choice(_ERROR_MESSAGES) + mod = r.choice(_MODULES) + idx = lambda: f"{r.randint(1000000, 9999999):07x}" # noqa: E731 + hunk_old, hunk_new = r.randint(1, 50), r.randint(1, 50) + commit_hash = f"{r.randint(1000000, 9999999):07x}" + + ctx = _DiffCtx( + cls=cls, + m1=m1, + m2=m2, + m3=m3, + v1=v1, + v2=v2, + v3=v3, + ln=ln, + ln2=ln2, + err=err, + mod=mod, + hunk_old=hunk_old, + hunk_new=hunk_new, + ) + builder = { + "python": _diff_hunks_python, + "go": _diff_hunks_go, + "rust": _diff_hunks_rust, + "typescript": _diff_hunks_typescript, + }.get(language, _diff_hunks_python) + hunk1, hunk2, hunk3 = builder(ctx) + + return f"""\ +commit {commit_hash} +Author: dev +Date: Mon Jan 15 14:32:00 2025 +0000 + + feat({mod}): add async {m1} with error handling + +diff --git a/{f1} b/{f1} +index {idx()}..{idx()} 100644 +--- a/{f1} ++++ b/{f1} +{hunk1} +diff --git a/{f2} b/{f2} +index {idx()}..{idx()} 100644 +--- a/{f2} ++++ b/{f2} +{hunk2} +diff --git a/{f3} b/{f3} +index {idx()}..{idx()} 100644 +--- a/{f3} ++++ b/{f3} +{hunk3} +""" diff --git a/src/aiperf/dataset/generator/_coding_go.py b/src/aiperf/dataset/generator/_coding_go.py new file mode 100644 index 0000000000..7986d67483 --- /dev/null +++ b/src/aiperf/dataset/generator/_coding_go.py @@ -0,0 +1,231 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Go code-template generators (mixin for CodingContentGenerator). + +Extracted from ``coding_content.py`` to keep that module under the +ergonomics file-size cap. Methods read ``self._template_rng`` and the +shared vocabulary tuples; behavior is unchanged. +""" + +from __future__ import annotations + +from aiperf.dataset.generator._coding_vocab import ( + _CLASSES, + _DB_TABLES, + _ERROR_MESSAGES, + _GO_PACKAGES, + _METHODS, + _MODULES, + _VARS, +) + + +class _GoMixin: + def _gen_go_code(self) -> str: + return self._template_rng.choice( + [ + self._gen_go_struct, + self._gen_go_http_handler, + self._gen_go_errors, + self._gen_go_test, + ] + )() + + def _gen_go_struct(self) -> str: + r = self._template_rng + pkg1, pkg2 = r.sample(list(_GO_PACKAGES), 2) + cls = r.choice(_CLASSES) + m1, m2 = r.sample(_METHODS, 2) + v1, v2, v3 = r.sample(_VARS, 3) + pkg_name = r.choice(_MODULES) + err = r.choice(_ERROR_MESSAGES) + + return f"""\ +package {pkg_name} + +import ( + "{pkg1}" + "{pkg2}" +) + +type {cls} struct {{{{ + {v1} string `json:"{v1}"` + {v2} int `json:"{v2},omitempty"` + {v3} bool `json:"-"` + mu sync.RWMutex +}}}} + +func New{cls}({v1} string) *{cls} {{{{ + return &{cls}{{{{{v1}: {v1}}}}} +}}}} + +func (s *{cls}) {m1.title()}(ctx context.Context) error {{{{ + s.mu.Lock() + defer s.mu.Unlock() + if s.{v1} == "" {{{{ + return {pkg1}.Errorf("{err}") + }}}} + s.{v2}++ + return nil +}}}} + +func (s *{cls}) {m2.title()}() (string, error) {{{{ + s.mu.RLock() + defer s.mu.RUnlock() + if !s.{v3} {{{{ + return "", {pkg1}.Errorf("%w: not initialized", Err{cls}) + }}}} + return {pkg2}.Sprintf("%s:%d", s.{v1}, s.{v2}), nil +}}}} +""" + + def _gen_go_http_handler(self) -> str: + r = self._template_rng + cls = r.choice(_CLASSES) + m1, m2 = r.sample(_METHODS, 2) + v1, v2 = r.sample(_VARS, 2) + pkg_name = r.choice(_MODULES) + table = r.choice(_DB_TABLES) + err = r.choice(_ERROR_MESSAGES) + status_code = r.choice( + ["http.StatusOK", "http.StatusCreated", "http.StatusAccepted"] + ) + + return f"""\ +package {pkg_name} + +import ( + "encoding/json" + "net/http" + "log/slog" +) + +type {m1.title()}Request struct {{{{ + {v1.title()} string `json:"{v1}" binding:"required"` + {v2.title()} int `json:"{v2}" binding:"gte=0"` +}}}} + +type {m1.title()}Response struct {{{{ + Items []map[string]any `json:"items"` + Total int `json:"total"` +}}}} + +func (h *{cls}) {m1.title()}Handler(w http.ResponseWriter, r *http.Request) {{{{ + var req {m1.title()}Request + if err := json.NewDecoder(r.Body).Decode(&req); err != nil {{{{ + slog.Error("{err}", "handler", "{m1}") + http.Error(w, err.Error(), http.StatusBadRequest) + return + }}}} + + items, err := h.svc.{m2.title()}(r.Context(), req.{v1.title()}) + if err != nil {{{{ + slog.Error("{err}", "table", "{table}") + http.Error(w, "{err}", http.StatusInternalServerError) + return + }}}} + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader({status_code}) + json.NewEncoder(w).Encode({m1.title()}Response{{{{Items: items, Total: len(items)}}}}) +}}}} +""" + + def _gen_go_errors(self) -> str: + r = self._template_rng + cls = r.choice(_CLASSES) + pkg_name = r.choice(_MODULES) + e1, e2, e3 = r.sample(_ERROR_MESSAGES, 3) + m1 = r.choice(_METHODS) + v1 = r.choice(_VARS) + + return f"""\ +package {pkg_name} + +import ( + "errors" + "fmt" +) + +var ( + Err{cls} = errors.New("{e1}") + ErrNot{m1.title()} = errors.New("{e2}") + ErrInvalid{v1.title()} = errors.New("{e3}") +) + +type {cls}Error struct {{{{ + Op string + {v1.title()} string + Err error +}}}} + +func (e *{cls}Error) Error() string {{{{ + return fmt.Sprintf("%s %s: %v", e.Op, e.{v1.title()}, e.Err) +}}}} + +func (e *{cls}Error) Unwrap() error {{{{ + return e.Err +}}}} + +func Wrap{cls}Error(op, {v1} string, err error) error {{{{ + return &{cls}Error{{{{Op: op, {v1.title()}: {v1}, Err: err}}}} +}}}} +""" + + def _gen_go_test(self) -> str: + r = self._template_rng + cls = r.choice(_CLASSES) + pkg_name = r.choice(_MODULES) + m1, m2 = r.sample(_METHODS, 2) + v1, v2 = r.sample(_VARS, 2) + + return f"""\ +package {pkg_name}_test + +import ( + "context" + "testing" +) + +func Test{cls}_{m1.title()}(t *testing.T) {{{{ + tests := []struct {{{{ + name string + {v1} string + want int + wantErr bool + }}}}{{{{ + {{{{"valid {v1}", "test_value", 42, false}}}}, + {{{{"empty {v1}", "", 0, true}}}}, + {{{{"long {v1}", "a]very_long_value_that_exceeds_limit", 0, true}}}}, + }}}} + + for _, tt := range tests {{{{ + t.Run(tt.name, func(t *testing.T) {{{{ + s := New{cls}(tt.{v1}) + got, err := s.{m1.title()}(context.Background()) + if (err != nil) != tt.wantErr {{{{ + t.Errorf("{m1.title()}() error = %v, wantErr %v", err, tt.wantErr) + return + }}}} + if got != tt.want {{{{ + t.Errorf("{m1.title()}() = %v, want %v", got, tt.want) + }}}} + }}}}) + }}}} +}}}} + +func Test{cls}_{m2.title()}_Concurrent(t *testing.T) {{{{ + s := New{cls}("{v2}") + ctx := context.Background() + errs := make(chan error, 10) + for i := 0; i < 10; i++ {{{{ + go func() {{{{ errs <- s.{m2.title()}(ctx) }}}}() + }}}} + for i := 0; i < 10; i++ {{{{ + if err := <-errs; err != nil {{{{ + t.Errorf("concurrent {m2}: %v", err) + }}}} + }}}} +}}}} +""" diff --git a/src/aiperf/dataset/generator/_coding_json.py b/src/aiperf/dataset/generator/_coding_json.py new file mode 100644 index 0000000000..c4336074ce --- /dev/null +++ b/src/aiperf/dataset/generator/_coding_json.py @@ -0,0 +1,122 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JSON-response template generators (mixin for CodingContentGenerator). + +Extracted from ``coding_content.py`` to keep that module under the +ergonomics file-size cap. Methods read ``self._template_rng`` and the +shared vocabulary tuples; behavior is unchanged. +""" + +from __future__ import annotations + +from aiperf.dataset.generator._coding_vocab import ( + _CLASSES, + _ERROR_MESSAGES, + _METHODS, + _MODULES, + _STATUS_CODES, + _VARS, +) + + +class _JsonMixin: + def _gen_json_response(self, language: str | None = None) -> str: + return self._template_rng.choice( + [ + self._gen_json_object, + self._gen_json_paginated, + self._gen_json_error, + ] + )() + + def _gen_json_object(self) -> str: + r = self._template_rng + m1, m2 = r.sample(_METHODS, 2) + v1, v2, v3 = r.sample(_VARS, 3) + cls = r.choice(_CLASSES) + id_suffix = r.randint(1000, 9999) + num_val = r.randint(0, 1000) + float_val = r.uniform(0, 1) + ts = r.randint(1700000000, 1800000000) + items = [ + f' {{{{"id": {r.randint(1, 999)}, "name": "{r.choice(_VARS)}"}}}}' + for _ in range(3) + ] + items_str = ",\n".join(items) + + return f"""\ +{{{{ + "status": "ok", + "data": {{{{ + "{v1}": "{cls.lower()}_{id_suffix}", + "{v2}": {num_val}, + "{v3}": {float_val:.4f}, + "metadata": {{{{ + "action": "{m1}", + "source": "{m2}", + "timestamp": "{ts}" + }}}}, + "items": [ +{items_str} + ] + }}}} +}}}} +""" + + def _gen_json_paginated(self) -> str: + r = self._template_rng + v1, v2 = r.sample(_VARS, 2) + cls = r.choice(_CLASSES) + total = r.randint(50, 5000) + page = r.randint(1, 20) + per_page = r.choice([10, 20, 50, 100]) + items = [ + f' {{{{"id": "{cls.lower()}_{r.randint(1000, 9999)}", "{v1}": "{r.choice(_MODULES)}", "{v2}": {r.randint(0, 100)}}}}}' + for _ in range(min(per_page, 5)) + ] + items_str = ",\n".join(items) + + return f"""\ +{{{{ + "data": [ +{items_str} + ], + "pagination": {{{{ + "page": {page}, + "per_page": {per_page}, + "total": {total}, + "total_pages": {(total + per_page - 1) // per_page}, + "has_next": {str(page * per_page < total).lower()}, + "has_prev": {str(page > 1).lower()} + }}}} +}}}} +""" + + def _gen_json_error(self) -> str: + r = self._template_rng + err = r.choice(_ERROR_MESSAGES) + status = r.choice(_STATUS_CODES) + code = status.split()[0] + trace_id = f"{r.randint(100000, 999999):06x}-{r.randint(100000, 999999):06x}" + v1 = r.choice(_VARS) + cls = r.choice(_CLASSES) + + return f"""\ +{{{{ + "error": {{{{ + "code": {code}, + "status": "{status}", + "message": "{err}", + "details": [ + {{{{ + "field": "{v1}", + "reason": "{err}", + "type": "{cls}" + }}}} + ], + "trace_id": "{trace_id}", + "documentation_url": "https://docs.example.com/errors/{code}" + }}}} +}}}} +""" diff --git a/src/aiperf/dataset/generator/_coding_ml.py b/src/aiperf/dataset/generator/_coding_ml.py new file mode 100644 index 0000000000..5e8010baf1 --- /dev/null +++ b/src/aiperf/dataset/generator/_coding_ml.py @@ -0,0 +1,261 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""ML/CUDA template generators (mixin for CodingContentGenerator). + +Extracted from ``coding_content.py`` to keep that module under the +ergonomics file-size cap. Methods read ``self._template_rng`` and the +shared vocabulary tuples; behavior is unchanged. +""" + +from __future__ import annotations + +from aiperf.dataset.generator._coding_vocab import ( + _CUDA_ERRORS, + _ML_CLASSES, + _ML_IMPORTS, + _ML_METHODS, + _ML_VARS, + _MODEL_NAMES, +) + + +class _MlMixin: + def _gen_ml_training_code(self) -> str: + r = self._template_rng + model = r.choice(_MODEL_NAMES) + imp1, imp2, imp3 = r.sample(list(_ML_IMPORTS), 3) + cls1, cls2 = r.sample(list(_ML_CLASSES), 2) + m1, m2 = r.sample(list(_ML_METHODS), 2) + v1, v2, v3, v4 = r.sample(list(_ML_VARS), 4) + lr = r.choice([1e-5, 2e-5, 5e-5, 1e-4, 3e-4]) + epochs = r.randint(1, 10) + bs = r.choice([1, 2, 4, 8, 16, 32]) + grad_accum = r.choice([1, 2, 4, 8]) + + return f"""\ +import {imp1} +import {imp2} +from {imp3} import {cls1}, {cls2} + +model_name = "{model}" +tokenizer = {cls2}.from_pretrained(model_name) +model = {cls1}.from_pretrained( + model_name, + torch_dtype=torch.bfloat16, + device_map="auto", +) + +train_dataset = datasets.load_dataset("json", data_files="train.jsonl", split="train") + +training_args = TrainingArguments( + output_dir="./checkpoints", + num_train_epochs={epochs}, + per_device_train_batch_size={bs}, + gradient_accumulation_steps={grad_accum}, + learning_rate={lr}, + max_grad_norm=1.0, + warmup_ratio=0.1, + bf16=True, + logging_steps=10, + save_strategy="epoch", + report_to="wandb", +) + +optimizer = torch.optim.AdamW(model.parameters(), lr={lr}, weight_decay=0.01) + +for epoch in range({epochs}): + model.train() + for step, batch in enumerate(train_loader): + {v1} = batch["{v1}"].to("cuda") + {v2} = batch["{v2}"].to("cuda") + outputs = model({m1}={v1}, {v2}={v2}) + {v3} = outputs.{v3} + {v3}.{m2}() + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) + optimizer.step() + optimizer.zero_grad() + + if step % 10 == 0: + print(f"Epoch {{epoch}} Step {{step}} {v3}: {{{{{v3}.item():.4f}}}}") + +model.save_pretrained("./final_model") +tokenizer.save_pretrained("./final_model") +""" + + def _gen_ml_inference_code(self) -> str: + r = self._template_rng + model = r.choice(_MODEL_NAMES) + cls1 = r.choice(("AutoModelForCausalLM", "AutoModelForSeq2SeqLM")) + v1, v2, v3 = r.sample(list(_ML_VARS), 3) + temp = r.choice([0.1, 0.3, 0.7, 1.0]) + top_p = r.choice([0.9, 0.95, 1.0]) + max_new = r.choice([128, 256, 512, 1024, 2048]) + + return f"""\ +import torch +from transformers import {cls1}, AutoTokenizer, GenerationConfig + +model_name = "{model}" +tokenizer = AutoTokenizer.from_pretrained(model_name) +model = {cls1}.from_pretrained( + model_name, + torch_dtype=torch.float16, + device_map="auto", + attn_implementation="flash_attention_2", +) + +generation_config = GenerationConfig( + max_new_tokens={max_new}, + temperature={temp}, + top_p={top_p}, + do_sample={"True" if temp > 0 else "False"}, + repetition_penalty=1.1, +) + +prompt = "Explain the architecture of a transformer model." +{v1} = tokenizer(prompt, return_tensors="pt").to(model.device) + +with torch.inference_mode(): + {v2} = model.generate( + **{v1}, + generation_config=generation_config, + pad_token_id=tokenizer.eos_token_id, + ) + +{v3} = tokenizer.batch_decode({v2}[:, {v1}["{v1}"].shape[-1]:], skip_special_tokens=True) +print({v3}[0]) +""" + + def _gen_ml_config(self) -> str: + r = self._template_rng + model = r.choice(_MODEL_NAMES) + lr = r.choice([1e-5, 2e-5, 5e-5, 1e-4, 3e-4]) + epochs = r.randint(1, 10) + bs = r.choice([1, 2, 4, 8, 16, 32]) + grad_accum = r.choice([1, 2, 4, 8]) + max_len = r.choice([512, 1024, 2048, 4096]) + warmup = r.choice([0.03, 0.05, 0.1]) + lora_r = r.choice([8, 16, 32, 64]) + lora_alpha = lora_r * 2 + quant_bits = r.choice([4, 8]) + + return f"""\ +{{{{ + "model_name_or_path": "{model}", + "torch_dtype": "bfloat16", + "attn_implementation": "flash_attention_2", + "max_seq_length": {max_len}, + "training": {{{{ + "num_train_epochs": {epochs}, + "per_device_train_batch_size": {bs}, + "gradient_accumulation_steps": {grad_accum}, + "learning_rate": {lr}, + "weight_decay": 0.01, + "warmup_ratio": {warmup}, + "lr_scheduler_type": "cosine", + "max_grad_norm": 1.0, + "bf16": true, + "gradient_checkpointing": true, + "optim": "adamw_torch_fused" + }}}}, + "lora": {{{{ + "r": {lora_r}, + "lora_alpha": {lora_alpha}, + "lora_dropout": 0.05, + "target_modules": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + "task_type": "CAUSAL_LM" + }}}}, + "quantization": {{{{ + "load_in_{quant_bits}bit": true, + "bnb_{quant_bits}bit_compute_dtype": "bfloat16", + "bnb_{quant_bits}bit_quant_type": "nf4", + "bnb_{quant_bits}bit_use_double_quant": true + }}}}, + "data": {{{{ + "dataset_name": "train.jsonl", + "max_length": {max_len}, + "packing": true, + "num_proc": 8 + }}}} +}}}} +""" + + def _gen_ml_training_log(self) -> str: + r = self._template_rng + model = r.choice(_MODEL_NAMES).split("/")[-1] + total_steps = r.randint(500, 10000) + epoch = r.randint(0, 5) + lines = [] + + for _ in range(r.randint(8, 15)): + step = r.randint(1, total_steps) + loss = r.uniform(0.3, 4.0) + lr_val = r.uniform(1e-6, 5e-4) + grad = r.uniform(0.1, 10.0) + tokens_per_sec = r.randint(1000, 50000) + lines.append( + f"{{{{'step': {step}, 'epoch': {epoch + step / total_steps:.2f}, " + f"'loss': {loss:.4f}, 'lr': {lr_val:.2e}, " + f"'grad_norm': {grad:.3f}, 'tokens_per_sec': {tokens_per_sec}}}}}" + ) + + gpu_mem = r.uniform(10, 80) + gpu_util = r.randint(80, 100) + eval_loss = r.uniform(0.5, 3.0) + eval_ppl = r.uniform(2.0, 20.0) + + lines.append( + f"\n[Eval] epoch={epoch + 1} loss={eval_loss:.4f} perplexity={eval_ppl:.2f}" + ) + lines.append(f"[GPU] memory_allocated={gpu_mem:.1f}GB utilization={gpu_util}%") + lines.append( + f"[GPU] peak_memory={gpu_mem + r.uniform(1, 10):.1f}GB " + f"reserved={gpu_mem + r.uniform(5, 20):.1f}GB" + ) + lines.append( + f"[Checkpoint] Saved model checkpoint to ./checkpoints/{model}/step-{total_steps}" + ) + + return "\n".join(lines) + "\n" + + def _gen_cuda_error(self) -> str: + r = self._template_rng + err = r.choice(_CUDA_ERRORS) + model = r.choice(_MODEL_NAMES).split("/")[-1] + rank = r.randint(0, 7) + gpu_id = r.randint(0, 7) + alloc_gb = r.uniform(0.5, 16.0) + total_gb = r.choice([24.0, 40.0, 48.0, 80.0]) + free_gb = r.uniform(0.01, 2.0) + cls1, cls2 = r.sample(list(_ML_CLASSES), 2) + m1, m2 = r.sample(list(_ML_METHODS), 2) + + return f"""\ +Traceback (most recent call last): + File "train.py", line {r.randint(50, 300)}, in main + outputs = model.{m1}(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + File "torch/nn/modules/module.py", line {r.randint(1400, 1600)}, in _wrapped_call_impl + return self._call_impl(*args, **kwargs) + File "transformers/models/llama/modeling_llama.py", line {r.randint(800, 1200)}, in {m1} + hidden_states = self.model(input_ids, attention_mask=attention_mask) + File "torch/nn/modules/module.py", line {r.randint(1400, 1600)}, in _wrapped_call_impl + return self._call_impl(*args, **kwargs) + File "transformers/models/llama/modeling_llama.py", line {r.randint(400, 800)}, in {m2} + layer_outputs = decoder_layer(hidden_states, attention_mask=attention_mask) +{err} + +|===========================================================================| +| PyTorch CUDA memory summary, device: {gpu_id} | +|---------------------------------------------------------------------------| +| CUDA OOMs: {r.randint(1, 5):>10} | +|---------------------------------------------------------------------------| +| Metric | Cur Usage | Peak Usage | Total Alloc | +|---------------------------------------------------------------------------| +| Allocated memory | {alloc_gb:>8.2f} GB | {total_gb - free_gb:>8.2f} GB | {total_gb * r.randint(2, 10):>9.2f} GB | +| Reserved memory | {total_gb - free_gb + 1:>8.2f} GB | {total_gb:>8.2f} GB | {total_gb * r.randint(2, 10):>9.2f} GB | +| Free memory | {free_gb:>8.2f} GB | | | +|===========================================================================| + +Model: {model} | Rank: {rank} | GPU: {gpu_id} (NVIDIA A100-SXM4-{int(total_gb)}GB) +""" diff --git a/src/aiperf/dataset/generator/_coding_prompts_conv.py b/src/aiperf/dataset/generator/_coding_prompts_conv.py new file mode 100644 index 0000000000..6139afcac7 --- /dev/null +++ b/src/aiperf/dataset/generator/_coding_prompts_conv.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""User-prompt, context, and coding-conversation generators (mixin). + +Extracted from ``coding_content.py`` to keep that module under the +ergonomics file-size cap. Methods read ``self._template_rng`` and the +shared vocabulary tuples; behavior is unchanged. +""" + +from __future__ import annotations + +from aiperf.dataset.generator._coding_text import ( + _USER_REQUESTS, +) +from aiperf.dataset.generator._coding_vocab import ( + _CLASSES, + _ERROR_MESSAGES, + _FILE_PATHS, + _METHODS, + _MODULES, + _TYPES, + _VARS, +) + + +class _SafeFormatMap(dict): + """Dict subclass that returns '{key}' for missing keys in str.format_map.""" + + def __missing__(self, key: str) -> str: + return f"{{{key}}}" + + +class _PromptsConvMixin: + def _gen_user_prompt(self) -> str: + r = self._template_rng + template = r.choice(_USER_REQUESTS) + base = template.format( + module=r.choice(_MODULES), + cls=r.choice(_CLASSES), + method=r.choice(_METHODS), + var=r.choice(_VARS), + error=r.choice(_ERROR_MESSAGES), + type=r.choice(_TYPES), + ) + + if r.random() < 0.3: + base += "\n\n" + self._gen_prompt_context() + return base + + def _gen_prompt_context(self) -> str: + r = self._template_rng + kind = r.choice(["snippet", "error_output", "constraint"]) + if kind == "snippet": + cls = r.choice(_CLASSES) + m1 = r.choice(_METHODS) + v1, v2 = r.sample(_VARS, 2) + f = r.choice(_FILE_PATHS) + return ( + f"Here's the relevant code from `{f}`:\n\n" + f"```\n" + f"class {cls}:\n" + f" def {m1}(self, {v1}):\n" + f" {v2} = self._{v1}\n" + f" return {v2}\n" + f"```" + ) + elif kind == "error_output": + err = r.choice(_ERROR_MESSAGES) + cls = r.choice(_CLASSES) + m1 = r.choice(_METHODS) + f = r.choice(_FILE_PATHS) + return ( + f"Error output:\n\n" + f"```\n" + f' File "{f}", line {r.randint(10, 300)}, in {m1}\n' + f' raise RuntimeError("{err}")\n' + f"RuntimeError: {err}\n" + f"```" + ) + else: + return r.choice( + ( + "Constraint: no new dependencies allowed in this PR.", + "This is on the hot path — keep allocations minimal.", + "Must remain backward-compatible with the v1 API.", + f"The {r.choice(_MODULES)} service is frozen — only touch {r.choice(_MODULES)}.", + f"Target is under {r.randint(5, 50)}ms p99 latency.", + "We need this for the release on Friday — keep it simple.", + ) + ) + + def _gen_coding_conversation(self) -> str: + r = self._template_rng + return r.choice( + [ + self._gen_conv_bugfix, + self._gen_conv_review, + self._gen_conv_feature, + self._gen_conv_debug, + self._gen_conv_qa, + self._gen_conv_refactor, + self._gen_conv_perf, + self._gen_conv_cicd, + self._gen_conv_ml_debug, + self._gen_conv_test_write, + self._gen_conv_migration, + self._gen_conv_deploy, + self._gen_conv_security, + self._gen_conv_distributed, + self._gen_conv_observability, + self._gen_conv_db_optimize, + self._gen_conv_architecture_review, + self._gen_conv_incident_response, + ] + )() + + def _conv_ids(self) -> dict[str, str]: + r = self._template_rng + return { + "cls": r.choice(_CLASSES), + "module": r.choice(_MODULES), + "method": r.choice(_METHODS), + "var": r.choice(_VARS), + "error": r.choice(_ERROR_MESSAGES), + } + + def _conv_bridge(self, pool: tuple[str, ...], ids: dict[str, str]) -> str: + r = self._template_rng + return r.choice(pool).format_map(_SafeFormatMap(ids)) + + def _conv_user_msg(self, ids: dict[str, str]) -> str: + r = self._template_rng + template = r.choice(_USER_REQUESTS) + return template.format_map(_SafeFormatMap(ids)) diff --git a/src/aiperf/dataset/generator/_coding_python.py b/src/aiperf/dataset/generator/_coding_python.py new file mode 100644 index 0000000000..635f9483f5 --- /dev/null +++ b/src/aiperf/dataset/generator/_coding_python.py @@ -0,0 +1,280 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Python code-template generators (mixin for CodingContentGenerator). + +Extracted from ``coding_content.py`` to keep that module under the +ergonomics file-size cap. Methods read ``self._template_rng`` and the +shared vocabulary tuples; behavior is unchanged. +""" + +from __future__ import annotations + +from aiperf.dataset.generator._coding_vocab import ( + _CLASSES, + _DB_TABLES, + _DECORATORS, + _ERROR_MESSAGES, + _HTTP_ROUTES, + _METHODS, + _MODULES, + _TYPES, + _VARS, +) + + +class _PythonMixin: + def _gen_python_code(self) -> str: + return self._template_rng.choice( + [ + self._gen_python_class, + self._gen_python_functions, + self._gen_python_test, + self._gen_python_http_handler, + self._gen_python_data_model, + ] + )() + + def _gen_python_class(self) -> str: + r = self._template_rng + cls = r.choice(_CLASSES) + mod = r.choice(_MODULES) + m1, m2, m3 = r.sample(_METHODS, 3) + v1, v2, v3 = r.sample(_VARS, 3) + t1, t2 = r.sample(_TYPES, 2) + dec = r.choice(_DECORATORS) + imp_mod = r.choice(_MODULES) + imp_cls = r.choice(_CLASSES) + err = r.choice(_ERROR_MESSAGES) + + return f"""\ +import {mod} +from {mod}.{imp_mod} import {imp_cls} + + +class {cls}: + \"\"\"Handles {m1} operations for {mod}.\"\"\" + + _default_{v3} = 64 + + def __init__(self, {v1}: {t1}, {v2}: {t2} = None): + self._{v1} = {v1} + self._{v2} = {v2} + self._{v3} = self._default_{v3} + self._initialized = False + + {dec} + async def {m1}(self, {v1}: {t1}) -> {t2}: + if not self._initialized: + raise RuntimeError("{cls} not initialized") + {v2} = await self._{m2}({v1}) + return {v2} + + async def _{m2}(self, {v1}: {t1}) -> {t2}: + try: + {v2} = {mod}.{m2}({v1}) + return {v2} + except Exception as e: + raise ValueError("{err}") from e + + def {m3}(self) -> None: + self._initialized = True + self._{v3} = 0 +""" + + def _gen_python_functions(self) -> str: + r = self._template_rng + m1, m2, m3 = r.sample(_METHODS, 3) + v1, v2, v3 = r.sample(_VARS, 3) + t1, t2, t3 = r.sample(_TYPES, 3) + mod = r.choice(_MODULES) + imp_mod = r.choice(_MODULES) + cls = r.choice(_CLASSES) + err = r.choice(_ERROR_MESSAGES) + + return f"""\ +from __future__ import annotations + +import asyncio +import logging +from contextlib import asynccontextmanager +from typing import AsyncIterator + +from {mod}.{imp_mod} import {cls} + +logger = logging.getLogger(__name__) + + +async def {m1}({v1}: {t1}, {v2}: {t2} | None = None) -> {t3}: + async with _acquire_{v3}({v1}) as {v3}: + {v2} = await {cls}().{m2}({v3}) + return [{v2} for _ in range(10) if {v2} is not None] + + +@asynccontextmanager +async def _acquire_{v3}({v1}: {t1}) -> AsyncIterator[{t2}]: + {v3} = {mod}.{m3}({v1}) + try: + yield {v3} + finally: + await {v3}.close() + + +def {m2}_sync({v1}: {t1}, *, max_retries: int = 3) -> {t2}: + for attempt in range(max_retries): + try: + return {mod}.{m2}({v1}) + except RuntimeError: + if attempt == max_retries - 1: + raise + logger.warning("{err}, attempt %d", attempt + 1) + raise AssertionError("unreachable") +""" + + def _gen_python_test(self) -> str: + r = self._template_rng + cls = r.choice(_CLASSES) + mod = r.choice(_MODULES) + m1, m2, m3 = r.sample(_METHODS, 3) + v1, v2 = r.sample(_VARS, 2) + err = r.choice(_ERROR_MESSAGES) + + return f"""\ +import pytest +from unittest.mock import AsyncMock, patch + +from {mod} import {cls} + + +class Test{cls}: + @pytest.fixture + def instance(self): + return {cls}({v1}="test_value") + + @pytest.mark.asyncio + async def test_{m1}_returns_expected(self, instance): + instance._{m2} = AsyncMock(return_value=42) + result = await instance.{m1}() + assert result == 42 + instance._{m2}.assert_awaited_once() + + @pytest.mark.parametrize("{v1}", ["alpha", "beta", "gamma"]) + def test_{m2}_with_values(self, instance, {v1}): + instance._{v1} = {v1} + result = instance.{m2}() + assert result is not None + + @pytest.mark.asyncio + async def test_{m3}_raises_on_{v2}(self, instance): + with pytest.raises(ValueError, match="{err}"): + await instance.{m3}(None) + + @pytest.mark.asyncio + async def test_{m1}_with_mock_dependency(self, instance): + with patch("{mod}.{m2}") as mock: + mock.return_value = {{{{"key": "{v2}"}}}}\n result = await instance.{m1}() + assert "{v2}" in str(result) +""" + + def _gen_python_http_handler(self) -> str: + r = self._template_rng + cls = r.choice(_CLASSES) + mod = r.choice(_MODULES) + m1, m2 = r.sample(_METHODS, 2) + v1, v2, v3 = r.sample(_VARS, 3) + route = r.choice(_HTTP_ROUTES) + table = r.choice(_DB_TABLES) + err = r.choice(_ERROR_MESSAGES) + + return f"""\ +from __future__ import annotations + +from fastapi import APIRouter, Depends, HTTPException, status +from pydantic import BaseModel, Field + +from {mod}.{cls.lower()} import {cls} + +router = APIRouter(prefix="{route}", tags=["{mod}"]) + + +class {cls}Request(BaseModel): + {v1}: str = Field(description="Primary {v1} identifier") + {v2}: int = Field(default=10, ge=1, le=100, description="Page size") + {v3}: str | None = Field(default=None, description="Optional filter") + + +class {cls}Response(BaseModel): + items: list[dict] = Field(description="Result items from {table}") + total: int = Field(description="Total count") + page: int = Field(description="Current page number") + + +@router.post("/", response_model={cls}Response, status_code=201) +async def {m1}( + body: {cls}Request, + svc: {cls} = Depends(), +) -> {cls}Response: + try: + items = await svc.{m1}(body.{v1}, page_size=body.{v2}) + except ValueError as exc: + raise HTTPException(status_code=422, detail=str(exc)) from exc + return {cls}Response(items=items, total=len(items), page=1) + + +@router.get("/{{{{{v1}}}}}") +async def {m2}({v1}: str, svc: {cls} = Depends()) -> dict: + result = await svc.{m2}({v1}) + if result is None: + raise HTTPException(status_code=404, detail="{err}") + return {{"status": "ok", "data": result}} +""" + + def _gen_python_data_model(self) -> str: + r = self._template_rng + cls = r.choice(_CLASSES) + v1, v2, v3, v4 = r.sample(_VARS, 4) + m1 = r.choice(_METHODS) + table = r.choice(_DB_TABLES) + + return f"""\ +from __future__ import annotations + +from datetime import datetime +from enum import StrEnum + +from pydantic import BaseModel, Field, field_validator + + +class {cls}Status(StrEnum): + PENDING = "pending" + ACTIVE = "active" + SUSPENDED = "suspended" + DELETED = "deleted" + + +class {cls}Config(BaseModel): + {v1}: str = Field(description="{cls} {v1} identifier") + {v2}: int = Field(default=0, ge=0, description="Current {v2} count") + {v3}: float = Field(default=1.0, gt=0, description="Rate limit for {m1}") + status: {cls}Status = Field(default={cls}Status.PENDING, description="Lifecycle status") + {v4}: dict[str, str] = Field(default_factory=dict, description="Arbitrary {v4}") + created_at: datetime = Field(default_factory=datetime.utcnow, description="Creation timestamp") + source_table: str = Field(default="{table}", description="Backing store table") + + @field_validator("{v1}") + @classmethod + def _validate_{v1}(cls, v: str) -> str: + if not v or len(v) > 256: + raise ValueError("{v1} must be 1-256 characters") + return v.strip() + + @field_validator("{v3}") + @classmethod + def _validate_{v3}(cls, v: float) -> float: + if v > 10_000: + raise ValueError("{v3} exceeds max rate") + return v + + def {m1}(self) -> bool: + return self.status == {cls}Status.ACTIVE and self.{v2} > 0 +""" diff --git a/src/aiperf/dataset/generator/_coding_rust.py b/src/aiperf/dataset/generator/_coding_rust.py new file mode 100644 index 0000000000..6dfdf157bb --- /dev/null +++ b/src/aiperf/dataset/generator/_coding_rust.py @@ -0,0 +1,235 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Rust code-template generators (mixin for CodingContentGenerator). + +Extracted from ``coding_content.py`` to keep that module under the +ergonomics file-size cap. Methods read ``self._template_rng`` and the +shared vocabulary tuples; behavior is unchanged. +""" + +from __future__ import annotations + +from aiperf.dataset.generator._coding_vocab import ( + _CLASSES, + _ERROR_MESSAGES, + _METHODS, + _MODULES, + _RUST_CRATES, + _VARS, +) + + +class _RustMixin: + def _gen_rust_code(self) -> str: + return self._template_rng.choice( + [ + self._gen_rust_struct, + self._gen_rust_http_handler, + self._gen_rust_errors, + self._gen_rust_test, + ] + )() + + def _gen_rust_struct(self) -> str: + r = self._template_rng + cr1, cr2 = r.sample(list(_RUST_CRATES), 2) + cls = r.choice(_CLASSES) + m1, m2 = r.sample(_METHODS, 2) + v1, v2, v3 = r.sample(_VARS, 3) + err = r.choice(_ERROR_MESSAGES) + + return f"""\ +use {cr1}; +use {cr2}; + +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct {cls} {{{{ + {v1}: String, + {v2}: Vec, + #[serde(default)] + {v3}: Option, + initialized: bool, +}}}} + +impl {cls} {{{{ + pub fn new({v1}: impl Into) -> Self {{{{ + Self {{{{ + {v1}: {v1}.into(), + {v2}: Vec::new(), + {v3}: None, + initialized: false, + }}}} + }}}} + + pub async fn {m1}(&mut self) -> Result<(), anyhow::Error> {{{{ + if !self.initialized {{{{ + anyhow::bail!("{err}"); + }}}} + self.{m2}().await + }}}} + + async fn {m2}(&self) -> Result<(), anyhow::Error> {{{{ + let _{v2} = self.{v1}.as_bytes(); + tracing::debug!("{m2} completed for {{}}", self.{v1}); + Ok(()) + }}}} +}}}} +""" + + def _gen_rust_http_handler(self) -> str: + r = self._template_rng + cls = r.choice(_CLASSES) + m1, m2 = r.sample(_METHODS, 2) + v1, v2 = r.sample(_VARS, 2) + mod = r.choice(_MODULES) + + return f"""\ +use axum::{{extract::{{Path, State}}, http::StatusCode, Json}}; +use serde::{{Deserialize, Serialize}}; +use std::sync::Arc; + +use crate::{mod}::{cls}; + +#[derive(Debug, Deserialize)] +pub struct {m1.title()}Request {{{{ + {v1}: String, + {v2}: Option, +}}}} + +#[derive(Debug, Serialize)] +pub struct {m1.title()}Response {{{{ + id: String, + {v1}: String, + created: bool, +}}}} + +pub async fn {m1}_handler( + State(svc): State>, + Json(body): Json<{m1.title()}Request>, +) -> Result, StatusCode> {{{{ + let result = svc + .{m1}(&body.{v1}) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + + Ok(Json({m1.title()}Response {{{{ + id: result.id.to_string(), + {v1}: body.{v1}, + created: true, + }}}})) +}}}} + +pub async fn {m2}_handler( + State(svc): State>, + Path({v1}): Path, +) -> Result, StatusCode> {{{{ + svc.{m2}(&{v1}) + .await + .map(|v| Json(serde_json::json!({{{{"status": "ok", "data": v}}}}))) + .map_err(|_| StatusCode::NOT_FOUND) +}}}} +""" + + def _gen_rust_errors(self) -> str: + r = self._template_rng + cls = r.choice(_CLASSES) + e1, e2, e3 = r.sample(_ERROR_MESSAGES, 3) + v1 = r.choice(_VARS) + mod = r.choice(_MODULES) + + return f"""\ +use thiserror::Error; + +#[derive(Debug, Error)] +pub enum {cls}Error {{{{ + #[error("{e1}")] + NotInitialized, + + #[error("{e2}: {{{{{v1}}}}}")] + InvalidInput {{{{ {v1}: String }}}}, + + #[error("{e3}")] + Internal(#[from] anyhow::Error), + + #[error("io error in {mod}")] + Io(#[from] std::io::Error), + + #[error("serialization failed")] + Serde(#[from] serde_json::Error), +}}}} + +impl {cls}Error {{{{ + pub fn is_retryable(&self) -> bool {{{{ + matches!(self, Self::Internal(_) | Self::Io(_)) + }}}} + + pub fn status_code(&self) -> u16 {{{{ + match self {{{{ + Self::NotInitialized => 503, + Self::InvalidInput {{{{ .. }}}} => 400, + Self::Internal(_) => 500, + Self::Io(_) => 502, + Self::Serde(_) => 422, + }}}} + }}}} +}}}} +""" + + def _gen_rust_test(self) -> str: + r = self._template_rng + cls = r.choice(_CLASSES) + m1, m2 = r.sample(_METHODS, 2) + v1, v2 = r.sample(_VARS, 2) + err = r.choice(_ERROR_MESSAGES) + cr = r.choice(_RUST_CRATES) + + return f"""\ +use {cr}; + +#[cfg(test)] +mod tests {{{{ + use super::*; + + fn make_{cls.lower()}() -> {cls} {{{{ + {cls}::new("{v1}_test") + }}}} + + #[tokio::test] + async fn test_{m1}_success() {{{{ + let mut svc = make_{cls.lower()}(); + svc.initialized = true; + let result = svc.{m1}().await; + assert!(result.is_ok(), "expected Ok, got {{:?}}", result); + }}}} + + #[tokio::test] + async fn test_{m1}_not_initialized() {{{{ + let mut svc = make_{cls.lower()}(); + let err = svc.{m1}().await.unwrap_err(); + assert!(err.to_string().contains("{err}")); + }}}} + + #[test] + fn test_{m2}_returns_bytes() {{{{ + let svc = make_{cls.lower()}(); + let {v2} = svc.{v1}.as_bytes(); + assert!(!{v2}.is_empty()); + }}}} + + #[tokio::test] + async fn test_{m1}_concurrent() {{{{ + let svc = std::sync::Arc::new(tokio::sync::Mutex::new(make_{cls.lower()}())); + let mut handles = vec![]; + for _ in 0..5 {{{{ + let svc = svc.clone(); + handles.push(tokio::spawn(async move {{{{ + svc.lock().await.{m1}().await + }}}})); + }}}} + for h in handles {{{{ + let _ = h.await.unwrap(); + }}}} + }}}} +}}}} +""" diff --git a/src/aiperf/dataset/generator/_coding_sql.py b/src/aiperf/dataset/generator/_coding_sql.py new file mode 100644 index 0000000000..aea1d80c5f --- /dev/null +++ b/src/aiperf/dataset/generator/_coding_sql.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""SQL query generator (mixin for CodingContentGenerator). + +Extracted from ``coding_content.py`` to keep that module under the +ergonomics file-size cap. Methods read ``self._template_rng`` and the +shared vocabulary tuples; behavior is unchanged. +""" + +from __future__ import annotations + +from aiperf.dataset.generator._coding_vocab import ( + _DB_TABLES, + _MODULES, + _STATUS_CODES, + _VARS, +) + + +class _SqlMixin: + def _gen_sql_query(self) -> str: + r = self._template_rng + t1, t2, t3 = r.sample(list(_DB_TABLES), 3) + v1, v2, v3 = r.sample(list(_VARS), 3) + kind = r.choice(["select_join", "insert", "create", "alter"]) + + if kind == "select_join": + return self._sql_select_join(t1, t2, t3, v1=v1, v2=v2, v3=v3) + if kind == "insert": + return self._sql_insert(t1, v1, v2) + if kind == "create": + return self._sql_create(t1, v1, v2, v3) + return self._sql_alter(t1, t2, v1, v2) + + def _sql_select_join( + self, + t1: str, + t2: str, + t3: str, + *, + v1: str, + v2: str, + v3: str, + ) -> str: + r = self._template_rng + limit = r.randint(10, 1000) + offset = r.randint(0, 500) + return f"""\ +SELECT + t1.id, + t1.{v1}, + t1.created_at, + t2.{v2}, + t2.{v3}, + COUNT(t3.id) AS {v3}_count +FROM {t1} t1 +INNER JOIN {t2} t2 ON t2.{t1}_id = t1.id +LEFT JOIN {t3} t3 ON t3.{t2}_id = t2.id +WHERE t1.status = 'active' + AND t1.created_at >= NOW() - INTERVAL '30 days' + AND t2.{v2} IS NOT NULL +GROUP BY t1.id, t1.{v1}, t1.created_at, t2.{v2}, t2.{v3} +HAVING COUNT(t3.id) > 0 +ORDER BY t1.created_at DESC +LIMIT {limit} OFFSET {offset}; +""" + + def _sql_insert(self, t1: str, v1: str, v2: str) -> str: + r = self._template_rng + n_rows = r.randint(1, 5) + rows = [] + for _ in range(n_rows): + rows.append( + f" ('{r.choice(_MODULES)}', {r.randint(1, 1000)}, " + f"'{r.choice(_STATUS_CODES).split()[0]}', NOW())" + ) + rows_str = ",\n".join(rows) + return f"""\ +INSERT INTO {t1} ({v1}, {v2}, status, created_at) +VALUES +{rows_str} +ON CONFLICT ({v1}) +DO UPDATE SET + {v2} = EXCLUDED.{v2}, + status = EXCLUDED.status, + updated_at = NOW() +RETURNING id, {v1}, {v2}; +""" + + @staticmethod + def _sql_create(t1: str, v1: str, v2: str, v3: str) -> str: + return f"""\ +CREATE TABLE IF NOT EXISTS {t1} ( + id BIGSERIAL PRIMARY KEY, + {v1} VARCHAR(256) NOT NULL, + {v2} INTEGER DEFAULT 0, + {v3} JSONB DEFAULT '{{}}'::jsonb, + status VARCHAR(32) DEFAULT 'pending', + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + CONSTRAINT {t1}_{v1}_unique UNIQUE ({v1}) +); + +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_{t1}_{v1} + ON {t1} ({v1}); +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_{t1}_status_created + ON {t1} (status, created_at DESC); +""" + + def _sql_alter(self, t1: str, t2: str, v1: str, v2: str) -> str: + r = self._template_rng + col_type = r.choice( + ["VARCHAR(256)", "INTEGER", "BOOLEAN", "JSONB", "TIMESTAMPTZ"] + ) + return f"""\ +BEGIN; + +ALTER TABLE {t1} + ADD COLUMN IF NOT EXISTS {v1} {col_type}, + ADD COLUMN IF NOT EXISTS {v2} INTEGER DEFAULT 0; + +UPDATE {t1} +SET {v1} = ( + SELECT {v2} FROM {t2} + WHERE {t2}.{t1}_id = {t1}.id + LIMIT 1 +) +WHERE {v1} IS NULL; + +ALTER TABLE {t1} + ALTER COLUMN {v1} SET NOT NULL; + +COMMIT; +""" diff --git a/src/aiperf/dataset/generator/_coding_text.py b/src/aiperf/dataset/generator/_coding_text.py new file mode 100644 index 0000000000..5fc4bd3079 --- /dev/null +++ b/src/aiperf/dataset/generator/_coding_text.py @@ -0,0 +1,329 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Long-form natural-language tuples for CodingContentGenerator. + +Split from ``_coding_vocab.py`` so each vocabulary module stays under +the file-size cap. Contains user-request prompts, multi-turn bridge +phrases, follow-up questions, and pool-block-count weights. +""" + +from __future__ import annotations + +# fmt: off + + +_USER_REQUESTS = ( + # simple one-liners (original) + "Fix the failing test in {module} — it returns {error}", + "Add retry logic to {cls}.{method}() with exponential backoff", + "Refactor the {method} function to use async/await instead of callbacks", + "The {cls} class is throwing {error} when {var} is None", + "Add input validation for the {var} parameter in {method}()", + "Write unit tests for {cls}.{method}() covering edge cases", + "Optimize the {method} query — it's taking too long with large datasets", + "Add logging to {cls} so we can debug {error} in production", + "Move the {method} logic from {module} to a shared utility", + "Implement caching for {cls}.{method}() to reduce database load", + "Update the {module} config to support environment variable overrides", + "Add a health check endpoint that verifies {cls} connectivity", + "The CI is failing because {module} import is broken after the refactor", + "Create a migration script for the {var} schema change", + "Add rate limiting to the {method} endpoint — we're getting hammered", + "Debug why {cls}.{method}() returns stale data after {method}()", + "Add pagination support to the {method}() response", + "Implement graceful shutdown for the {cls} worker pool", + "The {module} integration test is flaky — fix the race condition", + "Add type hints to all public methods in {cls}", + "Refactor {module} to use dependency injection instead of globals", + "Add metrics collection for {method}() latency and error rates", + "Fix the memory leak in {cls} — it's not releasing {var} properly", + "Implement {method} fallback when the primary {module} is unavailable", + "Add request/response logging middleware for the {module} API", + "Write a load test for {cls}.{method}() with concurrent connections", + "Add circuit breaker pattern to {cls} for external API calls", + "The {cls}.{method}() docstring is wrong — update it to match the code", + "Implement batch processing for {method}() to handle bulk {var} updates", + "Add WebSocket support to {cls} for real-time {var} updates", + # multi-step tasks + "Migrate {cls}.{method}() from sync to async — it's called in 3 places across {module} and needs backward compat", + "Split the {cls} class into two: one for {method} and one for the {var} lifecycle management", + "We need to add {method}() to {cls}, then wire it into the {module} pipeline and add an integration test", + "Extract the {method} logic from {cls} into a standalone service, update all callers, and add a deprecation warning to the old path", + "Rewrite the {module} retry logic: replace the sleep loop with a proper backoff strategy using {cls}", + # error context prompts + "Getting {error} after upgrading {module} to the latest version — only happens under load", + "The {cls}.{method}() call started returning {error} after we merged the {var} migration PR", + "Users are reporting {error} intermittently — the {module} logs show {var} is sometimes null", + "After deploying the {method} change, we see {error} on about 5%% of requests to {cls}", + "The staging environment throws {error} but prod is fine — suspect it's the {var} config difference", + # file path references + "Look at {module}/{cls}.{method}() — the {var} parameter is never validated before being passed to the database layer", + "In the {module} service, the {method}() function at line ~200 has a subtle bug with {var} boundary handling", + "The {cls} constructor in {module} initializes {var} too early — move it to the {method}() call site", + # constraint-carrying + "Add {method}() to {cls} without breaking the existing API contract — we have downstream consumers", + "Optimize {cls}.{method}() for the case where {var} has over 10K entries, but keep the simple path fast too", + "Fix the {error} in {module} — but don't change the public interface, we're in a code freeze for other modules", + "Add telemetry to {cls}.{method}() without adding any new dependencies to the {module} package", + # multi-sentence with background + "We profiled the {method} endpoint and {var} is growing unbounded in {cls}. We need to add eviction or cap the size. The 99th percentile latency spiked 3x last week.", + "The {cls} pool keeps hitting {error} during peak hours. We scaled horizontally but the issue persists. I think {method}() is holding a lock too long.", + "After the last {module} refactor, {cls}.{method}() no longer returns deterministic results. The old tests still pass but the integration tests are flaky. Might be a race condition on {var}.", + "We're moving from REST to gRPC for the {module} service. Start by converting {cls}.{method}() — it's the most latency-sensitive endpoint. Keep the REST handler as a thin adapter for backward compat.", + # review / debugging style + "Can you review the {cls}.{method}() implementation? I think the error handling around {var} is wrong", + "Why does {cls} create a new {var} on every call to {method}()? Seems wasteful", + "Walk me through the {method}() flow in {module} — I need to understand where {var} gets validated", + "Is there a reason {cls}.{method}() catches Exception instead of the specific {error}?", + # infra / DevOps + "Add a Dockerfile for the {module} service that runs {cls} on port 8080 with health checks", + "The k8s deployment for {module} keeps OOMKilling — add memory limits and check if {cls} leaks during {method}()", + "Set up a GitHub Action that runs the {module} tests, lints with ruff, and blocks merge on failure", + "Add Prometheus metrics for {cls}.{method}() — we need p50/p95/p99 latency and error rate by status code", + # data / schema + "Add a new {var} column to the {module} table with a default value and backfill script", + "The {cls} serializer is dropping {var} fields when they're empty lists — should preserve them as []", + "Normalize the {var} schema in {module}: split the nested object into its own table with a foreign key", +) + +# -- Bridge text for multi-turn conversations -- + +_BRIDGE_ANALYZE = ( + "Let me look at the relevant code.", + "I'll start by reading the file to understand the current implementation.", + "Let me search for where this is defined.", + "First, let me check the existing code.", + "Let me examine the implementation.", + "I'll read the source to understand what's happening.", + "Let me look at the file to see the current state.", + "I need to understand the existing logic first.", + "Let me check where {cls} is defined.", + "I'll look at the {method}() implementation first.", + "Let me find all the callers of {method}() so we know the impact.", + "I want to see the full {cls} class before making changes.", +) + +_BRIDGE_FIX = ( + "I can see the issue. Let me fix it.", + "The problem is in the error handling. Here's the fix:", + "This needs to be updated. Let me apply the change.", + "Found it. The logic is incorrect here. Let me correct it.", + "I see the bug. The condition is inverted. Here's the fix:", + "The issue is a missing null check. Let me add it.", + "This needs to be async. Let me update it.", + "The root cause is a race condition on the shared state. Here's a fix:", + "I see the problem -- {var} is being mutated after it's shared. Let me fix it.", + "The issue is that {method}() doesn't account for the empty case. Here's the change:", + "This is a classic off-by-one. Let me correct the boundary check.", + "The lock ordering is wrong here. Let me restructure it.", +) + +_BRIDGE_TEST = ( + "Let me run the tests to verify.", + "Now let me check if the tests pass.", + "Let me verify the fix with the test suite.", + "Running the tests to confirm the change works.", + "Let me make sure nothing else broke.", + "I'll add a test for the new behavior and run the suite.", + "Let me run just the relevant tests first.", + "Let me verify with both unit and integration tests.", +) + +_BRIDGE_EXPLAIN = ( + "Here's what's happening in this code:", + "The flow works like this:", + "This is structured as follows:", + "The key parts are:", + "Let me walk through the logic:", + "The architecture here is layered -- {cls} delegates to {module} for the heavy lifting.", + "There are two paths through this code depending on whether {var} is set.", + "The call chain is: {method}() -> {module}.{method}() -> the underlying store.", +) + +_BRIDGE_SUMMARY = ( + "The fix adds proper error handling for the {var} case.", + "I've updated {cls}.{method}() to handle the edge case.", + "The change ensures {var} is validated before use.", + "This should resolve the {error} issue. The root cause was missing validation on {var}.", + "Done. The {method}() call now correctly handles the {var} boundary condition.", + "Summary: added null check for {var} and updated the return type of {method}().", + "All tests pass. The change is backward-compatible since {method}() still returns the same type.", + "Fixed. The {cls} now properly cleans up {var} on both the happy path and the error path.", + "To summarize: {cls}.{method}() was holding a reference to {var} after the connection closed. " + "The fix moves the cleanup into a finally block.", +) + +_BRIDGE_SECURITY = ( + "This endpoint is vulnerable to SQL injection. The {var} parameter is interpolated directly into the query without sanitization.", + "The JWT validation is missing the audience claim check. An attacker could use a token issued for a different service.", + "Let me check the authentication middleware. The RBAC rules should prevent unauthorized access to {method}().", + "The TLS certificate is using an insecure cipher suite. Let me update the configuration.", + "I see the issue -- the CORS policy allows wildcard origins, which bypasses the CSRF protection.", + "The API key is being logged in plaintext. Let me add a secrets filter to the logging configuration.", + "The password hashing is using MD5. Let me migrate to bcrypt with a proper salt.", + "Let me verify the OAuth2 authorization code flow. The redirect URI validation looks incomplete.", +) + +_BRIDGE_DISTRIBUTED = ( + "The problem is a split-brain scenario. When the network partitions, both nodes think they're the leader.", + "This needs eventual consistency. Let me add a vector clock to track causal ordering of {var} updates.", + "The quorum calculation is wrong -- with 5 nodes you need at least 3 for a write quorum, not 2.", + "Let me add a distributed lock with a TTL to prevent the {method}() race condition across replicas.", + "The gossip protocol is flooding the network. Let me switch to a pull-based protocol with exponential backoff.", + "I see the issue -- the Raft log is not being compacted, so leader election takes increasingly long.", + "The shard rebalancing is not atomic. If it fails midway, some keys become unreachable.", + "Let me add a read-repair mechanism so stale replicas converge after the partition heals.", +) + +_BRIDGE_OBSERVABILITY = ( + "The trace spans are not being propagated across the {module} service boundary. Let me add the OpenTelemetry context injection.", + "I'll add a histogram metric for {method}() latency with buckets at p50/p90/p99 to track the SLO.", + "The structured logs are missing the correlation_id field, making it impossible to trace requests across services.", + "Let me set up a Prometheus alert that fires when the error rate exceeds the SLI threshold for 5 minutes.", + "The dashboard is missing the {cls} service panel. Let me add a Grafana query for the {method} latency distribution.", + "I see the problem -- the span context is being dropped at the async boundary. Let me propagate it through the task.", +) + +_BRIDGE_DATA_ARCHITECTURE = ( + "The EXPLAIN ANALYZE shows a sequential scan on {var} -- we need a composite index on ({var}, {method}).", + "This is a classic N+1 query problem. The ORM is issuing a separate SELECT for each {var} in the loop.", + "Let me batch the {method}() inserts into a single transaction. The current approach holds a lock per row.", + "The connection pool is exhausted because {cls}.{method}() opens a new connection without releasing it on error.", + "I'll denormalize the {var} join to avoid the cross-shard query. The read pattern is 100x more frequent than writes.", + "The transaction isolation level needs to be SERIALIZABLE here to prevent phantom reads on {var}.", + "Let me add a covering index so the query can be satisfied from the index alone without a table lookup.", + "The partition key is wrong -- hashing by {var} creates hot spots because the distribution is skewed.", +) + +_BRIDGE_ARCHITECTURE_TRADEOFF = ( + "There are two approaches here. Option A: add a caching layer in front of {cls}.{method}() with a TTL-based invalidation. " + "This gives us sub-millisecond reads but introduces a consistency window where stale {var} can be returned. " + "Option B: use a write-through cache that invalidates on every {method}() call. This maintains consistency but adds " + "latency to writes and complexity to the error handling path. Given the read-heavy workload (100:1 ratio), " + "I'd recommend Option A with a 30-second TTL and a manual invalidation endpoint for critical updates.", + + "The current architecture has {cls} calling {module} synchronously, which blocks the event loop during {method}(). " + "We could switch to a message queue (Redis Streams or Kafka) to decouple the producer and consumer. " + "The tradeoff is that we lose the synchronous error feedback -- if {method}() fails, the caller won't know until it " + "polls for the result. We'd need to add a dead-letter queue and a retry policy with exponential backoff. " + "For this use case, I think the decoupling is worth it because the {method}() latency varies 10x under load.", + + "Looking at this from a security perspective, the {var} field is user-controlled input that flows through " + "{cls}.{method}() into a SQL query. The ORM provides parameterized queries, so SQL injection isn't a risk, " + "but the {var} value is reflected in error messages which could leak internal table names. Additionally, " + "the rate limiter on this endpoint uses a per-IP strategy, but behind a load balancer all requests share " + "the same source IP. We should switch to a per-API-key rate limit and sanitize error responses.", + + "This is a classic CAP theorem tradeoff. The {module} service currently prioritizes consistency (CP) -- " + "if a network partition occurs, the service rejects writes rather than risk divergent state. For the {cls} " + "use case, availability matters more than strict consistency because {method}() is idempotent and clients " + "already handle retries. I'd recommend switching to an AP model with conflict resolution via last-write-wins " + "using the timestamp from the {var} field. We'd need to add a reconciliation job that runs hourly.", +) + +_BRIDGE_REFACTOR = ( + "Let me extract this into a separate method for clarity.", + "I'll restructure {cls} to separate the {method} concern from the lifecycle logic.", + "The current approach mixes IO with business logic. Let me split them.", + "I'll move {method}() into its own module since it's used across multiple services.", + "Let me introduce an interface so we can swap the {module} implementation later.", + "I'll consolidate the duplicate {method} logic into a shared helper.", + "The {cls} class is doing too much. Let me split it along the {var}/{method} boundary.", +) + +_BRIDGE_PERF = ( + "Let me profile {method}() to see where the time goes.", + "The bottleneck is likely in the {var} allocation. Let me check.", + "I'll add some timing instrumentation first.", + "The issue is that {cls} creates a new {var} on every call. Let me add pooling.", + "Let me check the query plan to see if we're missing an index.", + "This is doing N+1 queries. Let me batch the {method}() calls.", + "The {var} is being serialized on every request. Let me cache it.", + "I see the problem -- {method}() is called inside the lock, blocking all other workers.", +) + +_BRIDGE_DEPLOY = ( + "Let me check the deployment configuration.", + "I'll look at the Dockerfile and the k8s manifests.", + "Let me verify the environment variables are set correctly.", + "I'll check the CI pipeline configuration.", + "Let me look at the health check endpoint.", + "I see the issue in the resource limits. Let me update the deployment.", + "The liveness probe is too aggressive. Let me increase the timeout.", +) + +_BRIDGE_WRITE_TEST = ( + "Let me write tests for the new behavior.", + "I'll add test cases for both the happy path and the error cases.", + "Let me add a parametrized test to cover all the edge cases.", + "I'll write an integration test that exercises the full {method}() flow.", + "Let me add a regression test for this specific bug.", + "I'll mock the {module} dependency so the test is isolated.", + "Here's a test that verifies the fix -- it would have caught the original bug.", +) + +_FOLLOWUP_QUESTIONS = ( + "Can you also add a test for the edge case where {var} is None?", + "What about the {method} path -- does it need the same fix?", + "Should we add logging here too?", + "Can you explain why {cls} uses {var} instead of a local?", + "Is there a performance concern with this approach?", + "Should we also update the {method} docstring?", + "What happens if {var} is empty instead of None?", + "Can you also check if {method}() handles concurrent access correctly?", + "Does this need a database migration?", + "Should we add a feature flag for this change?", + "What about backward compatibility? The old callers pass {var} as a string.", + "Can you check if the {module} service needs the same fix?", + "Is this safe to deploy without a maintenance window?", + "Can you run the integration tests too?", + "Looks good. Can you also update the config to increase the default {var}?", + "One more thing -- can you make {method}() idempotent?", +) + +_LANGUAGES = ("python", "go", "rust", "typescript") + +_TEXT_POOL_BLOCKS = 200 +_BASELINE_POOL_TOKENS = 10_000_000 + +# Block counts per generator, weighted to reflect AI inference server workloads. +# ML/AI content (~12%) reflects the primary use case of benchmarking LLM inference +# servers, where MoE models route tokens based on content domain. Real library +# names (torch, numpy, etc.) activate correct expert pathways. +# ~28% code, ~11% ML/AI code, ~20% bash/output+training logs, ~11% JSON, +# ~9% errors, ~3% SQL, ~10% other (tool use, diffs, CI, config, docs, tests), +# ~8% user prompts (natural language coding requests) +_TOOL_POOL_BLOCK_COUNTS: dict[str, int] = { + # Code (~28%) + "_gen_python_code": 45, + "_gen_go_code": 45, + "_gen_rust_code": 45, + "_gen_typescript_code": 45, + # ML/AI code (~11%) + "_gen_ml_training_code": 30, + "_gen_ml_inference_code": 25, + "_gen_ml_config": 15, + # Bash/output + training logs (~20%) + "_gen_bash_output": 130, + "_gen_ml_training_log": 20, + # JSON (~11%) + "_gen_json_response": 80, + # Errors (~9%) + "_gen_error_traceback": 45, + "_gen_cuda_error": 20, + # SQL (~3%) + "_gen_sql_query": 20, + # User prompts (~6%) + "_gen_user_prompt": 35, + # Tool use / diffs / CI / config / docs / tests (~8%) + "_gen_tool_use_block": 25, + # Multi-turn conversations (~10%) + "_gen_coding_conversation": 90, + "_gen_git_diff": 15, + "_gen_cicd_output": 15, + "_gen_config_file": 15, + "_gen_markdown_doc": 15, + "_gen_test_output": 15, +} +# fmt: on diff --git a/src/aiperf/dataset/generator/_coding_tool.py b/src/aiperf/dataset/generator/_coding_tool.py new file mode 100644 index 0000000000..0819952b84 --- /dev/null +++ b/src/aiperf/dataset/generator/_coding_tool.py @@ -0,0 +1,368 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tool-use and bash-template generators (mixin for CodingContentGenerator). + +Extracted from ``coding_content.py`` to keep that module under the +ergonomics file-size cap. Methods read ``self._template_rng`` and the +shared vocabulary tuples; behavior is unchanged. +""" + +from __future__ import annotations + +from aiperf.dataset.generator._coding_vocab import ( + _CLASSES, + _CLI_COMMANDS, + _ERROR_MESSAGES, + _FILE_PATHS, + _GO_PACKAGES, + _LANG_FILE_PATHS, + _METHODS, + _MODULES, + _RUST_CRATES, + _VARS, +) + + +class _ToolMixin: + def _file_pool(self, language: str | None) -> tuple[str, ...]: + if language: + return _LANG_FILE_PATHS.get(language, _FILE_PATHS) + return _FILE_PATHS + + def _gen_tool_use_block(self, language: str | None = None) -> str: + r = self._template_rng + return r.choice( + [ + lambda: self._gen_tool_read(language=language), + lambda: self._gen_tool_edit(language=language), + lambda: self._gen_tool_search(language=language), + lambda: self._gen_tool_bash(language=language), + ] + )() + + def _gen_tool_read(self, language: str | None = None) -> str: + r = self._template_rng + file_pool = self._file_pool(language) + f = r.choice(file_pool) + start_line = r.randint(1, 200) + cls = r.choice(_CLASSES) + m1, m2 = r.sample(_METHODS, 2) + v1, v2 = r.sample(_VARS, 2) + mod = r.choice(_MODULES) + err = r.choice(_ERROR_MESSAGES) + + lang_lines: dict[str | None, list[str]] = { + "python": [ + f"def {m1}(self, {v1}):", + f"self._{v1} = {v1}", + f"{v2} = {mod}.{m2}({v1})", + f"if {v1} is None:", + f' raise ValueError("{err}")', + f"return {v2}", + f'logger.debug(f"{cls}.{m1}: {{{{{v1}}}}}")', + "", + ], + "go": [ + f"func (s *{cls}) {m1.title()}(ctx context.Context) error {{", + f"s.{v1} = {v1}", + f"{v2}, err := s.{m2.title()}(ctx)", + "if err != nil {", + f'return fmt.Errorf("{err}: %w", err)', + "}", + "return nil", + "", + ], + "rust": [ + f"pub async fn {m1}(&mut self) -> Result<()> {{", + f"let {v1} = self.{v2}.clone();", + f"let {v2} = self.{m2}(&{v1}).await?;", + f"if {v2}.is_empty() {{", + f'anyhow::bail!("{err}");', + "}", + "Ok(())", + "", + ], + "typescript": [ + f"async {m1}({v1}: string): Promise {{", + f"this.{v1} = {v1};", + f"const {v2} = await this.{m2}({v1});", + f"if (!{v2}) {{", + f" throw new Error('{err}');", + "}", + f"console.log(`{cls}.{m1}: ${{{{{v2}}}}}`);", + "", + ], + } + code_lines = lang_lines.get(language, lang_lines["python"]) + + lines = [] + for i in range(start_line, start_line + r.randint(15, 30)): + indent = " " if r.random() > 0.3 else " " + line_content = r.choice(code_lines) + lines.append(f"{i:>6}\t{indent}{line_content}") + + content = "\n".join(lines) + return f"""\ +read +{f} + +{content} + +""" + + def _gen_tool_edit(self, language: str | None = None) -> str: + r = self._template_rng + file_pool = self._file_pool(language) + f = r.choice(file_pool) + m1, m2 = r.sample(_METHODS, 2) + v1, v2 = r.sample(_VARS, 2) + cls = r.choice(_CLASSES) + err = r.choice(_ERROR_MESSAGES) + + edits: dict[str | None, tuple[str, str]] = { + "python": ( + f" def {m1}(self, {v1}):\n return self._{m2}({v1})", + f" async def {m1}(self, {v1}: str) -> dict:\n" + f" try:\n" + f" {v2} = await self._{m2}({v1})\n" + f" if {v2} is None:\n" + f' raise ValueError("{err}")\n' + f' return {{{{"status": "ok", "data": {v2}}}}}\n' + f" except Exception as exc:\n" + f' logger.error("{cls}.{m1} failed: %s", exc)\n' + f" raise", + ), + "go": ( + f"func (s *{cls}) {m1.title()}() error {{{{\n return nil\n}}}}", + f"func (s *{cls}) {m1.title()}(ctx context.Context) error {{{{\n" + f" {v2}, err := s.{m2.title()}(ctx)\n" + f" if err != nil {{{{\n" + f' return fmt.Errorf("{err}: %w", err)\n' + f" }}}}\n" + f" s.{v1} = {v2}\n" + f" return nil\n" + f"}}}}", + ), + "rust": ( + f"fn {m1}(&self) -> Result<()> {{{{\n Ok(())\n}}}}", + f"async fn {m1}(&mut self) -> Result<()> {{{{\n" + f" let {v2} = self.{m2}().await?;\n" + f' anyhow::ensure!(!{v2}.is_empty(), "{err}");\n' + f" self.{v1} = {v2};\n" + f" Ok(())\n" + f"}}}}", + ), + "typescript": ( + f"{m1}({v1}: string) {{{{\n return this.{m2}({v1});\n}}}}", + f"async {m1}({v1}: string): Promise> {{{{\n" + f" const {v2} = await this.{m2}({v1});\n" + f" if (!{v2}) throw new Error('{err}');\n" + f" return {{ status: 'ok', data: {v2} }};\n" + f"}}}}", + ), + } + old_str, new_str = edits.get(language, edits["python"]) + + return f"""\ +edit +{f} +{old_str} +{new_str} + +The file {f} has been updated successfully. + +""" + + def _gen_tool_search(self, language: str | None = None) -> str: + r = self._template_rng + file_pool = self._file_pool(language) + + lang_patterns: dict[str | None, list[str]] = { + "python": [ + f"class {r.choice(_CLASSES)}", + f"def {r.choice(_METHODS)}", + f"import {r.choice(_MODULES)}", + f"async def {r.choice(_METHODS)}", + ], + "go": [ + f"func {r.choice(_METHODS).title()}", + f"type {r.choice(_CLASSES)} struct", + f'"{r.choice(list(_GO_PACKAGES))}"', + f"func New{r.choice(_CLASSES)}", + ], + "rust": [ + f"fn {r.choice(_METHODS)}", + f"pub struct {r.choice(_CLASSES)}", + f"use {r.choice(list(_RUST_CRATES))}", + f"impl {r.choice(_CLASSES)}", + ], + "typescript": [ + f"class {r.choice(_CLASSES)}", + f"export function {r.choice(_METHODS)}", + f"import {{ {r.choice(_CLASSES)} }}", + f"interface {r.choice(_CLASSES)}", + ], + } + patterns = lang_patterns.get(language, lang_patterns["python"]) + pattern = r.choice([*patterns, r.choice(_ERROR_MESSAGES)]) + + files = r.sample(list(file_pool), min(r.randint(3, 6), len(file_pool))) + matches = [] + for f in files: + line_num = r.randint(1, 400) + ctx = r.choice(_VARS) + matches.append(f"{f}:{line_num}: {pattern}({ctx})") + + content = "\n".join(matches) + return f"""\ +search +{pattern} + +{content} + +""" + + def _gen_tool_bash(self, language: str | None = None) -> str: + r = self._template_rng + mod = r.choice(_MODULES) + cls = r.choice(_CLASSES) + methods = r.sample(list(_METHODS), 4) + n_pass = r.randint(10, 80) + n_fail = r.randint(0, 3) + dur = r.uniform(0.5, 30.0) + + lang_cmds: dict[str | None, str] = { + "python": "pytest -xvs tests/", + "go": "go test -v ./...", + "rust": "cargo test", + "typescript": "npx vitest run", + } + cmd = lang_cmds.get(language, r.choice(_CLI_COMMANDS)) + + test_lines = [] + for m in methods: + passed = r.random() > 0.2 + if language == "go": + status = "ok" if passed else "FAIL" + test_lines.append( + f"--- {status}: Test{m.title()} ({r.uniform(0.001, 2.0):.3f}s)" + ) + elif language == "rust": + status = "ok" if passed else "FAILED" + test_lines.append(f"test {mod}::{cls.lower()}::test_{m} ... {status}") + elif language == "typescript": + mark = "\u2713" if passed else "\u2717" + test_lines.append(f" {mark} {cls} > {m} ({r.randint(1, 500)} ms)") + else: + status = "PASSED" if passed else "FAILED" + test_lines.append(f"tests/test_{mod}.py::Test{cls}::test_{m} {status}") + test_output = "\n".join(test_lines) + + return f"""\ +bash +{cmd} + +{test_output} + +{n_pass} passed, {n_fail} failed in {dur:.2f}s + +""" + + def _gen_bash_output(self, language: str | None = None) -> str: + r = self._template_rng + return r.choice( + [ + lambda: self._gen_bash_file_explore(language=language), + lambda: self._gen_bash_build_test(language=language), + lambda: self._gen_bash_git_workflow(language=language), + ] + )() + + def _gen_bash_file_explore(self, language: str | None = None) -> str: + r = self._template_rng + file_pool = self._file_pool(language) + ext_cmds: dict[str | None, tuple[str, str]] = { + "python": ("find . -name '*.py'", "src/**/*.py"), + "go": ("find . -name '*.go'", "**/*.go"), + "rust": ("find . -name '*.rs'", "src/**/*.rs"), + "typescript": ("find . -name '*.ts'", "src/**/*.ts"), + } + find_cmd, glob_pat = ext_cmds.get(language, ext_cmds["python"]) + cmd = r.choice(("ls -la", find_cmd, "tree src/", "wc -l")) + files = r.sample(list(file_pool), min(r.randint(4, 8), len(file_pool))) + file_listing = "\n".join( + f" {f:<42} {r.randint(1, 500):>4} lines {r.randint(1, 50):>3}K" + for f in files + ) + total_lines = r.randint(500, 15000) + + return f"""\ +$ {cmd} +{file_listing} +$ wc -l {glob_pat} | tail -1 + {total_lines} total +$ du -sh . + {r.randint(1, 500)}M\t. +""" + + def _gen_bash_build_test(self, language: str | None = None) -> str: + r = self._template_rng + mod = r.choice(_MODULES) + n_pkgs = r.randint(10, 200) + build_time = r.uniform(0.5, 30.0) + n_pass = r.randint(20, 150) + n_fail = r.randint(0, 5) + test_time = r.uniform(1.0, 60.0) + + lang_build: dict[str | None, tuple[str, str]] = { + "python": ( + "pip install -e '.[dev]'", + f"pytest tests/ -x\n {n_pass} passed, {n_fail} failed in {test_time:.1f}s", + ), + "go": ( + f"go build ./cmd/{mod}\n Compiled {n_pkgs} packages in {build_time:.1f}s", + f"go test -v -race ./...\n {n_pass} passed, {n_fail} failed in {test_time:.1f}s", + ), + "rust": ( + f"cargo build --release\n Compiling {n_pkgs} crates\n Finished in {build_time:.1f}s", + f"cargo test\n {n_pass} passed, {n_fail} failed in {test_time:.1f}s", + ), + "typescript": ( + f"npm ci && npm run build\n Resolved {n_pkgs} packages in {build_time:.1f}s", + f"npx vitest run\n {n_pass} passed, {n_fail} failed in {test_time:.1f}s", + ), + } + build_cmd, test_cmd = lang_build.get(language, lang_build["python"]) + + return f"""\ +$ {build_cmd} +$ {test_cmd} +$ echo $? +{"0" if n_fail == 0 else "1"} +""" + + def _gen_bash_git_workflow(self, language: str | None = None) -> str: + r = self._template_rng + file_pool = self._file_pool(language) + branch = f"{r.choice(_MODULES)}/{r.choice(_METHODS)}-{r.choice(_VARS)}" + mod = r.choice(_MODULES) + files = r.sample(list(file_pool), min(3, len(file_pool))) + changed = "\n".join(f" M {f}" for f in files) + hash1 = f"{r.randint(1000000, 9999999):07x}" + hash2 = f"{r.randint(1000000, 9999999):07x}" + + return f"""\ +$ git checkout -b {branch} +Switched to a new branch '{branch}' +$ git status +On branch {branch} +Changes not staged for commit: +{changed} +$ git add -A && git commit -m "feat: {r.choice(_METHODS)} {r.choice(_VARS)} in {mod}" +[{branch} {hash1}] feat: {r.choice(_METHODS)} {r.choice(_VARS)} in {mod} + {len(files)} files changed, {r.randint(10, 200)} insertions(+), {r.randint(1, 50)} deletions(-) +$ git log --oneline -3 +{hash1} feat: {r.choice(_METHODS)} {r.choice(_VARS)} in {mod} +{hash2} fix: {r.choice(_ERROR_MESSAGES)} +""" diff --git a/src/aiperf/dataset/generator/_coding_tool_long.py b/src/aiperf/dataset/generator/_coding_tool_long.py new file mode 100644 index 0000000000..17aab3b50a --- /dev/null +++ b/src/aiperf/dataset/generator/_coding_tool_long.py @@ -0,0 +1,376 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Long-form tool generators: tool_read_long, tool_bash_verbose, tool_search_verbose. + +Extracted from ``coding_content.py`` to keep that module under the +ergonomics file-size cap. Methods read ``self._template_rng`` and the +shared vocabulary tuples; behavior is unchanged. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from aiperf.dataset.generator._coding_vocab import ( + _CLASSES, + _CLI_COMMANDS, + _ERROR_MESSAGES, + _METHODS, + _MODULES, + _TYPES, + _VARS, +) + + +@dataclass(frozen=True, slots=True) +class _ReadLongCtx: + """Pre-sampled randomness shared across per-language tool-read-long builders.""" + + cls: str + m1: str + m2: str + m3: str + v1: str + v2: str + v3: str + mod: str + err: str + t1: str + t2: str + + +def _read_long_python(c: _ReadLongCtx) -> list[str]: + return [ + f"class {c.cls}:", + f' """{c.cls} handles {c.m1} operations for {c.mod}."""', + "", + f" _default_{c.v3} = 64", + "", + f" def __init__(self, {c.v1}: {c.t1}, {c.v2}: {c.t2} = None):", + f" self._{c.v1} = {c.v1}", + f" self._{c.v2} = {c.v2}", + f" self._{c.v3} = self._default_{c.v3}", + " self._initialized = False", + " self._lock = asyncio.Lock()", + "", + f" async def {c.m1}(self, {c.v1}: {c.t1}) -> {c.t2}:", + " if not self._initialized:", + f' raise RuntimeError("{c.cls} not initialized")', + " async with self._lock:", + f" {c.v2} = await self._{c.m2}({c.v1})", + f" if {c.v2} is None:", + f' raise ValueError("{c.err}")', + f" return {c.v2}", + "", + f" async def _{c.m2}(self, {c.v1}: {c.t1}) -> {c.t2}:", + " try:", + f" {c.v2} = await {c.mod}.{c.m2}({c.v1})", + f' logger.debug(f"{c.cls}.{c.m2}: {{{{{c.v1}}}}}")', + f" return {c.v2}", + " except Exception as e:", + f' logger.error("{c.err}: %s", e)', + f' raise ValueError("{c.err}") from e', + "", + f" async def {c.m3}(self, {c.v1}: {c.t1}, {c.v2}: {c.t2}) -> None:", + f" if {c.v1} is None:", + " return", + f" existing = await self._{c.m2}({c.v1})", + " if existing is not None:", + f" existing.{c.v3} = {c.v2}", + " await existing.save()", + " else:", + f" await {c.mod}.{c.m3}({c.v1}, {c.v2})", + "", + f" def {c.m1}_sync(self) -> None:", + " self._initialized = True", + f" self._{c.v3} = 0", + ] + + +def _read_long_go(c: _ReadLongCtx) -> list[str]: + return [ + f"type {c.cls} struct {{", + f"\t{c.v1} {c.t1}", + f"\t{c.v2} {c.t2}", + "\tmu sync.RWMutex", + "\tlog *zap.Logger", + "}", + "", + f"func New{c.cls}({c.v1} {c.t1}, log *zap.Logger) *{c.cls} {{", + f"\treturn &{c.cls}{{", + f"\t\t{c.v1}: {c.v1},", + "\t\tlog: log,", + "\t}", + "}", + "", + f"func (s *{c.cls}) {c.m1.title()}(ctx context.Context) error {{", + "\ts.mu.Lock()", + "\tdefer s.mu.Unlock()", + "", + f"\t{c.v2}, err := s.{c.m2.title()}(ctx)", + "\tif err != nil {", + f'\t\treturn fmt.Errorf("{c.err}: %w", err)', + "\t}", + f"\ts.{c.v1} = {c.v2}", + "\treturn nil", + "}", + "", + f"func (s *{c.cls}) {c.m2.title()}(ctx context.Context) ({c.t2}, error) {{", + f'\ts.log.Debug("{c.cls}.{c.m2.title()}", zap.String("{c.v1}", s.{c.v1}))', + f"\tresult, err := {c.mod}.{c.m2.title()}(ctx, s.{c.v1})", + "\tif err != nil {", + f'\t\treturn "", fmt.Errorf("{c.err}: %w", err)', + "\t}", + "\treturn result, nil", + "}", + ] + + +def _read_long_rust(c: _ReadLongCtx) -> list[str]: + return [ + f"pub struct {c.cls} {{", + f" {c.v1}: {c.t1},", + f" {c.v2}: Option<{c.t2}>,", + " initialized: bool,", + "}", + "", + f"impl {c.cls} {{", + f" pub fn new({c.v1}: {c.t1}) -> Self {{", + f" Self {{ {c.v1}, {c.v2}: None, initialized: false }}", + " }", + "", + f" pub async fn {c.m1}(&mut self) -> Result<{c.t2}> {{", + f' anyhow::ensure!(self.initialized, "{c.cls} not initialized");', + f" let {c.v2} = self.{c.m2}().await?;", + f" if {c.v2}.is_empty() {{", + f' anyhow::bail!("{c.err}");', + " }", + f" Ok({c.v2})", + " }", + "", + f" async fn {c.m2}(&self) -> Result<{c.t2}> {{", + f" let {c.v2} = {c.mod}::{c.m2}(&self.{c.v1}).await?;", + f' tracing::debug!("{c.cls}.{c.m2}: {{}}", self.{c.v1});', + f" Ok({c.v2})", + " }", + "", + f" pub async fn {c.m3}(&mut self, {c.v1}: {c.t1}) -> Result<()> {{", + f" let existing = self.{c.m2}().await.ok();", + " match existing {", + " Some(val) if !val.is_empty() => {", + f" self.{c.v2} = Some(val);", + " }", + " _ => {", + f" {c.mod}::{c.m3}(&{c.v1}).await?;", + " }", + " }", + " Ok(())", + " }", + "}", + ] + + +def _read_long_typescript(c: _ReadLongCtx) -> list[str]: + return [ + f"export class {c.cls} {{", + f" private {c.v1}: {c.t1};", + f" private {c.v2}: {c.t2} | null = null;", + " private initialized = false;", + "", + f" constructor({c.v1}: {c.t1}) {{", + f" this.{c.v1} = {c.v1};", + " }", + "", + f" async {c.m1}({c.v1}: {c.t1}): Promise<{c.t2}> {{", + " if (!this.initialized) {", + f" throw new Error('{c.cls} not initialized');", + " }", + f" const {c.v2} = await this.{c.m2}({c.v1});", + f" if (!{c.v2}) {{", + f" throw new Error('{c.err}');", + " }", + f" return {c.v2};", + " }", + "", + f" private async {c.m2}({c.v1}: {c.t1}): Promise<{c.t2} | null> {{", + " try {", + f" const {c.v2} = await {c.mod}.{c.m2}({c.v1});", + f" console.debug(`{c.cls}.{c.m2}: ${{{{{c.v1}}}}}`);", + f" return {c.v2};", + " } catch (err) {", + f" console.error('{c.err}:', err);", + " throw err;", + " }", + " }", + "", + f" async {c.m3}({c.v1}: {c.t1}, {c.v2}: {c.t2}): Promise {{", + f" const existing = await this.{c.m2}({c.v1}).catch(() => null);", + " if (existing) {", + f" Object.assign(existing, {{ {c.v3}: {c.v2} }});", + " await existing.save();", + " } else {", + f" await {c.mod}.{c.m3}({c.v1}, {c.v2});", + " }", + " }", + "}", + ] + + +class _ToolLongMixin: + def _gen_tool_read_long(self, language: str | None = None) -> str: + """Like _gen_tool_read but with 40-80 lines for realistic large file reads.""" + r = self._template_rng + file_pool = self._file_pool(language) + f = r.choice(file_pool) + start_line = r.randint(1, 200) + cls = r.choice(_CLASSES) + m1, m2, m3 = r.sample(_METHODS, 3) + v1, v2, v3 = r.sample(_VARS, 3) + mod = r.choice(_MODULES) + err = r.choice(_ERROR_MESSAGES) + t1, t2 = r.sample(_TYPES, 2) + + ctx = _ReadLongCtx( + cls=cls, + m1=m1, + m2=m2, + m3=m3, + v1=v1, + v2=v2, + v3=v3, + mod=mod, + err=err, + t1=t1, + t2=t2, + ) + builder = { + "python": _read_long_python, + "go": _read_long_go, + "rust": _read_long_rust, + "typescript": _read_long_typescript, + }.get(language, _read_long_python) + code_lines = builder(ctx) + + lines = [] + for i, content in enumerate(code_lines, start=start_line): + lines.append(f"{i:>6}\t{content}") + + content = "\n".join(lines) + return f"""\ +read +{f} + +{content} + +""" + + def _gen_tool_bash_verbose(self, language: str | None = None) -> str: + """Like _gen_tool_bash but with longer, more realistic test output.""" + r = self._template_rng + mod = r.choice(_MODULES) + cls = r.choice(_CLASSES) + methods = r.sample(list(_METHODS), r.randint(8, 15)) + n_pass = r.randint(30, 150) + n_fail = r.randint(0, 3) + dur = r.uniform(2.0, 45.0) + + lang_cmds: dict[str | None, str] = { + "python": "pytest -xvs tests/", + "go": "go test -v ./...", + "rust": "cargo test", + "typescript": "npx vitest run", + } + cmd = lang_cmds.get(language, r.choice(_CLI_COMMANDS)) + + test_lines = [] + for m in methods: + passed = r.random() > 0.15 + t = r.uniform(0.001, 3.0) + if language == "go": + status = "ok" if passed else "FAIL" + test_lines.append(f"--- {status}: Test{m.title()} ({t:.3f}s)") + if not passed: + v = r.choice(_VARS) + test_lines.append( + f" {mod}_test.go:{r.randint(20, 300)}: " + f"expected {v} to be non-nil" + ) + elif language == "rust": + status = "ok" if passed else "FAILED" + test_lines.append(f"test {mod}::{cls.lower()}::test_{m} ... {status}") + if not passed: + err = r.choice(_ERROR_MESSAGES) + test_lines.append(f" thread '{m}' panicked at '{err}'") + elif language == "typescript": + mark = "\u2713" if passed else "\u2717" + test_lines.append(f" {mark} {cls} > {m} ({r.randint(1, 800)} ms)") + if not passed: + test_lines.append(" Expected: true\n Received: false") + else: + status = "PASSED" if passed else "FAILED" + test_lines.append(f"tests/test_{mod}.py::Test{cls}::test_{m} {status}") + if not passed: + err = r.choice(_ERROR_MESSAGES) + v = r.choice(_VARS) + test_lines.extend( + [ + f" FAILED tests/test_{mod}.py::Test{cls}::test_{m}", + f" AssertionError: assert {v} == expected", + f" where {v} = {cls}().{m}()", + f" {err}", + ] + ) + test_output = "\n".join(test_lines) + + warnings = "" + if r.random() < 0.4: + w_count = r.randint(1, 5) + warnings = f"\n\n{w_count} warning(s)" + + return f"""\ +bash +{cmd} + +{test_output} +{warnings} +========================= {n_pass} passed, {n_fail} failed in {dur:.2f}s ========================= + +""" + + def _gen_tool_search_verbose(self, language: str | None = None) -> str: + """Like _gen_tool_search but returns many matches across multiple files.""" + r = self._template_rng + file_pool = self._file_pool(language) + pattern = r.choice(_CLASSES) + + files = r.sample(list(file_pool), min(r.randint(6, 12), len(file_pool))) + matches = [] + for f in files: + n_hits = r.randint(1, 4) + for _ in range(n_hits): + line_num = r.randint(1, 500) + v = r.choice(_VARS) + m = r.choice(_METHODS) + ctx = r.choice( + [ + f"class {pattern}({r.choice(_CLASSES)}):", + f" {m} = {pattern}({v})", + f"from {r.choice(_MODULES)} import {pattern}", + f" self._{v} = {pattern}.{m}()", + f" result: {pattern} = await svc.{m}({v})", + f"# TODO: refactor {pattern} to use async", + ] + ) + matches.append(f"{f}:{line_num}:{ctx}") + + content = "\n".join(matches) + return f"""\ +search +{pattern} + +Found {len(matches)} matches in {len(files)} files: + +{content} + +""" diff --git a/src/aiperf/dataset/generator/_coding_typescript.py b/src/aiperf/dataset/generator/_coding_typescript.py new file mode 100644 index 0000000000..923756f474 --- /dev/null +++ b/src/aiperf/dataset/generator/_coding_typescript.py @@ -0,0 +1,221 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""TypeScript code-template generators (mixin for CodingContentGenerator). + +Extracted from ``coding_content.py`` to keep that module under the +ergonomics file-size cap. Methods read ``self._template_rng`` and the +shared vocabulary tuples; behavior is unchanged. +""" + +from __future__ import annotations + +from aiperf.dataset.generator._coding_vocab import ( + _CLASSES, + _ERROR_MESSAGES, + _HTTP_ROUTES, + _METHODS, + _MODULES, + _TS_IMPORTS, + _VARS, +) + + +class _TypeScriptMixin: + def _gen_typescript_code(self) -> str: + return self._template_rng.choice( + [ + self._gen_typescript_class, + self._gen_typescript_http_handler, + self._gen_typescript_types, + self._gen_typescript_test, + ] + )() + + def _gen_typescript_class(self) -> str: + r = self._template_rng + imp = r.choice(_TS_IMPORTS) + imp_cls = r.choice(_CLASSES) + cls = r.choice(_CLASSES) + m1, m2, m3 = r.sample(_METHODS, 3) + v1, v2, v3 = r.sample(_VARS, 3) + err = r.choice(_ERROR_MESSAGES) + + return f"""\ +import {{{{ {imp_cls} }}}} from '{imp}'; + +interface {cls}Config {{{{ + {v1}: string; + {v2}?: number; + timeout: number; +}}}} + +export class {cls} {{{{ + #{v1}: string; + #{v2}: number; + readonly {v3}: string; + + constructor(config: {cls}Config) {{{{ + this.#{v1} = config.{v1}; + this.#{v2} = config.{v2} ?? 0; + this.{v3} = crypto.randomUUID(); + }}}} + + async {m1}({v1}: string): Promise {{{{ + try {{{{ + const {v2} = await this.{m2}({v1}); + console.log(`${{{{this.#{v1}}}}}: ${{{{{v2}}}}}`); + }}}} catch (err) {{{{ + throw new Error(`{err}`); + }}}} + }}}} + + async {m3}(): Promise {{{{ + return this.#{v2} > 0; + }}}} + + private async {m2}({v1}: string): Promise {{{{ + return this.#{v2}; + }}}} +}}}} +""" + + def _gen_typescript_http_handler(self) -> str: + r = self._template_rng + cls = r.choice(_CLASSES) + m1, m2 = r.sample(_METHODS, 2) + v1, v2 = r.sample(_VARS, 2) + route = r.choice(_HTTP_ROUTES) + err = r.choice(_ERROR_MESSAGES) + + return f"""\ +import {{ Hono }} from 'hono'; +import {{ z }} from 'zod'; +import {{ {cls} }} from './{cls.lower()}'; + +const {m1}Schema = z.object({{{{ + {v1}: z.string().min(1).max(256), + {v2}: z.number().int().positive().optional(), +}}}}); + +type {m1.title()}Input = z.infer; + +const app = new Hono(); + +app.post('{route}', async (c) => {{{{ + const body = {m1}Schema.safeParse(await c.req.json()); + if (!body.success) {{{{ + return c.json({{{{ error: body.error.flatten() }}}}, 400); + }}}} + + const svc = new {cls}(); + try {{{{ + const result = await svc.{m1}(body.data.{v1}); + return c.json({{{{ status: 'ok', data: result }}}}, 201); + }}}} catch (err) {{{{ + return c.json({{{{ error: '{err}' }}}}, 500); + }}}} +}}}}); + +app.get('{route}/:id', async (c) => {{{{ + const id = c.req.param('id'); + const svc = new {cls}(); + const item = await svc.{m2}(id); + if (!item) return c.json({{{{ error: 'not found' }}}}, 404); + return c.json({{{{ status: 'ok', data: item }}}}); +}}}}); + +export default app; +""" + + def _gen_typescript_types(self) -> str: + r = self._template_rng + cls = r.choice(_CLASSES) + v1, v2, v3 = r.sample(_VARS, 3) + m1, m2 = r.sample(_METHODS, 2) + err = r.choice(_ERROR_MESSAGES) + + return f"""\ +export type {cls}Status = 'pending' | 'active' | 'failed' | 'completed'; + +export interface {cls}Event {{{{ + kind: '{m1}' | '{m2}' | 'error'; + {v1}: string; + timestamp: number; +}}}} + +export type {m1.title()}Event = Extract<{cls}Event, {{{{ kind: '{m1}' }}}}>; +export type ErrorEvent = Extract<{cls}Event, {{{{ kind: 'error' }}}}>; + +export interface {cls}Config {{{{ + readonly {v1}: string; + readonly {v2}: number; + readonly {v3}?: Record; +}}}} + +export type Partial{cls} = Partial<{cls}Config> & Pick<{cls}Config, '{v1}'>; + +export function is{cls}Event(e: unknown): e is {cls}Event {{{{ + return ( + typeof e === 'object' && + e !== null && + 'kind' in e && + typeof (e as {cls}Event).{v1} === 'string' + ); +}}}} + +export function assert{cls}Status(s: string): asserts s is {cls}Status {{{{ + const valid: {cls}Status[] = ['pending', 'active', 'failed', 'completed']; + if (!valid.includes(s as {cls}Status)) {{{{ + throw new Error(`{err}: ${{{{s}}}}`); + }}}} +}}}} +""" + + def _gen_typescript_test(self) -> str: + r = self._template_rng + cls = r.choice(_CLASSES) + m1, m2, m3 = r.sample(_METHODS, 3) + v1, v2 = r.sample(_VARS, 2) + err = r.choice(_ERROR_MESSAGES) + mod = r.choice(_MODULES) + + return f"""\ +import {{ describe, it, expect, beforeEach, vi }} from 'vitest'; +import {{ {cls} }} from '../{mod}'; + +describe('{cls}', () => {{{{ + let instance: {cls}; + + beforeEach(() => {{{{ + instance = new {cls}({{{{ {v1}: 'test', timeout: 5000 }}}}); + vi.clearAllMocks(); + }}}}); + + describe('{m1}', () => {{{{ + it('should return expected value', async () => {{{{ + const result = await instance.{m1}('{v2}'); + expect(result).toBeDefined(); + expect(typeof result).toBe('object'); + }}}}); + + it('should throw on invalid input', async () => {{{{ + await expect(instance.{m1}('')).rejects.toThrow('{err}'); + }}}}); + }}}}); + + describe('{m2}', () => {{{{ + it('should call dependency', async () => {{{{ + const spy = vi.spyOn(instance as any, '{m3}'); + await instance.{m2}('{v1}'); + expect(spy).toHaveBeenCalledOnce(); + }}}}); + }}}}); + + it('should handle concurrent calls', async () => {{{{ + const promises = Array.from({{{{ length: 5 }}}}, () => instance.{m1}('{v1}')); + const results = await Promise.all(promises); + expect(results).toHaveLength(5); + }}}}); +}}}}); +""" diff --git a/src/aiperf/dataset/generator/_coding_vocab.py b/src/aiperf/dataset/generator/_coding_vocab.py new file mode 100644 index 0000000000..12c1a2bbd3 --- /dev/null +++ b/src/aiperf/dataset/generator/_coding_vocab.py @@ -0,0 +1,345 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Core vocabulary tuples for CodingContentGenerator. + +Extracted from ``coding_content.py`` to stay under the file-size cap. +Contains identifier vocabularies (modules, classes, methods, types, +vars, file paths, routes, error messages, language imports, ML/CUDA +content). Long natural-language tuples live in ``_coding_text.py``. +""" + +from __future__ import annotations + +# fmt: off +# -- Vocabulary tuples for template fills -- + +_MODULES = ( + "auth", "cache", "config", "database", "events", "handler", "logger", + "metrics", "middleware", "pipeline", "processor", "registry", "router", + "scheduler", "serializer", "service", "storage", "transport", "validator", + "worker", "adapter", "broker", "collector", "dispatcher", "encoder", + "factory", "gateway", "indexer", "manager", "monitor", "notifier", + "observer", "parser", "provider", "queue", "resolver", "scanner", + "session", "sink", "source", "stream", "transformer", "uploader", + # web / HTTP + "api", "webhook", "cors", "oauth", "graphql", "grpc", "websocket", + "rate_limiter", "proxy", "load_balancer", "reverse_proxy", + # database / data + "migration", "schema", "repository", "connection_pool", "query_builder", + "data_loader", "orm", "replication", "sharding", "backup", + # ML / data science + "inference", "tokenizer", "embedding", "feature_store", "model_registry", + "trainer", "evaluator", "dataset", "sampler", "checkpoint", + # DevOps / infra + "deployer", "provisioner", "orchestrator", "health_check", "autoscaler", + "dns_resolver", "cert_manager", "secret_store", "telemetry", "alerter", + # security + "firewall", "encryptor", "key_manager", "audit", "compliance", + # real libraries / frameworks + "torch", "numpy", "pandas", "sqlalchemy", "fastapi", "pydantic", + "celery", "redis", "boto3", "transformers", "datasets", "accelerate", + "flask", "django", "requests", +) + +_CLASSES = ( + "RequestHandler", "DataProcessor", "EventEmitter", "CacheManager", + "ConnectionPool", "TaskScheduler", "MessageBroker", "StateManager", + "ConfigLoader", "MetricsCollector", "RateLimiter", "CircuitBreaker", + "RetryPolicy", "BatchProcessor", "StreamReader", "TokenValidator", + "SessionStore", "PermissionChecker", "ResourceAllocator", "HealthMonitor", + "LoadBalancer", "QueueConsumer", "IndexBuilder", "SchemaValidator", + "PipelineStage", "WorkerPool", "ContextManager", "PluginLoader", + "TemplateEngine", "SignalHandler", "ProtocolAdapter", "BufferManager", + "ThrottleController", "RegistryClient", "LockManager", "SnapshotStore", + "AuditLogger", "FeatureToggle", "MigrationRunner", "DeploymentManager", + # HTTP layer + "HttpClient", "RouteResolver", "CorsMiddleware", "AuthMiddleware", + "ResponseSerializer", "RequestParser", "WebSocketManager", "ApiGateway", + # data layer + "QueryExecutor", "TransactionManager", "MigrationEngine", "PoolManager", + "ReplicaSelector", "ShardRouter", "CursorIterator", "ChangeStream", + # error types + "RetryableError", "ValidationError", "TimeoutError", "QuotaExceeded", + "ConflictError", "NotFoundError", "AuthorizationError", "RateLimitError", + # ML / inference + "ModelLoader", "TokenEncoder", "EmbeddingStore", "FeatureExtractor", + "InferenceEngine", "BatchScheduler", "GradientAccumulator", "Checkpoint", + # infra / orchestration + "ServiceMesh", "HealthProbe", "AutoScaler", "SecretProvider", + "CertRotator", "DnsCache", "TelemetryExporter", "AlertDispatcher", + # real framework classes + "Tensor", "DataFrame", "Series", "Session", "Engine", "Router", + "Pipeline", "Trainer", "Dataset", "DataLoader", "Optimizer", + "Tokenizer", +) + +_METHODS = ( + "process", "handle", "validate", "transform", "execute", "initialize", + "configure", "dispatch", "resolve", "serialize", "deserialize", "encode", + "decode", "publish", "subscribe", "notify", "aggregate", "partition", + "schedule", "allocate", "release", "acquire", "flush", "compress", + "decompress", "authenticate", "authorize", "revoke", "checkpoint", + "rollback", "migrate", "replicate", "synchronize", "reconcile", + "invalidate", "prefetch", "evict", "rebalance", "throttle", "retry", + "render", "persist", "hydrate", "prune", "drain", "backfill", + "enqueue", "dequeue", "broadcast", "handshake", "negotiate", "probe", + "rotate", "shard", "merge", "split", "compact", "snapshot", + "finalize", "abort", "resume", "suspend", "escalate", "demote", + "promote", "quarantine", "scrub", "warm_up", "cool_down", "heal", + "reclaim", "tombstone", "seal", "unseal", "bootstrap", "teardown", + # real library methods + "forward", "backward", "train", "evaluate", "predict", "fit", + "load_state_dict", "save_pretrained", "from_pretrained", "to_dict", +) + +_TYPES = ( + "str", "int", "float", "bool", "bytes", "dict", "list", "tuple", "set", + "None", "Any", "Optional", "Sequence", "Mapping", "Iterator", "Callable", + "Awaitable", "Coroutine", "AsyncIterator", "Generator", "TypeVar", + "Protocol", "ClassVar", "Final", "Literal", "Union", "Type", + "NamedTuple", "TypedDict", "Annotated", "ParamSpec", "Self", +) + +_VARS = ( + "result", "data", "config", "context", "payload", "response", "request", + "buffer", "cursor", "offset", "count", "total", "index", "batch", + "chunk", "token", "record", "entry", "item", "value", "key", "state", + "status", "event", "message", "signal", "metric", "timestamp", "duration", + "timeout", "retries", "threshold", "capacity", "interval", "priority", + "sequence", "channel", "endpoint", "header", "session", "connection", + "pipeline", "schema", "trace_id", "tenant_id", "batch_size", "page_size", + "shard_key", "replica_id", "worker_id", "partition_key", "ttl", + "max_retries", "backoff", "jitter", "watermark", "checkpoint_id", + "correlation_id", "span_id", "parent_id", "depth", "fanout", + "concurrency", "rate", "window", "lag", "drift", "skew", + "epoch", "generation", "version", "revision", "digest", "nonce", +) + +_FILE_PATHS = ( + "src/main.py", "src/config.py", "src/models.py", "src/routes.py", + "src/utils.py", "src/middleware.py", "src/database.py", "src/auth.py", + "tests/test_main.py", "tests/test_models.py", "tests/conftest.py", + "lib/core.go", "lib/handler.go", "lib/service.go", "lib/types.go", + "pkg/api/server.go", "pkg/api/client.go", "pkg/store/store.go", + "cmd/server/main.go", "internal/config/config.go", + "src/lib.rs", "src/main.rs", "src/config.rs", "src/error.rs", + "src/handler.rs", "src/models.rs", "src/routes.rs", + "src/index.ts", "src/app.ts", "src/types.ts", "src/api.ts", + "src/components/App.tsx", "src/components/Form.tsx", + "Dockerfile", "Makefile", "docker-compose.yml", "pyproject.toml", + ".github/workflows/ci.yml", "kubernetes/deployment.yaml", +) + +_HTTP_ROUTES = ( + "/api/v1/users", "/api/v1/items", "/api/v1/orders", "/api/v1/auth/login", + "/api/v1/auth/refresh", "/api/v2/search", "/api/v2/analytics", + "/health", "/ready", "/metrics", "/api/v1/webhooks", "/api/v1/uploads", + "/api/v1/notifications", "/api/v1/settings", "/api/v1/billing", + "/api/v1/teams/{team_id}/members", "/api/v1/projects/{project_id}/runs", + "/api/v1/tenants/{tenant_id}/quota", "/internal/gc", "/internal/debug/pprof", +) + +_DB_TABLES = ( + "users", "orders", "items", "sessions", "audit_log", "migrations", + "api_keys", "rate_limits", "notifications", "webhooks", "tenants", + "permissions", "invitations", "uploads", "billing_events", + "job_queue", "dead_letter", "feature_flags", "schema_versions", "locks", +) + +_STATUS_CODES = ( + "200 OK", "201 Created", "204 No Content", "301 Moved Permanently", + "400 Bad Request", "401 Unauthorized", "403 Forbidden", "404 Not Found", + "409 Conflict", "429 Too Many Requests", "500 Internal Server Error", + "502 Bad Gateway", "503 Service Unavailable", "504 Gateway Timeout", +) + +_LANG_FILE_PATHS: dict[str, tuple[str, ...]] = { + "python": ( + "src/main.py", "src/config.py", "src/models.py", "src/routes.py", + "src/utils.py", "src/middleware.py", "src/database.py", "src/auth.py", + "tests/test_main.py", "tests/test_models.py", "tests/conftest.py", + "pyproject.toml", "Dockerfile", "Makefile", + "src/api/v1/endpoints.py", "src/api/v1/schemas.py", "src/api/deps.py", + "src/core/security.py", "src/core/events.py", "src/services/worker.py", + "src/repositories/base.py", "tests/integration/test_api.py", + ), + "go": ( + "lib/core.go", "lib/handler.go", "lib/service.go", "lib/types.go", + "pkg/api/server.go", "pkg/api/client.go", "pkg/store/store.go", + "cmd/server/main.go", "internal/config/config.go", + "go.mod", "go.sum", "Makefile", + "internal/middleware/auth.go", "internal/middleware/ratelimit.go", + "internal/repository/postgres.go", "internal/service/worker.go", + "pkg/api/middleware.go", "pkg/api/routes.go", + "internal/telemetry/tracing.go", "internal/health/probe.go", + ), + "rust": ( + "src/lib.rs", "src/main.rs", "src/config.rs", "src/error.rs", + "src/handler.rs", "src/models.rs", "src/routes.rs", + "Cargo.toml", "Cargo.lock", + "src/middleware/auth.rs", "src/middleware/tracing.rs", + "src/repository/mod.rs", "src/repository/postgres.rs", + "src/service/mod.rs", "src/service/worker.rs", + "tests/integration/api_test.rs", "benches/throughput.rs", + ), + "typescript": ( + "src/index.ts", "src/app.ts", "src/types.ts", "src/api.ts", + "src/components/App.tsx", "src/components/Form.tsx", + "src/utils.ts", "src/middleware.ts", "src/routes.ts", + "package.json", "tsconfig.json", "Dockerfile", + "src/services/auth.service.ts", "src/services/worker.service.ts", + "src/middleware/rate-limiter.ts", "src/middleware/error-handler.ts", + "src/models/user.model.ts", "src/models/order.model.ts", + "src/repositories/base.repository.ts", "tests/integration/api.test.ts", + ), +} + +_ERROR_MESSAGES = ( + "connection refused", "timeout exceeded", "permission denied", + "resource not found", "invalid argument", "out of memory", + "deadlock detected", "rate limit exceeded", "authentication failed", + "schema validation error", "serialization error", "buffer overflow", + "index out of range", "null pointer dereference", "type mismatch", + "missing required field", "duplicate key", "constraint violation", + "circular dependency detected", "maximum recursion depth exceeded", + "transaction aborted", "lock timeout after 30s", "quota exceeded", + "connection pool exhausted", "certificate expired", "DNS resolution failed", + "checksum mismatch", "payload too large", "stale read", + "leader election in progress", "shard unavailable", "replica lag exceeded", + "write conflict detected", "token revoked", "session expired", + "circuit breaker open", "backpressure applied", "partition offline", + "consensus timeout", "snapshot corrupted", "migration in progress", + # GPU / infra errors + "CUDA out of memory", "NCCL timeout", "connection reset by peer", + "relation does not exist", "broken pipe", "no route to host", + "too many open files", "disk quota exceeded", +) + +_CLI_COMMANDS = ( + "git status", "git diff HEAD~1", "git log --oneline -10", + "docker build -t app .", "docker compose up -d", + "kubectl get pods -n default", "kubectl apply -f deployment.yaml", + "cargo build --release", "cargo test -- --nocapture", + "go build ./...", "go test -v ./...", "go vet ./...", + "npm run build", "npm test", "npx tsc --noEmit", + "pytest -xvs tests/", "ruff check .", "mypy src/", + "make build", "make test", "make lint", + "curl -s http://localhost:8080/health", + "ps aux | grep python", "top -bn1 | head -20", + # k8s / infra + "kubectl describe pod app-7d4b8f-xz9k", "kubectl logs -f deploy/api --tail=100", + "kubectl rollout status deploy/worker", "kubectl top nodes", + "helm upgrade --install app ./chart -f values.yaml", + "terraform plan -out=tfplan", "terraform apply tfplan", + # redis / data stores + "redis-cli INFO memory", "redis-cli --latency-history -i 1", + "pg_dump -Fc mydb > backup.dump", "mongosh --eval 'db.stats()'", + # perf / profiling + "perf stat -e cache-misses,cache-references ./bin/server", + "strace -c -p $(pgrep server)", "valgrind --tool=memcheck ./bin/app", + "pprof -http=:6060 http://localhost:6060/debug/pprof/heap", + # load testing + "wrk -t12 -c400 -d30s http://localhost:8080/api/v1/items", + "hey -n 10000 -c 100 http://localhost:8080/health", + "ab -n 5000 -c 50 http://localhost:8080/", + # misc dev + "find . -name '*.py' | xargs wc -l | tail -1", + "du -sh node_modules/ target/ dist/", + "lsof -i :8080", "ss -tlnp | grep 8080", + "journalctl -u myapp --since '1 hour ago'", +) + +_GO_PACKAGES = ( + "fmt", "os", "io", "net", "http", "context", "sync", "time", + "strings", "strconv", "encoding/json", "log", "errors", "math", + "sort", "bytes", "crypto", "regexp", "path/filepath", "database/sql", + # popular third-party packages + "github.com/gin-gonic/gin", "go.uber.org/zap", + "github.com/spf13/viper", "github.com/spf13/cobra", + "gorm.io/gorm", "google.golang.org/grpc", + "github.com/prometheus/client_golang/prometheus", + "github.com/redis/go-redis/v9", "github.com/nats-io/nats.go", + "github.com/jackc/pgx/v5", +) + +_RUST_CRATES = ( + "std::io", "std::fs", "std::collections", "std::sync", "std::fmt", + "serde", "serde_json", "tokio", "anyhow", "thiserror", "tracing", + "clap", "reqwest", "axum", "sqlx", "uuid", "chrono", "regex", + # additional popular crates + "tower", "hyper", "diesel", "sea_orm", "tonic", "prost", + "async_trait", "futures", +) + +_TS_IMPORTS = ( + "express", "axios", "lodash", "zod", "prisma", "next", + "react", "react-dom", "typescript", "jest", "vitest", + "node:fs", "node:path", "node:http", "node:crypto", + # additional popular packages + "@nestjs/common", "typeorm", "drizzle-orm", "bullmq", + "@trpc/server", "ioredis", "pg", "knex", +) + +_DECORATORS = ( + "@staticmethod", "@classmethod", "@property", "@abstractmethod", + "@override", "@cached_property", "@dataclass", "@lru_cache", + "@pytest.mark.asyncio", "@pytest.mark.parametrize", + "@app.route", "@app.get", "@app.post", "@router.get", + # ML framework decorators + "@torch.no_grad()", "@torch.inference_mode()", "@torch.compile", + "@torch.jit.script", "@torch.cuda.amp.autocast", +) + +_ML_IMPORTS = ( + "torch", "torch.nn", "torch.optim", "torch.utils.data", + "torch.cuda", "torch.distributed", "torch.amp", + "transformers", "datasets", "accelerate", "peft", + "numpy", "safetensors", "wandb", "tensorboard", + "deepspeed", "bitsandbytes", "trl", "vllm", "triton", +) + +_ML_CLASSES = ( + "Linear", "Conv2d", "MultiheadAttention", "LayerNorm", "Embedding", + "CrossEntropyLoss", "AdamW", "CosineAnnealingLR", "DataLoader", + "DistributedDataParallel", "AutoModelForCausalLM", "AutoTokenizer", + "TrainingArguments", "Trainer", "GenerationConfig", + "BitsAndBytesConfig", "LoraConfig", "PeftModel", + "StoppingCriteria", "LogitsProcessor", +) + +_ML_METHODS = ( + "forward", "backward", "zero_grad", "step", "state_dict", + "load_state_dict", "save_pretrained", "from_pretrained", + "generate", "encode", "decode", "batch_decode", + "to", "cuda", "cpu", +) + +_ML_VARS = ( + "logits", "hidden_states", "attention_mask", "input_ids", + "labels", "loss", "grad_norm", "learning_rate", "num_epochs", + "batch_size", "max_length", "temperature", "top_p", "top_k", + "model_name", +) + +_MODEL_NAMES = ( + "meta-llama/Llama-3.1-8B", "meta-llama/Llama-3.1-70B", + "mistralai/Mixtral-8x7B-v0.1", "mistralai/Mistral-7B-v0.1", + "google/gemma-2-9b", "Qwen/Qwen2.5-72B", + "nvidia/Llama-3.1-Nemotron-70B-Instruct", + "deepseek-ai/DeepSeek-V3", "microsoft/phi-4", +) + +_CUDA_ERRORS = ( + "CUDA out of memory. Tried to allocate 2.00 GiB", + "RuntimeError: Expected all tensors to be on the same device", + "torch.cuda.OutOfMemoryError: CUDA out of memory", + "NCCL error: unhandled system error, NCCL version 2.18.5", + "RuntimeError: NCCL communicator was aborted on rank 0", + "RuntimeError: cuDNN error: CUDNN_STATUS_NOT_SUPPORTED", + "RuntimeError: FlashAttention only supports Ampere GPUs or newer", + "torch.distributed.DistBackendError: NCCL error", + "RuntimeError: Deterministic behavior was enabled", + "CUDA error: device-side assert triggered", +) +# fmt: on diff --git a/src/aiperf/dataset/generator/coding_content.py b/src/aiperf/dataset/generator/coding_content.py new file mode 100644 index 0000000000..06f7f4f2b4 --- /dev/null +++ b/src/aiperf/dataset/generator/coding_content.py @@ -0,0 +1,243 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Coding content generator for realistic coding trace replay. + +Generates structurally plausible coding content (code, bash output, JSON, +errors, git diffs, CI output, configs, markdown, test output, user prompts) +using template-based generation with random identifiers. + +Unlike PromptGenerator which uses Shakespeare as its corpus, this generator +builds two token pools from structural templates: +- text_pool: user prompts (natural language coding requests) +- tool_pool: mixed technical content (code, errors, diffs, configs, etc.) + +Generation uses window slicing from pre-built token pools, same as PromptGenerator. +""" + +from __future__ import annotations + +from aiperf.common import random_generator as rng +from aiperf.common.exceptions import ConfigurationError, NotInitializedError +from aiperf.common.hash_id_random_generator import HashIdRandomGenerator +from aiperf.common.tokenizer import Tokenizer +from aiperf.config import PromptConfig +from aiperf.config.dataset.defaults import InputTokensDefaults +from aiperf.dataset.generator._coding_cicd_docs import _CicdDocsMixin +from aiperf.dataset.generator._coding_conversations import _ConversationsMixin +from aiperf.dataset.generator._coding_conversations_advanced import ( + _ConversationsAdvancedMixin, +) +from aiperf.dataset.generator._coding_errors_diff import _ErrorsDiffMixin +from aiperf.dataset.generator._coding_go import _GoMixin +from aiperf.dataset.generator._coding_json import _JsonMixin +from aiperf.dataset.generator._coding_ml import _MlMixin +from aiperf.dataset.generator._coding_prompts_conv import _PromptsConvMixin +from aiperf.dataset.generator._coding_python import _PythonMixin +from aiperf.dataset.generator._coding_rust import _RustMixin +from aiperf.dataset.generator._coding_sql import _SqlMixin +from aiperf.dataset.generator._coding_text import ( + _BASELINE_POOL_TOKENS, + _TEXT_POOL_BLOCKS, + _TOOL_POOL_BLOCK_COUNTS, +) +from aiperf.dataset.generator._coding_tool import _ToolMixin +from aiperf.dataset.generator._coding_tool_long import _ToolLongMixin +from aiperf.dataset.generator._coding_typescript import _TypeScriptMixin +from aiperf.dataset.generator.base import BaseGenerator +from aiperf.dataset.generator.prompt import sample_tokens_from_corpus + + +class CodingContentGenerator( + _PythonMixin, + _GoMixin, + _RustMixin, + _TypeScriptMixin, + _ToolMixin, + _JsonMixin, + _ErrorsDiffMixin, + _CicdDocsMixin, + _MlMixin, + _SqlMixin, + _PromptsConvMixin, + _ToolLongMixin, + _ConversationsMixin, + _ConversationsAdvancedMixin, + BaseGenerator, +): + """Generator for structurally plausible coding content. + + Builds two pre-tokenized pools from template-based content: + - text_pool: natural language coding requests (built lazily on first use) + - tool_pool: mixed technical content - code, errors, diffs, etc. + (~500K tokens, built eagerly at init) + + Exposes the PromptGenerator-compatible interface (``generate`` / + ``generate_prompt`` / ``calculate_num_tokens``); prompt sampling draws + from the tool pool. + """ + + def __init__( + self, + config: PromptConfig, + tokenizer: Tokenizer, + pool_tokens_target: int | None = None, + **kwargs, + ): + self.config = config + self.tokenizer = tokenizer + self._pool_scale = max( + 1.0, (pool_tokens_target or _BASELINE_POOL_TOKENS) / _BASELINE_POOL_TOKENS + ) + + self._template_rng = rng.derive("dataset.coding_content.template") + self._corpus_rng = rng.derive("dataset.coding_content.corpus") + self._length_rng = rng.derive("dataset.coding_content.length") + + # Hash-ID-based RNG for deterministic per-hash_id generation in + # _build_token_sequence; also read directly by the graph adapters' + # CorpusContentSynthesizer (graph/adapters/shared/content.py), which + # reseeds it per (trace_id, hash_id) for byte-deterministic blocks. + self._hash_id_corpus_rng = HashIdRandomGenerator.from_base_rng(self._corpus_rng) + + super().__init__(config=config, tokenizer=tokenizer, **kwargs) + + self._text_pool: list[int] | None = None + self._tool_pool: list[int] = [] + self._cache: dict[int, list[int]] = {} + self._decoded_cache: dict[tuple[tuple[int, ...], int, int], str] = {} + # No stable terminator probe for the coding corpus; the empty list + # means "append no terminator". Nothing consumes this attribute on + # this branch. + self._bpe_stable_terminator_tokens: list[int] = [] + + self._build_tool_pool() + + # PromptGenerator-parity surface read by the graph adapters' + # CorpusContentSynthesizer (corpus hoisting / shared-memory attach). + self._tokenized_corpus = self._tool_pool + self._corpus_size = len(self._tool_pool) + + def generate( + self, + mean: int | None = None, + stddev: int | None = None, + hash_ids: list[int] | None = None, + block_size: int | None = None, + ) -> str: + if hash_ids: + if mean is None: + raise ValueError("mean must be provided when hash_ids is set.") + # config.block_size defaults to None; fall back to the same + # default PromptGenerator.generate uses so hash_ids work with a + # default config instead of raising TypeError. + bs = block_size or self.config.block_size or InputTokensDefaults.BLOCK_SIZE + return self._generate_cached_prompt(mean, hash_ids, bs) + num_tokens = self.calculate_num_tokens(mean, stddev) + return self.generate_prompt(num_tokens) + + def generate_prompt(self, num_tokens: int) -> str: + tokens = self._sample_tokens(num_tokens, self._tool_pool) + return self.tokenizer.decode(tokens) + + def calculate_num_tokens( + self, + mean: int | None = None, + stddev: int | None = None, + ) -> int: + return self._length_rng.sample_positive_normal_integer(mean, stddev) + + def _ensure_text_pool(self) -> list[int]: + if self._text_pool is None: + self._build_text_pool() + assert self._text_pool is not None + return self._text_pool + + def _build_text_pool(self) -> None: + blocks: list[str] = [] + for _ in range(int(_TEXT_POOL_BLOCKS * self._pool_scale)): + blocks.append(self._gen_user_prompt()) + text = "\n\n".join(blocks) + self._text_pool = self.tokenizer.encode(text) + pool = self._text_pool + self.debug( + lambda: f"Built text pool with {len(pool)} tokens from {len(blocks)} blocks" + ) + + def _build_tool_pool(self) -> None: + blocks: list[str] = [] + for gen_name, count in _TOOL_POOL_BLOCK_COUNTS.items(): + gen_fn = getattr(self, gen_name) + for _ in range(int(count * self._pool_scale)): + blocks.append(gen_fn()) + self._template_rng.shuffle(blocks) + text = "\n\n".join(blocks) + self._tool_pool = self.tokenizer.encode(text) + self.debug( + lambda: ( + f"Built tool pool with {len(self._tool_pool)} tokens " + f"from {len(blocks)} blocks" + ) + ) + + def _sample_tokens(self, num_tokens: int, pool: list[int]) -> list[int]: + if not pool: + raise NotInitializedError("Token pool is not initialized.") + pool_size = len(pool) + if num_tokens <= 0: + return [] + start_idx = self._corpus_rng.randrange(pool_size) + end_idx = start_idx + num_tokens + tokens = pool[start_idx:end_idx] + if end_idx > pool_size: + tokens += pool[: end_idx - pool_size] + return tokens + + def _generate_cached_prompt( + self, + num_tokens: int, + hash_ids: list[int], + block_size: int, + ) -> str: + cache_key = (tuple(hash_ids), num_tokens, block_size) + if cache_key in self._decoded_cache: + return self._decoded_cache[cache_key] + + final_prompt = self._build_token_sequence(num_tokens, hash_ids, block_size) + decoded = self.tokenizer.decode(final_prompt, skip_special_tokens=False) + self._decoded_cache[cache_key] = decoded + return decoded + + def _build_token_sequence( + self, + num_tokens: int, + hash_ids: list[int], + block_size: int, + ) -> list[int]: + final_prompt: list[int] = [] + current_block_size = block_size + + final_block_size = num_tokens - ((len(hash_ids) - 1) * block_size) + if final_block_size <= 0 or block_size < final_block_size: + raise ConfigurationError( + f"Input length: {num_tokens}, Hash IDs: {hash_ids}, Block size: {block_size} " + f"are not compatible. The final hash block size: {final_block_size} must be " + f"greater than 0 and less than or equal to {block_size}." + ) + + for index, hash_id in enumerate(hash_ids): + if index == len(hash_ids) - 1: + current_block_size = final_block_size + + if hash_id not in self._cache: + self._hash_id_corpus_rng.reseed_for_hash_id(hash_id) + self._cache[hash_id] = sample_tokens_from_corpus( + self._tool_pool, + current_block_size, + self._hash_id_corpus_rng, + self.tokenizer.block_separation_token_id, + ) + + final_prompt.extend(self._cache[hash_id]) + + return final_prompt diff --git a/src/aiperf/dataset/generator/image.py b/src/aiperf/dataset/generator/image.py index 40b2115a19..8dcdabeb09 100644 --- a/src/aiperf/dataset/generator/image.py +++ b/src/aiperf/dataset/generator/image.py @@ -116,7 +116,9 @@ def generate(self, *args, **kwargs) -> str: image = self._create_source_image(width, height) self.debug( - lambda: f"Generated image from {self.config.source} with width={width}, height={height}" + lambda: ( + f"Generated image from {self.config.source} with width={width}, height={height}" + ) ) base64_image = utils.encode_image(image, image_format) return f"data:image/{image_format.name.lower()};base64,{base64_image}" diff --git a/src/aiperf/dataset/generator/prompt.py b/src/aiperf/dataset/generator/prompt.py index f375817bf5..c24d8ffc9f 100644 --- a/src/aiperf/dataset/generator/prompt.py +++ b/src/aiperf/dataset/generator/prompt.py @@ -19,6 +19,48 @@ DEFAULT_CORPUS_FILE = "assets/shakespeare.txt" +def sample_tokens_from_corpus( + corpus, + num_tokens: int, + rng_to_use, + sep_token: int | None = None, +) -> list[int]: + """Sample ``num_tokens`` tokens from ``corpus`` using ``rng_to_use``. + + Mirrors :meth:`PromptGenerator._sample_tokens` but takes the RNG explicitly, + so callers (worker processes, the hash-id reseed path) can drive sampling + without sharing PromptGenerator state. Wraps the slice if it overflows the + corpus end so the result is always exactly ``num_tokens`` long. + + Args: + corpus: Token corpus as a sequence of token IDs. + num_tokens: Number of tokens to sample. + rng_to_use: Random generator with a ``randrange(n)`` method. + sep_token: Optional block-separation token prepended at index 0 + (consumes one slot of ``num_tokens``). + + Returns: + List of sampled token IDs of length ``num_tokens``. + """ + corpus_len = len(corpus) + tokens: list[int] = [] + + if sep_token is not None: + tokens.append(sep_token) + num_tokens -= 1 + + start = rng_to_use.randrange(corpus_len) + end = start + num_tokens + + if end <= corpus_len: + tokens.extend(corpus[start:end]) + else: + tokens.extend(corpus[start:]) + tokens.extend(corpus[: end - corpus_len]) + + return tokens + + class PromptGenerator(BaseGenerator): """A class for generating synthetic prompts from a text corpus. @@ -148,8 +190,10 @@ def tokenize_chunk(chunk): ] self._corpus_size = len(self._tokenized_corpus) self.debug( - lambda: f"Initialized corpus with {self._corpus_size} tokens " - f"from {len(chunks)} chunks using {num_threads} thread(s)" + lambda: ( + f"Initialized corpus with {self._corpus_size} tokens " + f"from {len(chunks)} chunks using {num_threads} thread(s)" + ) ) def _create_prefix_prompt_pool(self) -> None: @@ -164,7 +208,9 @@ def _create_prefix_prompt_pool(self) -> None: pool_size = self.prefix_prompts.pool_size or 0 self._prefix_prompts = [self.generate_prompt(length) for _ in range(pool_size)] self.debug( - lambda: f"Initialized prefix prompts pool with {len(self._prefix_prompts)} prompts" + lambda: ( + f"Initialized prefix prompts pool with {len(self._prefix_prompts)} prompts" + ) ) def generate( @@ -433,8 +479,10 @@ def generate_user_context_prompt(self, session_index: int) -> str: new_prompt = self.generate_prompt(length) self._user_context_prompts.append(new_prompt) self.debug( - lambda: f"Generated user context prompt #{len(self._user_context_prompts) - 1} " - f"for session {len(self._user_context_prompts) - 1}" + lambda: ( + f"Generated user context prompt #{len(self._user_context_prompts) - 1} " + f"for session {len(self._user_context_prompts) - 1}" + ) ) return self._user_context_prompts[session_index] diff --git a/src/aiperf/dataset/graph/__init__.py b/src/aiperf/dataset/graph/__init__.py new file mode 100644 index 0000000000..83786551e3 --- /dev/null +++ b/src/aiperf/dataset/graph/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Graph loader package - schema models, parser, validator. + +See `docs/benchmark-modes/agentic.md` for the canonical spec. Import from the +submodules directly (`models`, `parser`, `validator`, `workload_detect`, ...); +this package intentionally re-exports nothing to keep submodule imports light. +""" diff --git a/src/aiperf/dataset/graph/adapters/__init__.py b/src/aiperf/dataset/graph/adapters/__init__.py new file mode 100644 index 0000000000..edbef38130 --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/__init__.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Adapters that convert third-party benchmark formats into graph ParsedGraph. + +Each format adapter exposes a `from_X(path)` callable for direct use (the +native facade has none; it delegates to ``parser.parse_native``), and every +adapter class is registered under the `graph_adapter` plugin category. +""" + +from aiperf.dataset.graph.adapters.dynamo.trace import DynamoTraceAdapter +from aiperf.dataset.graph.adapters.native import NativeGraphAdapter +from aiperf.dataset.graph.adapters.protocols import GraphAdapterProtocol +from aiperf.dataset.graph.adapters.weka.trace import WekaTraceAdapter + +__all__ = [ + "DynamoTraceAdapter", + "GraphAdapterProtocol", + "NativeGraphAdapter", + "WekaTraceAdapter", +] diff --git a/src/aiperf/dataset/graph/adapters/dag_jsonl/__init__.py b/src/aiperf/dataset/graph/adapters/dag_jsonl/__init__.py new file mode 100644 index 0000000000..1b8cf32c80 --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/dag_jsonl/__init__.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Native ``dag_jsonl`` graph adapter. + +Reuses the legacy ``dag_jsonl`` loader and expands its Conversation output into +per-root instanced trees (:mod:`.tree`), then lowers those trees onto the +unified-segment-store graph IR (:mod:`.lowering`). Consumers import the leaf +modules (:mod:`.tree`, :mod:`.lowering`) directly rather than the legacy loader +internals. +""" diff --git a/src/aiperf/dataset/graph/adapters/dag_jsonl/lowering.py b/src/aiperf/dataset/graph/adapters/dag_jsonl/lowering.py new file mode 100644 index 0000000000..4310326880 --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/dag_jsonl/lowering.py @@ -0,0 +1,247 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Lowering: expanded ``dag_jsonl`` trees -> unified-segment-store ParsedGraph. + +The byte-parity core of the ``dag_jsonl`` graph adapter. Each +:class:`~aiperf.dataset.graph.adapters.dag_jsonl.tree.DagTree` becomes one +per-trace :class:`GraphRecord` (mirroring the native lowering's per-trace +graphs) against ONE shared content-addressed :class:`SegmentPool`: + +* Every authored message is interned VERBATIM via ``SegmentPool.add_raw_message`` + (key order and extra keys preserved), parent-chained across the whole per-node + walk so shared conversation prefixes dedup into a real KV-cache prefix. +* A node with message-context lineage carries an ordered assembly program + (``metadata["trie"]["assembly"]``) interleaving interned static segments with + ``{"s": {"src": }}`` live-reply slots -- the worker splices each slot + as ``{"role": "assistant", "content": }``, byte-matching the + legacy captured-assistant message for text responses. Slot producers are + stamped ``capture: true``; lineage-free nodes carry no assembly key (pure + static -> bytes path at dispatch). +* Model / stream / token cap / tools ride the NATIVE node fields (``model`` / + ``streaming`` / ``max_tokens`` / ``raw_tools``, Turn naming); + ``extra_body`` carries the merged vendor keys + (``**endpoint_extra, **turn extra``). The run's + ``--extra-inputs`` (``endpoint_extra``) are folded at parse with the + legacy precedence -- the later turn-``extra`` update wins on overlap -- and + every node is stamped + ``metadata["dispatch"]["endpoint_extra_applied"] = True`` so the worker skips + its own ``endpoint.extra`` re-merge (parse-time folding is authoritative even + when the run has no extras). Legacy-vs-graph wire parity is proven + order-insensitively (canonical sorted-keys body comparison), so the legacy + body KEY ORDER is not reproduced -- keys and values are. +* Firing topology: one completion-anchored ``StaticEdge`` per predecessor + (authored turn delay in microseconds), ``START`` edges for predecessor-less + nodes, ``END`` edges from sink nodes, and AND-fan-in ``ChannelRequirement`` + gates for SPAWN_JOIN inputs plus every spliced lineage producer (credits are + minted only after the captures they splice have landed). +""" + +from __future__ import annotations + +from typing import Any + +from aiperf.dataset.graph.adapters.dag_jsonl.tree import DagNodeSpec, DagTree +from aiperf.dataset.graph.models import ( + END_NODE_ID, + START_NODE_ID, + ChannelRequirement, + ChannelSpec, + GraphRecord, + LlmNode, + ParsedGraph, + ProvenanceSpec, + StaticEdge, + TraceRecord, +) +from aiperf.dataset.graph.segment_ir.envelope import stamp_prompt_segment_ids +from aiperf.dataset.graph.segment_ir.pool import SegmentPool + +__all__ = ["lower_dag_trees"] + +_PROVENANCE = ProvenanceSpec(source="dag_jsonl", tool="aiperf-dag-jsonl-lowering/1") + + +def lower_dag_trees( + trees: list[DagTree], + *, + default_model: str | None, + run_streaming: bool, + endpoint_extra: list[tuple[str, Any]] | None = None, +) -> ParsedGraph: + """Lower expanded dag trees onto the unified interned segment store. + + One :class:`GraphRecord` per tree in ``parsed.graphs[trace_id]`` with a + matching ``TraceRecord(id=trace_id, graph_ref=trace_id)``; ``parsed.graph`` + aliases the first tree's record. The lowering is a pure function of its + input (content-addressed ids, no RNG, no clock), so dataset-plane and + timing-plane re-parses stamp identical graphs. ``endpoint_extra`` (the run's + ``--extra-inputs`` pairs) is folded into every node's ``extra_body`` + at the legacy precedence (see :func:`_lower_node`). + """ + if not trees: + raise NotImplementedError( + "dag_jsonl workload: the file expands to zero root trees; " + "there is nothing to lower" + ) + pool = SegmentPool() + graphs: dict[str, GraphRecord] = {} + traces: list[TraceRecord] = [] + for tree in trees: + graphs[tree.trace_id] = _lower_tree( + pool, + tree, + default_model=default_model, + run_streaming=run_streaming, + endpoint_extra=endpoint_extra, + ) + traces.append(TraceRecord(id=tree.trace_id, graph_ref=tree.trace_id)) + return ParsedGraph( + graph=graphs[traces[0].id], + graphs=graphs, + traces=traces, + segment_pool=pool, + ) + + +def _lower_tree( + pool: SegmentPool, + tree: DagTree, + *, + default_model: str | None, + run_streaming: bool, + endpoint_extra: list[tuple[str, Any]] | None, +) -> GraphRecord: + capture_ids = {aid for spec in tree.nodes.values() for aid in spec.lineage} + nodes: dict[str, LlmNode] = {} + edges: list[StaticEdge] = [] + for spec in tree.nodes.values(): + nodes[spec.node_id] = _lower_node( + pool, + spec, + tree, + capture_ids, + default_model=default_model, + run_streaming=run_streaming, + endpoint_extra=endpoint_extra, + ) + if spec.predecessors: + edges.extend( + StaticEdge( + source=pred_id, + target=spec.node_id, + delay_after_predecessor_us=float(delay_us) if delay_us else None, + ) + for pred_id, delay_us in spec.predecessors + ) + else: + edges.append(StaticEdge(source=START_NODE_ID, target=spec.node_id)) + sources = {edge.source for edge in edges} + edges.extend( + StaticEdge(source=nid, target=END_NODE_ID) + for nid in nodes + if nid not in sources + ) + # Every LlmNode writes its response into a per-node ``{node_id}_out`` + # channel; the executor's channel store rejects writes to undeclared + # channels, and successors read them only to enforce the AND-fan-in WAIT + # (prompt content comes from the segment pool / dynamic slots, never the + # channel value), so the default single-producer overwrite spec suffices. + state = {f"{nid}_out": ChannelSpec() for nid in nodes} + return GraphRecord( + version="2.0", + provenance=_PROVENANCE, + state=state, + nodes=nodes, + edges=edges, + ) + + +def _lower_node( + pool: SegmentPool, + spec: DagNodeSpec, + tree: DagTree, + capture_ids: set[str], + *, + default_model: str | None, + run_streaming: bool, + endpoint_extra: list[tuple[str, Any]] | None, +) -> LlmNode: + # The full conversation prefix is interned at build time: each lineage + # ancestor contributes its verbatim authored messages plus one live-reply + # slot, then the node's own messages follow. ``parent`` threads through the + # WHOLE walk (slots intern nothing) so shared prefixes dedup across nodes. + assembly: list[dict[str, Any]] = [] + parent: str | None = None + for ancestor_id in spec.lineage: + for message in tree.nodes[ancestor_id].turn.raw_messages: + parent = pool.add_raw_message(message=message, parent_id=parent) + assembly.append({"seg": parent}) + assembly.append({"s": {"src": ancestor_id}}) + for message in spec.turn.raw_messages: + parent = pool.add_raw_message(message=message, parent_id=parent) + assembly.append({"seg": parent}) + + # Model / stream / token cap / tools ride the NATIVE node fields (Turn + # naming); ``extra_body`` carries only the merged vendor keys. + # Legacy merge: endpoint extras (--extra-inputs) update first, the turn's + # own extras update last -- the turn value wins on an overlapping key. + extra_body: dict[str, Any] = {} + extra_body.update(endpoint_extra or []) + extra_body.update(spec.turn.extra_body or {}) + + # SPAWN_JOIN gates plus one completion gate per spliced lineage producer + # (disjoint sets by construction: join leaves are spawned children, never + # in the joining node's own message-context lineage). + inputs = [ + ChannelRequirement(channel=f"{join_id}_out", count=1) + for join_id in spec.join_inputs + ] + inputs.extend( + ChannelRequirement(channel=f"{ancestor_id}_out", count=1) + for ancestor_id in spec.lineage + ) + + node = LlmNode( + prompt=[], + output=f"{spec.node_id}_out", + streaming=run_streaming, + model=spec.turn.model or default_model, + max_tokens=spec.turn.max_tokens, + raw_tools=_effective_tools(spec, tree), + extra_body=extra_body, + inputs=inputs, + arrival_offset_us=0, + # Unconditional: parse-time folding owns the run's --extra-inputs even + # when there are none, so the worker's endpoint.extra re-merge (which + # would clobber turn-extra values) is always skipped for dag nodes. + # The "dag" stamp carries the instance's legacy record identity; the + # timing plane's re-parse reads it to stamp agent_depth / + # parent_correlation_id on every credit (parse-plane metadata only -- + # the worker's store envelope never carries it). + metadata={ + "dispatch": {"endpoint_extra_applied": True}, + "dag": { + "agent_depth": spec.agent_depth, + "parent_node": spec.parent_node_id, + }, + }, + ) + segment_ids = [token["seg"] for token in assembly if "seg" in token] + extra: dict[str, Any] = {} + if spec.lineage: + extra["assembly"] = assembly + if spec.node_id in capture_ids: + extra["capture"] = True + return stamp_prompt_segment_ids(node, segment_ids, extra=extra or None) + + +def _effective_tools(spec: DagNodeSpec, tree: DagTree) -> list[dict[str, Any]] | None: + """Own turn's ``raw_tools``, else the nearest lineage ancestor's (legacy + parity: ``raw_tools`` is the lone per-turn field that walks history).""" + if spec.turn.raw_tools is not None: + return spec.turn.raw_tools + for ancestor_id in reversed(spec.lineage): + tools = tree.nodes[ancestor_id].turn.raw_tools + if tools is not None: + return tools + return None diff --git a/src/aiperf/dataset/graph/adapters/dag_jsonl/trace.py b/src/aiperf/dataset/graph/adapters/dag_jsonl/trace.py new file mode 100644 index 0000000000..4e58ba590b --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/dag_jsonl/trace.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Adapter entry + detection for the ``dag_jsonl`` graph workload. + +Thin registration/selection seam over the ``dag_jsonl`` core: :func:`from_dag_jsonl` +composes the three already-tested layers -- the legacy loader +(:func:`~aiperf.dataset.graph.adapters.dag_jsonl.tree.load_dag_conversations`), +per-root tree expansion +(:func:`~aiperf.dataset.graph.adapters.dag_jsonl.tree.expand_trees`), and the +unified-segment-store lowering +(:func:`~aiperf.dataset.graph.adapters.dag_jsonl.lowering.lower_dag_trees`) -- +threading its run-derived knobs through. :class:`DagJsonlGraphAdapter` implements +:class:`~aiperf.dataset.graph.adapters.protocols.GraphAdapterProtocol` and is +registered under ``graph_adapter.dag_jsonl`` (``detection_priority: 90``). + +Explicit-only: excluded from workload auto-detection (see +``workload_detect._AUTODETECT_EXCLUDED``) so legacy +``--custom-dataset-type dag_jsonl`` runs stay on the linear pipeline. Selected via +``--graph-format dag_jsonl``. ``can_load`` remains a real, tested sniff for a +future autodetect flip. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import orjson + +from aiperf.dataset.graph.adapters.dag_jsonl.lowering import lower_dag_trees +from aiperf.dataset.graph.adapters.dag_jsonl.tree import ( + expand_trees, + load_dag_conversations, +) +from aiperf.dataset.graph.models import ParsedGraph +from aiperf.dataset.graph.parse_context import GraphParseContext + +__all__ = ["DagJsonlGraphAdapter", "from_dag_jsonl"] + +# Dynamo's ``dynamo.request.trace.v1`` discriminator; a dag sniff rejects any +# line carrying it so the two ``.jsonl`` sniffs stay mutually exclusive. +_DYNAMO_SCHEMA = "dynamo.request.trace.v1" + + +def from_dag_jsonl( + path: str, + *, + default_model: str | None = None, + run_streaming: bool = True, + delay_cap_seconds: float | None = None, + endpoint_extra: list[tuple[str, Any]] | None = None, +) -> ParsedGraph: + """Load, expand, and lower a ``dag_jsonl`` file into a :class:`ParsedGraph`. + + Composes the three ``dag_jsonl`` layers with no state of its own: the legacy + loader (delay-cap clamp) -> per-root instanced trees -> unified interned + segment store. Every knob is a pure function of its arguments, so build-plane + and schedule-plane parses called with identical arguments stamp + byte-identical graphs. + + Args: + path: Path to a ``dag_jsonl`` file. + default_model: Fallback stamped onto every node's native ``model`` field + when a turn has no authored ``model`` (the worker's own dispatch + fallback -- ``ModelEndpointInfo.from_run(run).primary_model_name``). + run_streaming: Resolved endpoint ``stream`` flag, stamped onto every + node's ``streaming`` / body ``stream``. + delay_cap_seconds: Per-turn inter-turn delay cap (seconds); the same + value the legacy loader clamps authored delays with. + endpoint_extra: The run's ``--extra-inputs`` pairs + (``ModelEndpointInfo.from_run(run).endpoint.extra``), folded into + every node's ``extra_body`` at the legacy + precedence (turn ``extra`` wins on overlap). + """ + conversations = load_dag_conversations( + Path(path), delay_cap_seconds=delay_cap_seconds + ) + trees = expand_trees(conversations) + return lower_dag_trees( + trees, + default_model=default_model, + run_streaming=run_streaming, + endpoint_extra=endpoint_extra, + ) + + +class DagJsonlGraphAdapter: + """DAG-conversation-JSONL graph adapter (``graph_adapter.dag_jsonl``).""" + + @classmethod + def can_load(cls, path: Path) -> bool: + """True iff ``path`` is a ``.jsonl`` whose first non-blank line is a dag line. + + A dag line is a JSON object carrying both ``session_id`` and ``turns`` + and WITHOUT the dynamo schema discriminator (dynamo's sink envelope + unwrapped first) -- keeping this sniff mutually exclusive with + ``DynamoTraceAdapter.can_load``. Bounded to the first non-blank line; + never raises (any error -> ``False``). + """ + try: + if not path.is_file() or path.suffix.lower() != ".jsonl": + return False + with path.open("rb") as f: + for raw in f: + stripped = raw.strip() + if not stripped: + continue + rec = orjson.loads(stripped) + return ( + isinstance(rec, dict) + and "session_id" in rec + and "turns" in rec + and not _is_dynamo_record(rec) + ) + except (OSError, orjson.JSONDecodeError): + return False + return False + + @classmethod + def parse(cls, path: Path, ctx: GraphParseContext | None = None) -> ParsedGraph: + """Convert ``path`` into a :class:`ParsedGraph` via :func:`from_dag_jsonl`. + + ``ctx`` carries the four run-derived dispatch knobs (``default_model`` + / ``run_streaming`` / ``delay_cap_seconds`` / ``endpoint_extra``), each + forwarded ONLY when set so a ctx-less parse (CLI tooling / direct + callers) keeps the :func:`from_dag_jsonl` defaults — in particular a + partial ctx never clobbers ``run_streaming=True`` with ``None``. + + Wraps the entry call in :func:`_assert_dag_zero_arrival_offsets`: + production always dispatches through this seam (the registry-driven + ``parse_graph``), so the all-zero arrival-offset invariant the + t*/dynamic-slot gate's carve-out relies on is enforced on every parse. + """ + kwargs: dict[str, Any] = {} + if ctx is not None: + if ctx.default_model is not None: + kwargs["default_model"] = ctx.default_model + if ctx.run_streaming is not None: + kwargs["run_streaming"] = ctx.run_streaming + if ctx.delay_cap_seconds is not None: + kwargs["delay_cap_seconds"] = ctx.delay_cap_seconds + if ctx.endpoint_extra is not None: + kwargs["endpoint_extra"] = ctx.endpoint_extra + parsed = from_dag_jsonl(str(path), **kwargs) + _assert_dag_zero_arrival_offsets(parsed) + return parsed + + +def _assert_dag_zero_arrival_offsets(parsed: ParsedGraph) -> None: + """Hold the dag_jsonl invariant that every node has a zero arrival offset. + + The t*/dynamic-slot gate (``workload_detect._gate_dynamic_slots_vs_tstar``) + carves out graphs whose EVERY node carries an explicit-zero + ``arrival_offset_us`` -- the shape dag lowering stamps -- because all-zero + offsets degenerate the t* snapshot chop to a no-op (see the carve-out + comment there for the full argument). This guard wraps ``from_dag_jsonl`` + inside :meth:`DagJsonlGraphAdapter.parse` (production always dispatches + through this seam post-ladder), walking the parsed graph once (O(nodes)) + and raising if the lowering ever stamps a nonzero offset -- that would + engage t* against a graph the lowering did not time, silently + mis-partitioning nodes into warmup. If dag ever emits recorded offsets + intentionally, delete this guard and let ``_gate_dynamic_slots_vs_tstar`` + gate it like any recorded workload (the carve-out stops matching on its + own). + """ + for record in (parsed.graph, *parsed.graphs.values()): + for node_id, node in record.nodes.items(): + if node.arrival_offset_us: + raise ValueError( + f"dag_jsonl node '{node_id}' has arrival_offset_us=" + f"{node.arrival_offset_us}, but dag_jsonl lowering must " + "stamp a zero arrival offset on every node -- the " + "t*/dynamic-slot gate's explicit-zero carve-out is " + "justified by that invariant." + ) + + +def _is_dynamo_record(rec: dict) -> bool: + """True iff ``rec`` (dynamo sink envelope unwrapped) is a dynamo trace record.""" + from aiperf.dataset.graph.adapters.dynamo.trace_reader import unwrap_sink_envelope + + return unwrap_sink_envelope(rec).get("schema") == _DYNAMO_SCHEMA diff --git a/src/aiperf/dataset/graph/adapters/dag_jsonl/tree.py b/src/aiperf/dataset/graph/adapters/dag_jsonl/tree.py new file mode 100644 index 0000000000..39c343ce5d --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/dag_jsonl/tree.py @@ -0,0 +1,325 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Pure-topology middle layer for the ``dag_jsonl`` graph adapter. + +Reuses the legacy ``dag_jsonl`` loader output (``Conversation`` objects with +``branches``, per-turn ``prerequisites`` / ``branch_ids``, and ``is_root``) and +expands each root conversation into an instanced tree of :class:`DagNodeSpec`s. +Each node captures its lineage (message-context ancestors), completion-anchored +firing predecessors with microsecond delays, and SPAWN_JOIN fan-in gates. No +graph IR types are produced here; :mod:`.lowering` lowers these trees onto the +graph IR. + +All dag.md load-time validation (cycles, single-parent FORK, system-message +placement, join_at bounds) is delegated to the legacy loader — this module +does NOT re-validate. The only additional gate is that a session id must not +contain the framing characters used by the graph correlation-id scheme. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path + +from aiperf.common.enums import ConversationBranchMode, PrerequisiteKind +from aiperf.common.models.branch import ConversationBranchInfo +from aiperf.common.models.dataset_models import Conversation, Turn +from aiperf.dataset.loader._delay_cap import DelayCapTracker +from aiperf.dataset.loader.dag_jsonl import DagJsonlLoader + + +@dataclass(slots=True) +class DagNodeSpec: + """One dispatchable turn instance inside one root tree.""" + + node_id: str + """Unique per tree: ``":"``; ```` is the + session_id, suffixed ``"#"`` (n>=2) for repeated SPAWN instantiations of + one template.""" + turn: Turn + """The legacy Turn (raw_messages/raw_tools/model/max_tokens/extra_body/delay).""" + lineage: list[str] + """node_ids whose (messages + live reply) precede this node's own messages, + oldest first. FORK child turn 0: parent's full lineage + the forking parent + node. SPAWN child turn 0 and roots: []. Turn k>0: lineage(k-1) + [k-1's id].""" + predecessors: list[tuple[str, int]] + """(pred node_id, delay_us) completion-anchored firing deps. Sequential edge + carries this turn's authored delay (ms -> us); fork/spawn child turn-0 edges + carry 0 (legacy dispatches children immediately at parent credit return, and + root turn-0 delay is NOT honored by the legacy rate loop - both preserved).""" + join_inputs: list[str] = field(default_factory=list) + """Leaf node_ids of SPAWN branches gating this turn (SPAWN_JOIN fan-in).""" + agent_depth: int = 0 + """Legacy record identity: 0 for root-chain turns, owner depth + 1 for + FORK/SPAWN child instances, 1 for pre-session spawn children (mirrors + ``conversation_source.start_branch_child`` / ``start_pre_session_child``). + Every turn of one child instance carries the instance's depth.""" + parent_node_id: str | None = None + """The triggering parent node (the fork/spawn turn) for child instances; + None for roots AND pre-session children (legacy: no parent session exists + at pre-dispatch). Same value on every turn of the instance.""" + + +@dataclass(slots=True) +class DagTree: + """All instanced turn nodes reachable from one root conversation.""" + + trace_id: str + """Root session_id; namespaces this tree's node_ids.""" + nodes: dict[str, DagNodeSpec] + """node_id -> spec, insertion-ordered by a deterministic depth-first walk.""" + + +def load_dag_conversations( + path: Path, *, delay_cap_seconds: float | None +) -> dict[str, list[Conversation]]: + """Load a ``dag_jsonl`` file into the loader's session_id -> [Conversation] map. + + Wraps the legacy :class:`DagJsonlLoader` standalone construction so later + tasks never touch loader internals. The standalone constructor hard-codes + the delay cap to disabled (it has no ``run`` to read it from), so when a cap + is requested we replace the tracker before the first ``load_dataset()`` call + parses turns and applies the clamp. + """ + loader = DagJsonlLoader(path) + if delay_cap_seconds is not None: + loader._delay_cap_tracker = DelayCapTracker(cap_seconds=delay_cap_seconds) + return loader.load_dataset() + + +def expand_trees(conversations: dict[str, list[Conversation]]) -> list[DagTree]: + """Expand each root conversation into an instanced :class:`DagTree`. + + Roots are the conversations with ``is_root=True``, walked in the loader's + declaration (file) order; each yields one tree. SPAWN children are + instantiated fresh per reference, FORK children exactly once (single-parent + guaranteed by the loader). + """ + conv_by_sid: dict[str, Conversation] = {} + for convs in conversations.values(): + for conv in convs: + conv_by_sid[conv.session_id] = conv + + for sid in conv_by_sid: + if "|" in sid or "#" in sid: + raise NotImplementedError( + f"conversation '{sid}': session ids containing '|' or '#' " + "collide with graph correlation-id framing" + ) + + trees: list[DagTree] = [] + for convs in conversations.values(): + for conv in convs: + if conv.is_root: + trees.append(_expand_tree(conv, conv_by_sid)) + return trees + + +def _resolve_join_inputs( + *, turn: Turn, branch_leaves: dict[str, list[str]] +) -> list[str]: + """SPAWN_JOIN fan-in leaves gating ``turn``. + + Each SPAWN_JOIN prerequisite naming an already-fired branch contributes + that branch's post-spawn child-instance leaves, in prerequisite order. + Branches that have not fired yet resolve to no leaves. + """ + join_inputs: list[str] = [] + for prereq in turn.prerequisites: + if prereq.kind == PrerequisiteKind.SPAWN_JOIN and prereq.branch_id is not None: + join_inputs.extend(branch_leaves.get(prereq.branch_id, [])) + return join_inputs + + +def _child_entry_context( + *, + branch: ConversationBranchInfo, + is_pre: bool, + node_id: str, + lineage: list[str], + predecessors: list[tuple[str, int]], +) -> tuple[list[str], list[tuple[str, int]]]: + """Turn-0 seed ``(lineage, predecessors)`` for a child branch. + + Pre-session spawns fire on the owner turn-0's own trigger, so they copy its + predecessors and start context-free. Forks inherit the parent's message + context and fire at parent credit return (delay 0). Plain spawns start + context-free and also fire at parent credit return. + """ + if is_pre: + return [], list(predecessors) + if branch.mode == ConversationBranchMode.FORK: + return [*lineage, node_id], [(node_id, 0)] + return [], [(node_id, 0)] + + +@dataclass(slots=True) +class _TreeBuilder: + """Mutable per-tree expansion state for one root conversation. + + Holds the template lookup plus the growing node map and per-template + instantiation counter so the recursive branch walk can stay a set of small + methods rather than one deeply nested closure. + """ + + conv_by_sid: dict[str, Conversation] + nodes: dict[str, DagNodeSpec] = field(default_factory=dict) + # Per-tree instantiation counts: the first instance of a template session + # id is bare, subsequent SPAWN references get a "#" (n>=2) suffix. + instance_counts: dict[str, int] = field(default_factory=dict) + + def alloc_instance(self, sid: str) -> str: + """Mint this tree's instance id for template ``sid``. + + The first instance of a template is bare; subsequent SPAWN references + get a ``"#"`` (n>=2) suffix. + """ + count = self.instance_counts.get(sid, 0) + 1 + self.instance_counts[sid] = count + return sid if count == 1 else f"{sid}#{count}" + + def _turn_context( + self, + *, + instance_id: str, + turn: Turn, + k: int, + turn0_lineage: list[str], + turn0_predecessors: list[tuple[str, int]], + ) -> tuple[list[str], list[tuple[str, int]]]: + """Lineage and firing predecessors for one turn node. + + Turn 0 inherits the caller-supplied seeds (fork/spawn/root entry); a + later turn chains sequentially off ``k-1`` carrying its authored delay + (ms -> us). + """ + if k == 0: + return list(turn0_lineage), list(turn0_predecessors) + prev_id = f"{instance_id}:{k - 1}" + lineage = [*self.nodes[prev_id].lineage, prev_id] + predecessors = [(prev_id, int((turn.delay or 0) * 1000))] + return lineage, predecessors + + def _expand_branches( + self, + *, + turn: Turn, + branch_by_id: dict[str, ConversationBranchInfo], + branch_leaves: dict[str, list[str]], + node_id: str, + lineage: list[str], + predecessors: list[tuple[str, int]], + agent_depth: int, + ) -> None: + """Instantiate and recurse into every child branch fired by ``turn``. + + Records post-spawn branch leaves (so a later turn's SPAWN_JOIN can + resolve its fan-in) before recursing with per-branch-kind entry context. + Child instances carry the legacy identity: FORK/SPAWN children get + ``agent_depth + 1`` with this firing node as their parent; pre-session + children get depth 1 with no parent (``start_pre_session_child``). + """ + for branch_id in turn.branch_ids: + branch = branch_by_id.get(branch_id) + if branch is None: + raise ValueError( + f"node {node_id!r} fires branch {branch_id!r}, but the loader " + f"emitted no matching branch on this conversation " + f"(loader-consistency invariant violated)" + ) + is_pre = branch.dispatch_timing == "pre" + for child_sid in branch.child_conversation_ids: + child_conv = self.conv_by_sid[child_sid] + child_instance = self.alloc_instance(child_sid) + child_leaf = f"{child_instance}:{len(child_conv.turns) - 1}" + if branch.mode == ConversationBranchMode.SPAWN and not is_pre: + branch_leaves.setdefault(branch_id, []).append(child_leaf) + + child_lineage, child_predecessors = _child_entry_context( + branch=branch, + is_pre=is_pre, + node_id=node_id, + lineage=lineage, + predecessors=predecessors, + ) + self.expand( + child_conv, + child_instance, + turn0_lineage=child_lineage, + turn0_predecessors=child_predecessors, + agent_depth=1 if is_pre else agent_depth + 1, + parent_node_id=None if is_pre else node_id, + ) + + def expand( + self, + conv: Conversation, + instance_id: str, + *, + turn0_lineage: list[str], + turn0_predecessors: list[tuple[str, int]], + agent_depth: int, + parent_node_id: str | None, + ) -> None: + """Emit every turn node for one conversation instance, depth-first. + + Turn nodes are created in order; each fired branch is expanded + immediately after its owning turn so child node_ids interleave with the + parent's remaining turns exactly as the legacy walk produced them. + ``agent_depth`` / ``parent_node_id`` are the INSTANCE's legacy identity, + stamped identically on every turn node the instance emits. + """ + branch_by_id = {b.branch_id: b for b in conv.branches} + # branch_id -> leaf node_ids of its post-SPAWN child instances, filled as + # branches fire so a later turn's SPAWN_JOIN can resolve its fan-in. + branch_leaves: dict[str, list[str]] = {} + + for k, turn in enumerate(conv.turns): + node_id = f"{instance_id}:{k}" + lineage, predecessors = self._turn_context( + instance_id=instance_id, + turn=turn, + k=k, + turn0_lineage=turn0_lineage, + turn0_predecessors=turn0_predecessors, + ) + self.nodes[node_id] = DagNodeSpec( + node_id=node_id, + turn=turn, + lineage=lineage, + predecessors=predecessors, + join_inputs=_resolve_join_inputs( + turn=turn, branch_leaves=branch_leaves + ), + agent_depth=agent_depth, + parent_node_id=parent_node_id, + ) + self._expand_branches( + turn=turn, + branch_by_id=branch_by_id, + branch_leaves=branch_leaves, + node_id=node_id, + lineage=lineage, + predecessors=predecessors, + agent_depth=agent_depth, + ) + + +def _expand_tree(root: Conversation, conv_by_sid: dict[str, Conversation]) -> DagTree: + builder = _TreeBuilder(conv_by_sid=conv_by_sid) + root_instance = builder.alloc_instance(root.session_id) + # Recursive descent: expand -> _expand_branches -> expand adds ~3 stack frames + # per fork/spawn nesting level, so recursion depth scales with authored nesting, + # not tree width. Realistic DAGs nest a handful of levels; a pathological + # >300-level chain would approach CPython's default recursion limit and want an + # explicit iterative walk instead. + builder.expand( + root, + root_instance, + turn0_lineage=[], + turn0_predecessors=[], + agent_depth=0, + parent_node_id=None, + ) + return DagTree(trace_id=root.session_id, nodes=builder.nodes) diff --git a/src/aiperf/dataset/graph/adapters/dynamo/__init__.py b/src/aiperf/dataset/graph/adapters/dynamo/__init__.py new file mode 100644 index 0000000000..b90834b1e2 --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/dynamo/__init__.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dynamo trace-format loader (``dynamo.request.trace.v1``). + +The public entry points (and the chain-grouping / session-tree build) live in +:mod:`.trace`; the remaining modules are the internal reader +(:mod:`.trace_reader`), trie-lowering (:mod:`.trie_lowering`), fused-parallel +build (:mod:`.trace_parallel`), and segment-pool shim (:mod:`.store_backed_pool`) +stages. +""" + +from aiperf.dataset.graph.adapters.dynamo.trace import ( + DynamoTraceAdapter, + from_dynamo_trace, +) + +__all__ = ["DynamoTraceAdapter", "from_dynamo_trace"] diff --git a/src/aiperf/dataset/graph/adapters/dynamo/store_backed_pool.py b/src/aiperf/dataset/graph/adapters/dynamo/store_backed_pool.py new file mode 100644 index 0000000000..84d69c20e8 --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/dynamo/store_backed_pool.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dynamo-scoped :class:`SegmentPool` shims for the two dynamo build routes. + +Both dynamo routes re-list every message chain into each node's +``prompt_segment_ids``, so absent interning the same ``segment_id`` hexdigest is +born as one fresh str per re-listing (``~H/N`` copies per node) even though +:meth:`SegmentPool.add` already stores exactly one canonical +:class:`~aiperf.dataset.graph.segment_ir.pool.Segment` per unique value. These +two shims make ``add`` return that first-born canonical object so every +re-listing shares it -- values are byte-identical (interning changes str +identity only), so the store bytes, sidecar, and envelope are unaffected. + +* :class:`InterningSegmentPool` -- the EAGER route + (``from_dynamo_trace(direct_store=None)``). It subclasses + :class:`SegmentPool` and returns ``self._by_id[sid].id`` (the first-born + ``Segment.id``) from every ``add*``. The canonical is the object already + retained as the ``_by_id`` key/Segment, so the intern table is FREE. + +* :class:`StoreBackedSegmentPool` -- the DIRECT write-through route + (``GraphStoreBuilder._build_graph_store_streaming``), which constructs the + unified store BEFORE the parse and threads it into + ``from_dynamo_trace(direct_store=...)``. Every ``pool.add`` interns its segment + STRAIGHT INTO the store at parse time (no second RAM pool copy) -- the store's + own ``_ids`` first-occurrence dedup gives the same handle stream the eager + ``build_unified_trie_store_interned`` drain would assign, so the produced store + is byte-identical (both routes intern in ``build_trie_ir``'s content-loop + first-occurrence order, the single ordering authority). A handle-indexed + ``_sids`` list holds one canonical str per unique value so repeats return the + first-born object instead of the fresh hexdigest. + +The direct shim's ``_by_id`` stays empty (segments live in the store, not the +pool), so the returned ``ParsedGraph.segment_pool`` no-ops the interned drain's +put loop and ``strip_replay_text`` replaces it with a fresh empty ``SegmentPool`` +before the content-free sidecar is msgpack-encoded (nothing ever encodes a live +shim). + +Both shims live here rather than in ``pool.py`` so ``pool.py`` stays a +dependency-free stdlib leaf: the store type is referenced only under +``TYPE_CHECKING`` here. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from aiperf.dataset.graph.segment_ir.pool import Segment, SegmentPool, segment_id + +if TYPE_CHECKING: + from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + ) + +_DIRECT_ONLY = ( + "StoreBackedSegmentPool is the dynamo DIRECT write-through shim: the dynamo " + "content path interns segments ONLY via add() (add_message_chain + " + "trie_content emit/small-prompt), and nothing reads the pool back after " + "build_trie_ir. {method}() is unsupported -- an adapter that reaches it " + "(text/raw segments, or a read-back materialize/get) must build through the " + "eager SegmentPool/InterningSegmentPool route, not the direct store route." +) + + +class InterningSegmentPool(SegmentPool): + """Eager-route :class:`SegmentPool` that returns the FIRST-BORN canonical sid. + + Every ``add*`` returns ``self._by_id[sid].id`` -- the ``Segment.id`` string + stored on the segment's first occurrence -- instead of the freshly computed + hexdigest :meth:`SegmentPool.add` returns. On a dedup hit the fresh duplicate + dies immediately (refcount) and the canonical is what flows into every node's + ``prompt_segment_ids``, so a value re-listed across turns/sessions shares one + str object. On the first occurrence the returned object IS the fresh sid + (``_by_id[sid].id is sid``), so there is zero cost beyond one dict hit. + + This depends on ``Segment(id=sid)`` being stored with the FIRST-BORN sid + (:meth:`SegmentPool.add` inserts only when ``sid not in self._by_id``), which + is stable while ``pool.py`` is frozen. Values are byte-identical to a plain + ``SegmentPool`` (identity-only change), so store/sidecar/envelope bytes hold. + + ``SegmentPool`` is ``@dataclass(slots=True)``; ``__slots__ = ()`` keeps this + subclass slot-only (no per-pool ``__dict__``). + """ + + __slots__ = () + + def add( + self, *, role: str, content: str, tokens: list[int], parent_id: str | None + ) -> str: + sid = super().add( + role=role, content=content, tokens=tokens, parent_id=parent_id + ) + return self._by_id[sid].id + + def add_text(self, *, role: str, content: str, parent_id: str | None) -> str: + sid = super().add_text(role=role, content=content, parent_id=parent_id) + return self._by_id[sid].id + + def add_raw_message(self, *, message: dict[str, Any], parent_id: str | None) -> str: + sid = super().add_raw_message(message=message, parent_id=parent_id) + return self._by_id[sid].id + + +class StoreBackedSegmentPool(SegmentPool): + """A :class:`SegmentPool` whose :meth:`add` writes through to the unified store. + + Only :meth:`add` is a real operation (the dynamo path's sole pool call); every + other pool method raises :class:`NotImplementedError` naming the dynamo-only + write-through contract, so any non-dynamo adopter fails loud instead of + silently interning into an empty ``_by_id`` the store never sees. + + A handle-indexed ``_sids`` list holds one canonical sid str per unique value + (in first-occurrence order) so a repeated segment returns the first-born + object, mirroring :class:`InterningSegmentPool` on the direct route. + """ + + __slots__ = ("_store", "_sids") + + def __init__(self, store: GraphSegmentUnifiedBackingStore) -> None: + super().__init__() + self._store = store + # Handle -> canonical sid str. ``put_segment`` returns dense insertion + # indexes on first occurrence, so index i holds the sid first interned at + # handle i; a repeat's handle indexes back into this list to return the + # first-born object (the canonical intern table for the direct route). + self._sids: list[str] = [] + + def add( + self, + *, + role: str, + content: str, + tokens: list[int], + parent_id: str | None, + ) -> str: + """Compute the content-addressed id, write it through, return the canonical. + + The sid VALUE is what :meth:`SegmentPool.add` returns. ``store.put_segment`` + dedups first-occurrence on its ``_ids`` map (a repeat ``sid`` is a no-op + that returns the existing handle) and derives the persisted blob as + ``orjson.dumps({"role", "content"})`` -- byte-identical to what the eager + drain writes for the same token/text segment. The returned OBJECT is the + first-born canonical: on a first occurrence ``handle == len(self._sids)`` + so the fresh sid is appended and returned; on a repeat ``handle`` indexes + the earlier canonical out of ``_sids``. The defensive third branch (a + ``handle`` outrunning ``_sids`` because the store was pre-populated before + this shim existed -- never in production, where the shim is the store's + sole parse-time writer) degrades to returning the fresh value-correct sid. + """ + sid = segment_id(parent_id, role, tokens) + handle = self._store.put_segment(sid, role, content) + if handle == len(self._sids): + self._sids.append(sid) + return sid + if handle < len(self._sids): + return self._sids[handle] + return sid + + def add_text(self, *, role: str, content: str, parent_id: str | None) -> str: + raise NotImplementedError(_DIRECT_ONLY.format(method="add_text")) + + def add_raw_message(self, *, message: dict[str, Any], parent_id: str | None) -> str: + raise NotImplementedError(_DIRECT_ONLY.format(method="add_raw_message")) + + def get(self, sid: str) -> Segment: + raise NotImplementedError(_DIRECT_ONLY.format(method="get")) + + def materialize(self, path_ids: list[str]) -> list[dict]: + raise NotImplementedError(_DIRECT_ONLY.format(method="materialize")) + + +__all__ = ["InterningSegmentPool", "StoreBackedSegmentPool"] diff --git a/src/aiperf/dataset/graph/adapters/dynamo/trace.py b/src/aiperf/dataset/graph/adapters/dynamo/trace.py new file mode 100644 index 0000000000..fcd0c49923 --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/dynamo/trace.py @@ -0,0 +1,1090 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Adapter: convert Dynamo agent-trace files to the flat segment-trie ParsedGraph. + +One trace file (or segmented prefix / directory) carries ``request_end`` +records for one or more ``agent_context.session_id`` trajectories. Sessions are +grouped into SESSION-TREES (a root session plus every descendant linked by +``agent_context.parent_session_id``; see :func:`group_chains_into_trees`) and +each tree is lowered INDEPENDENTLY via the reusable per-tree seam +(:func:`_build_graph_from_chains`) into its OWN single-graph ``ParsedGraph`` +(``LlmNode``s ``{session_id}:{k}``, weka-identical IR). Those per-tree +``ParsedGraph``s are merged via +:func:`~aiperf.dataset.graph.merge.merge_parsed_graphs` (the SAME helper the +weka multi-item path uses) into ONE MULTI-GRAPH workload: one ``TraceRecord`` +per tree, keyed by its root session id into ``graphs[root_id]`` with +``graph_ref=root_id``. Building per tree drops CROSS-PARENT interval-order edges +(between independent trees) by construction while preserving every WITHIN-tree +edge (parent<->subagent + intra-session); the shared build order (hash-prefix +content parents, frozen block tags, covered-count ISL gate, interval-order +finished-before edges) is owned by +:func:`~aiperf.dataset.graph.segment_ir.trie_content.build_trie_ir` and runs +over each tree's node set. Prompt content is content-addressed into a +:class:`~aiperf.dataset.graph.segment_ir.pool.SegmentPool` at parse +time -- there is no separate channel-replay build pass. The per-tree seam is +process-local so a later increment can fan trees to worker processes; the +segment pool / block size / content seed are all resolved ONCE and pinned into +every tree so the union is byte-identical to a single global build over +disjoint-content trees. + +Hash source (see +:func:`~aiperf.dataset.graph.adapters.dynamo.trie_lowering.dynamo_trie_nodes`): +the recorded ``request.replay.input_sequence_hashes`` at the recorded +``trace_block_size`` whenever a record carries replay metadata (the +recorded-when-present rule the weka path follows). Current dynamo emits +replay on EVERY ``request_end`` (it skips the record entirely when the KV +block size is unavailable), so records without replay metadata only occur in +older captures or hand-authored traces; those fall back to +per-session VIRTUAL negative ids sized to ``input_tokens`` (block size +:data:`~aiperf.dataset.graph.adapters.dynamo.trie_lowering.DEFAULT_VIRTUAL_BLOCK_SIZE`) +and tag the trace ``virtual-hash-fallback``; consecutive turns of a session +still share a content prefix. + +Schema source (dynamo repo): lib/llm/src/request_trace/types.rs + and lib/llm/src/protocols/common/extensions.rs (AgentContext) +""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Container, Iterable +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import msgspec +import orjson + +from aiperf.common.environment import Environment +from aiperf.dataset.graph.adapters.dynamo.trace_reader import ( + DYNAMO_TRACE_SCHEMA_V1, + AgentTraceRecord, + DynamoTraceAdapterError, + EmptyDynamoTraceError, + _dir_segment_sort_key, + iter_trace_records, + unwrap_sink_envelope, +) +from aiperf.dataset.graph.adapters.shared.idle_gap import ( + IDLE_GAP_CAP_USE_DEFAULT, + resolve_idle_gap_cap, +) +from aiperf.dataset.graph.adapters.shared.selection import SelectionStats +from aiperf.dataset.graph.merge import merge_parsed_graphs +from aiperf.dataset.graph.models import ( + GraphRecord, + LlmNode, + NodeUnion, + ParsedGraph, + ProvenanceSpec, + StaticEdge, + TraceRecord, +) +from aiperf.dataset.graph.parse_context import ( + UNSET, + GraphParseContext, + publish_ctx_tokenizer_env, +) + +if TYPE_CHECKING: + from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + ) + +_DEFAULT_TAG = "from-dynamo-trace" +_SOURCE = "dynamo_trace" + + +def _collect_chains( + path: str | Path, + session_id_filter: str | None, + *, + max_depth: int, +) -> dict[str, _Chain]: + """Read trace records, sort per session, and reduce to non-empty chains.""" + by_session, parent_link, skipped_no_context = _collect_records( + path, session_id_filter + ) + if not by_session: + if skipped_no_context: + raise EmptyDynamoTraceError( + f"{path}: {skipped_no_context} records had no agent_context " + f"(replay-only trace?); dynamo agent-trace parsing requires " + f"session identity" + ) + raise EmptyDynamoTraceError(f"{path}: no trace records found") + + for sid in by_session: + by_session[sid].sort(key=lambda r: r.event_time_unix_ms) + + chains: dict[str, _Chain] = {} + for sid, recs in by_session.items(): + chain = _records_to_chain( + recs, session_id=sid, parent_session_id=parent_link.get(sid) + ) + if chain.turns: + chains[sid] = chain + + if not chains: + raise EmptyDynamoTraceError( + f"{path}: no request_end events found across {len(by_session)} sessions" + ) + + _guard_chain_forest(chains, parent_link, max_depth=max_depth) + return chains + + +def from_dynamo_trace( + path: str | Path, + *, + tag: str = _DEFAULT_TAG, + session_id_filter: str | None = None, + idle_gap_cap_seconds: float | None | object = IDLE_GAP_CAP_USE_DEFAULT, + content_root_seed: int | None = None, + content_tokenizer: str | None = None, + prompt_corpus: str = "coding", + release_replay: bool = False, + direct_store: GraphSegmentUnifiedBackingStore | None = None, + num_dataset_entries: int | None = None, + max_context_length: int | None = None, + selection_out: list[SelectionStats] | None = None, +) -> ParsedGraph: + """Parse a Dynamo agent-trace file/dir/prefix into the flat segment-trie IR. + + Every session's turns lower to one ``LlmNode`` per ``request_end`` + (``{session_id}:{k}``, the recorded session id verbatim + 0-based turn) + in a SINGLE flat graph. Sessions are first grouped + into session-trees (a root plus its ``parent_session_id`` descendants) and + each tree is lowered independently, so parent and subagent sessions of ONE + tree coexist with edges whose concurrency is EMERGENT from the recorded + intervals (a child overlapping its parent gets no edge between them; + disjoint intervals get a finished-before edge), while INDEPENDENT trees + never gain a cross-parent edge -- their interval order is derived only over + their own nodes. Prompt content is + reconstructed deterministically from ``(content_tokenizer, + prompt_corpus, content_root_seed)`` and content-addressed into the + returned ``ParsedGraph.segment_pool``; each node's ISL is the block- + aligned COVERED count of its hashes (never history + full-reconstruction + double counting). The eager route interns via an + :class:`~aiperf.dataset.graph.adapters.dynamo.store_backed_pool.InterningSegmentPool` + so equal ``prompt_segment_ids`` across turns/sessions share one canonical + str object (values byte-identical to a plain ``SegmentPool``). + + Args: + path: Path to a ``.jsonl`` file, ``.jsonl.gz`` file, segmented prefix, + or a directory containing one of the above. + tag: Base tag added to EVERY emitted trace (each session-tree is its own + single-root trace, so there is no ``"multi-root"`` tag). + session_id_filter: When set, restrict to records whose + ``agent_context.session_id`` matches. + idle_gap_cap_seconds: Per-tree idle-gap cap (seconds) on the recorded + timeline that stamps node ``arrival_offset_us`` and the edge delays + (``delay_after_predecessor_us`` on the binding predecessor, + ``min_start_delay_us`` on START edges) -- the SAME + ``ActiveIdleWarp`` the weka adapter applies, so recorded + end-to-start gaps replay with over-long idle gaps compressed to the + cap. The run path resolves this from the dataset's + ``synthesis.idle_gap_cap_seconds`` (``--synthesis-idle-gap-cap``); + left unset here it defaults to 60.0. Pass ``None`` to disable + warping (raw recorded gaps replay unchanged). + content_root_seed: Seed pinning the deterministic content synthesis; + ``None`` resolves via the ambient bootstrap root seed, else fresh + per-run OS entropy (identical ladder to the weka adapter). + content_tokenizer: Tokenizer for content synthesis (``None`` selects + the builtin deterministic tokenizer). + prompt_corpus: Corpus the content synthesizer samples from. + release_replay: When ``True``, free each record's ``request.replay`` + hash lists inside :func:`dynamo_trie_nodes` once they are copied + into the trie IR (a build-time RAM adjunct). SAFE ONLY because this + function lowers freshly-read ``chains`` exactly once per call; + defaults ``False`` so a caller that re-lowers the same in-memory + records never silently degrades to the virtual-hash fallback. The + production store build (:meth:`DynamoTraceAdapter.parse`) opts in. + direct_store: When set, the build plane's live + :class:`~aiperf.dataset.graph_segment_unified_store.GraphSegmentUnifiedBackingStore`. + The returned ``ParsedGraph.segment_pool`` becomes a + :class:`~aiperf.dataset.graph.adapters.dynamo.store_backed_pool.StoreBackedSegmentPool` + whose ``add()`` write-throughs intern each segment STRAIGHT INTO the + store at parse time (no second RAM pool copy) -- the Stage B direct + write-through route. ``None`` (every direct caller / tooling) keeps + the eager ``SegmentPool``. Threaded only via the build-plane seam + (``GraphStoreBuilder`` -> ``parse_graph_workload`` kwargs passthrough), + never through the format-agnostic ``GraphParseContext``. + num_dataset_entries: ``--num-dataset-entries`` cap on the number of + session-TREES built (the graph-plane fix for ai-dynamo/aiperf#1106). + ``None`` (default) builds every eligible tree. A tree is a root + session plus its subagent descendants -- the selection unit, so a + capped load never splits a tree. + max_context_length: ``--max-context-length`` per-tree ceiling on peak + context (max ``input_length + output_tokens`` over the tree's + records, via :func:`dynamo_tree_peak_context`). Over-limit trees are + rejected BEFORE the build; the first ``num_dataset_entries`` eligible + trees (root-sorted) are then kept (filter THEN cap). ``None`` applies + no context filter. Selection applies on BOTH the serial and + fused-parallel build paths. + selection_out: Optional sink; when provided AND a knob is set, the single + :class:`SelectionStats` for the filter-then-cap scan is appended for + the caller's report. Untouched when both knobs are ``None``. + + Returns: + A MULTI-GRAPH ``ParsedGraph`` (via :func:`merge_parsed_graphs`): one + ``TraceRecord`` per session-tree, each keyed by its root session id into + ``graphs[root_id]`` and selected by ``graph_ref=root_id``; traces are + id-sorted and the content-addressed ``segment_pool`` is the union of the + per-tree pools. ``graph`` is the FIRST tree (lex-min root, the default + single-graph slot), NOT the whole-capture union. Each tree's graph holds its sessions' + ``LlmNode``s + WITHIN-tree interval-order edges (cross-parent edges + absent by construction). + """ + max_depth = Environment.DYNAMO.MAX_SUBAGENT_DEPTH + + # Same seed ladder as the weka routes (explicit seed -> ambient bootstrap + # root seed -> per-run OS entropy): dynamo shares the weka content + # synthesizer, so an unresolved None here would mean ambient-RNG content + # in-process but offline-default content in tooling -- the divergence the + # weka unification removed. Resolved ONCE here and pinned into every tree + # build so all trees share the same seed -> byte-identical pools. + from aiperf.dataset.graph.adapters.shared.content import ( + resolve_effective_root_seed, + ) + + content_root_seed = resolve_effective_root_seed(content_root_seed) + + # Resolve the shared weka/dynamo idle-gap sentinel ONCE (60s default / + # explicit float / None = warp disabled) so the fused-parallel and serial + # builds warp identically. + idle_gap_cap = resolve_idle_gap_cap(idle_gap_cap_seconds) + + # Fused read+build parallel path: a cheap grouping scan (session ids + + # parent links only, NO ``input_sequence_hashes`` parse) decides tree + # membership in the parent, then raw record lines are shuffled to per-batch + # temp files and READ+BUILT inside worker processes, so the giant recorded + # hash arrays never cross a process boundary -- only the compact built graph + # returns. Gated to the pure tooling / direct-caller route (no live + # write-through store, which cannot cross processes) and a whole-corpus build + # (a single-session filter is one tree -> never parallel). Returns None below + # the tree-count threshold / for a single tree, and the serial + # read-then-build path below runs instead. + if direct_store is None and session_id_filter is None: + from aiperf.dataset.graph.adapters.dynamo import trace_parallel + + fused = trace_parallel.maybe_build_fused_parallel( + path, + content_root_seed=content_root_seed, + idle_gap_cap_seconds=idle_gap_cap, + content_tokenizer=content_tokenizer, + prompt_corpus=prompt_corpus, + release_replay=release_replay, + max_depth=max_depth, + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + selection_out=selection_out, + ) + if fused is not None: + return _finalize_parsed_graph(fused, tag=tag) + + chains = _collect_chains(path, session_id_filter, max_depth=max_depth) + if num_dataset_entries is not None or max_context_length is not None: + chains = _select_chains_filter_then_cap( + chains, + path=path, + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + selection_out=selection_out, + ) + per_tree = _build_trees_flat( + chains, + content_root_seed=content_root_seed, + idle_gap_cap_seconds=idle_gap_cap, + content_tokenizer=content_tokenizer, + prompt_corpus=prompt_corpus, + release_replay=release_replay, + direct_store=direct_store, + ) + return _finalize_parsed_graph(per_tree, tag=tag) + + +def _select_chains_filter_then_cap( + chains: dict[str, _Chain], + *, + path: str | Path, + num_dataset_entries: int | None, + max_context_length: int | None, + selection_out: list[SelectionStats] | None, +) -> dict[str, _Chain]: + """Filter-then-cap the collected chains at the session-TREE granularity. + + The SERIAL-path selection seam (the fused path selects in the parent scan): + groups the already-collected chains into session-trees (root-sorted), screens + each by :func:`dynamo_tree_peak_context` against ``max_context_length``, and + keeps the first ``num_dataset_entries`` eligible trees. Reuses that helper + directly over the flattened per-tree ``request_end`` records, so the selected + set matches the fused path's hash-free scan selection. Returns the union of + the kept trees' chains; raises :class:`EmptyDynamoTraceError` when every tree + is filtered out (a ceiling below the whole capture is a user error, not an + empty graph). + """ + from aiperf.dataset.graph.adapters.shared.peak_context import ( + dynamo_tree_peak_context, + ) + from aiperf.dataset.graph.adapters.shared.selection import ( + filter_then_cap, + log_selection_summary, + ) + + parent_link = { + sid: c.parent_session_id + for sid, c in chains.items() + if c.parent_session_id is not None + } + trees = group_chains_into_trees(chains, parent_link) + + def _candidates() -> Iterable[tuple[dict[str, _Chain], int]]: + for tree in trees: + records = [turn.record for chain in tree.values() for turn in chain.turns] + yield tree, dynamo_tree_peak_context(records) + + kept_trees, stats = filter_then_cap( + _candidates(), + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + ) + # Parent-side finalize point for the SERIAL path (and the fused-decline + # fallback): log the summary once here. The fused build path logs its own + # when it actually fans out; the two paths are mutually exclusive per build. + log_selection_summary( + stats, + source=str(path), + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + ) + if selection_out is not None: + selection_out.append(stats) + if not kept_trees: + raise EmptyDynamoTraceError( + f"{path}: all {stats.scanned} session-trees exceeded " + f"max_context_length={max_context_length}; nothing to build" + ) + selected: dict[str, _Chain] = {} + for tree in kept_trees: + selected.update(tree) + return selected + + +def _finalize_parsed_graph(per_tree: list[ParsedGraph], *, tag: str) -> ParsedGraph: + """Fold the base ``tag`` into each per-tree trace, then merge into ONE workload. + + Both build paths hand this the list of per-tree SINGLE-graph + ``ParsedGraph``s (each = one root ``TraceRecord`` with ``graph_ref=None``, + that tree's ``GraphRecord``, that tree's pool). This stamps the base ``tag`` + onto each trace (each tree is its own single-root trace, so there is no + ``"multi-root"`` tag) and returns :func:`merge_parsed_graphs`, which -- the + SAME helper the weka multi-item path uses -- keys each trace's graph under + ``graphs[trace.id]``, sets ``graph_ref=trace.id``, guards duplicate ids, + id-sorts the traces, unions the content-addressed pools, and keeps ``graph`` + as the FIRST tree (the default single-graph slot). + """ + tagged = ( + msgspec.structs.replace( + pg, + traces=[ + msgspec.structs.replace(trace, tags=sorted({tag, *trace.tags})) + for trace in pg.traces + ], + ) + for pg in per_tree + ) + return merge_parsed_graphs(tagged) + + +# --- session-tree grouping + per-tree build --------------------------------- + + +def group_chains_into_trees( + chains: dict[str, _Chain], + parent_link: dict[str, str], +) -> list[dict[str, _Chain]]: + """Partition chains into session-trees by ``parent_session_id`` linkage. + + A tree is a root session (one whose parent is absent, self, or external to + ``chains``) plus every descendant reachable through ``parent_link``. Each + session is unioned to its in-set root by walking ``parent_link`` to a + fixpoint; a session with no parent and no children is its own singleton + tree. Independent trees never share a dict, so lowering each one on its own + node set drops cross-parent interval-order edges by construction while + keeping every within-tree edge. + + ``parent_link`` maps ``session_id -> parent_session_id`` (the caller derives + it from each chain's ``parent_session_id``); a link whose parent is not in + ``chains`` marks the child as a forest root. The walk carries a per-session + ``seen`` guard so a ``parent_link`` cycle (already rejected upstream by + :func:`_guard_chain_forest`) can never spin here -- it halts at the repeat. + + Deterministic: trees are returned sorted by root session id and each tree's + dict is insertion-ordered by session id, so the downstream union is + byte-stable across runs. + """ + root_of = root_of_sessions(chains, parent_link) + trees: dict[str, dict[str, _Chain]] = defaultdict(dict) + for sid in sorted(chains): + trees[root_of[sid]][sid] = chains[sid] + return [trees[root] for root in sorted(trees)] + + +def root_of_sessions( + session_ids: Container[str] | Iterable[str], + parent_link: dict[str, str], +) -> dict[str, str]: + """Map each session id to its in-set tree root by walking ``parent_link``. + + ``session_ids`` need only support membership (``in``) and iteration -- a + ``chains`` dict (keyed by session id) or a plain ``set`` both work, so the + parent's cheap grouping scan (which only knows session ids, never the built + chains) and :func:`group_chains_into_trees` share ONE walk and therefore + ONE grouping decision. The walk stops at the first parent absent from the + set (a forest root) and carries a per-session ``seen`` guard so a + ``parent_link`` cycle halts at the repeat rather than spinning. + """ + root_of: dict[str, str] = {} + for sid in session_ids: + seen: set[str] = {sid} + cur = sid + while True: + parent = parent_link.get(cur) + if parent is None or parent not in session_ids or parent in seen: + break + seen.add(parent) + cur = parent + root_of[sid] = cur + return root_of + + +def _build_trees_flat( + chains: dict[str, _Chain], + *, + content_root_seed: int, + idle_gap_cap_seconds: float | None, + content_tokenizer: str | None, + prompt_corpus: str, + release_replay: bool, + direct_store: GraphSegmentUnifiedBackingStore | None, +) -> list[ParsedGraph]: + """Group ``chains`` into session-trees and build one ``ParsedGraph`` per tree. + + Returns the per-tree single-graph ``ParsedGraph`` list (the caller merges + them via :func:`merge_parsed_graphs`). The block size is resolved ONCE + across the WHOLE capture (fail-loud on a mix, exactly as the single global + build did) and pinned into every tree; the content seed is likewise pinned + by the caller, so every tree build is byte-identical to a single global + build over disjoint-content trees. + + This is the SERIAL build over already-collected ``chains`` -- the path taken + when the fused read+build parallel dispatch in :func:`from_dynamo_trace` + declines (below the tree-count threshold, a single-session filter, or a live + ``direct_store`` write-through sink that cannot cross process boundaries). + The parallel decision lives in :func:`from_dynamo_trace` (before the + expensive full read), not here, because the fused path must AVOID collecting + the hash-bearing chains in the parent in the first place. + """ + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + DEFAULT_VIRTUAL_BLOCK_SIZE, + _resolve_block_size, + ) + + # One block size across the whole capture (fail-loud on a mix); pinned into + # every tree so a replay-free tree never diverges to the virtual default. + block_size = _resolve_block_size(chains) or DEFAULT_VIRTUAL_BLOCK_SIZE + parent_link = { + sid: c.parent_session_id + for sid, c in chains.items() + if c.parent_session_id is not None + } + trees = group_chains_into_trees(chains, parent_link) + + return _build_trees_sequential( + trees, + block_size=block_size, + content_root_seed=content_root_seed, + idle_gap_cap_seconds=idle_gap_cap_seconds, + content_tokenizer=content_tokenizer, + prompt_corpus=prompt_corpus, + release_replay=release_replay, + direct_store=direct_store, + ) + + +def _build_trees_sequential( + trees: list[dict[str, _Chain]], + *, + block_size: int, + content_root_seed: int, + idle_gap_cap_seconds: float | None, + content_tokenizer: str | None, + prompt_corpus: str, + release_replay: bool, + direct_store: GraphSegmentUnifiedBackingStore | None, +) -> list[ParsedGraph]: + """Build each already-grouped session-tree into its OWN single-graph ParsedGraph. + + The in-process build loop, called BOTH by the serial dispatch of + :func:`_build_trees_flat` and by every pool worker (which receives its + contiguous BATCH of trees as ``trees``). Each returned ``ParsedGraph`` holds + one tree's ``GraphRecord``, one ``TraceRecord`` (id = the tree's root session + id, ``graph_ref=None``) carrying that tree's extra tags, and that tree's own + pool. The caller folds the base tag and merges the list via + :func:`merge_parsed_graphs` -- the SAME helper the weka multi-item path uses, + so ``.graphs[root_id]`` / ``graph_ref`` / id-sorted traces / content-pool + union are all byte-identical to weka. ``block_size`` is pinned by the caller + (resolved ONCE across the whole capture in the parent), so a replay-free tree + in a batch never diverges to the virtual default while a sibling batch + carries a recorded size. + + Pool strategy: with ``direct_store`` set the trees SHARE one + :class:`~aiperf.dataset.graph.adapters.dynamo.store_backed_pool.StoreBackedSegmentPool` + (the store is a single write-through sink whose ``by_id`` stays empty), so the + merged pool is the empty union the interned drain no-ops over. Otherwise each + tree builds into its OWN + :class:`~aiperf.dataset.graph.adapters.dynamo.store_backed_pool.InterningSegmentPool` + and its per-tree ``ParsedGraph`` carries a plain ``SegmentPool`` over that + tree's ``by_id``; :func:`merge_parsed_graphs` unions them by content-addressed + id (identical ids carry identical content under the pinned seed). This + per-tree-ParsedGraph shape is the seam the parallel path fans to worker + processes -- each worker returns its batch's LIST of per-tree ParsedGraph + blobs. + """ + from aiperf.dataset.graph.adapters.dynamo import store_backed_pool + from aiperf.dataset.graph.segment_ir.pool import SegmentPool + + shared_pool: SegmentPool | None = ( + store_backed_pool.StoreBackedSegmentPool(direct_store) + if direct_store is not None + else None + ) + + per_tree: list[ParsedGraph] = [] + for tree in trees: + # Shared write-through store, or a fresh interning pool per tree that + # this tree's ParsedGraph carries (the process-parallel seam: a worker + # returns its batch's per-tree ParsedGraph blobs, merged in the parent). + pool = ( + shared_pool + if shared_pool is not None + else store_backed_pool.InterningSegmentPool() + ) + graph, extra_tags = _build_graph_from_chains( + tree, + pool=pool, + content_root_seed=content_root_seed, + block_size=block_size, + idle_gap_cap_seconds=idle_gap_cap_seconds, + content_tokenizer=content_tokenizer, + prompt_corpus=prompt_corpus, + release_replay=release_replay, + ) + tree_pool: SegmentPool = ( + shared_pool if shared_pool is not None else SegmentPool(_by_id=pool.by_id) + ) + per_tree.append( + ParsedGraph( + graph=graph, + traces=[ + TraceRecord(id=_tree_root_id(tree), tags=sorted(set(extra_tags))) + ], + segment_pool=tree_pool, + ) + ) + return per_tree + + +def _tree_root_id(tree: dict[str, _Chain]) -> str: + """The tree's root session id = its lex-min session whose parent is out-of-tree. + + A well-formed session-tree has exactly ONE such session (every non-root + session's parent is in the tree by construction of :func:`root_of_sessions`); + the ``sorted(...)[0]`` is a defensive tiebreak and the ``min(tree)`` fallback + covers the degenerate all-have-parents case (a cycle, already rejected + upstream by :func:`_guard_chain_forest`). This is the SAME id the + whole-capture root sort used as the single trace id before the multi-graph + split, so a single-tree capture keeps its historical trace id. + """ + roots = sorted(sid for sid, c in tree.items() if c.parent_session_id not in tree) + return roots[0] if roots else min(tree) + + +def _build_graph_from_chains( + chains: dict[str, _Chain], + *, + pool: Any, + content_root_seed: int, + block_size: int, + idle_gap_cap_seconds: float | None, + content_tokenizer: str | None, + prompt_corpus: str, + release_replay: bool, +) -> tuple[GraphRecord, list[str]]: + """Lower ONE set of chains (a single session-tree) into a flat ``GraphRecord``. + + The reusable per-tree seam -- everything between ``dynamo_trie_nodes`` and + ``assemble_trie_graph`` -- so a later increment can fan each tree's chains + to a worker process that calls exactly this. Passing the WHOLE chain set + reproduces the pre-tree single global build (the differential oracle the + unit tests pin the tree-scoped union against). ``pool`` is caller-owned so + the orchestrator can share one write-through store across trees or union + per-tree pools; ``content_root_seed`` and ``block_size`` are resolved ONCE + by the caller and pinned here so every tree build is byte-identical. + + Heavy imports stay lazy: the content synthesizer / trie core are only needed + on a real parse, and ``DynamoTraceAdapter.can_load`` sniffing must not pull + the corpus machinery onto the detection path. + """ + from aiperf.common.tokenizer import BUILTIN_TOKENIZER_NAME + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + build_dynamo_llm_node, + dynamo_recon_callbacks, + dynamo_trie_nodes, + ) + from aiperf.dataset.graph.segment_ir.trie_content import ( + TrieNode, + assemble_trie_graph, + build_trie_ir, + ) + + nodes, bs, extra_tags = dynamo_trie_nodes( + chains, release_replay=release_replay, block_size=block_size + ) + callbacks = dynamo_recon_callbacks( + content_tokenizer or BUILTIN_TOKENIZER_NAME, + prompt_corpus, + content_root_seed, + block_size=bs, + trace_scope=_tree_root_id(chains), + ) + result = build_trie_ir( + nodes, + block_size=bs, + callbacks=callbacks, + pool=pool, + idle_gap_cap_seconds=idle_gap_cap_seconds, + small_prompt_fallback=True, + ) + + # build_dynamo_llm_node reads node.start (= warped_start), which is only + # stamped by build_trie_ir's warp pass -- assembly must run AFTER the build. + def build_node(node: TrieNode) -> LlmNode: + return build_dynamo_llm_node( + node, + build=result.builds[node.node_id], + ) + + graph = assemble_trie_graph( + nodes, + result, + build_node=build_node, + provenance=ProvenanceSpec(source=_SOURCE, tool="aiperf-dynamo-trie/1"), + ) + return graph, extra_tags + + +def _union_graphs(graphs: list[GraphRecord]) -> GraphRecord: + """Flatten per-tree ``GraphRecord``s into ONE (node maps unioned, edges concat). + + TEST-ONLY oracle helper as of the multi-graph restructure: the production + emit path no longer collapses trees into one union graph (it returns a + multi-graph ``ParsedGraph`` via :func:`merge_parsed_graphs`). This is + retained so the byte-equivalence guard can re-union the per-tree graphs and + assert the content matches the pre-change single union. + + Node ids are globally unique (``{sha256(session_id)[:16]}_a{turn}``, agent-id + collision-guarded per tree), so a cross-tree node-id collision is a real + corruption and fails loud. Output channel declarations (``state``) key on the + same unique node ids, so their union is a plain ``dict.update``. Edges are + already within-tree only, so concatenation introduces no cross-tree edge. + """ + nodes: dict[str, NodeUnion] = {} + edges: list[StaticEdge] = [] + state: dict[str, Any] = {} + for g in graphs: + for nid, node in g.nodes.items(): + if nid in nodes: + raise DynamoTraceAdapterError( + f"node id collision across session-trees: {nid!r}" + ) + nodes[nid] = node + edges.extend(g.edges) + state.update(g.state) + return GraphRecord( + version="2.0", + provenance=ProvenanceSpec(source=_SOURCE, tool="aiperf-dynamo-trie/1"), + state=state, + nodes=nodes, + edges=edges, + ) + + +_DYNAMO_TRACE_EVENT_TYPES = frozenset( + {"request_end", "tool_start", "tool_end", "tool_error"} +) + + +class DynamoTraceAdapter: + """Dynamo agent-trace v1 workload adapter (JSONL/JSONL.gz files, segmented dirs).""" + + @classmethod + def can_load(cls, path: Path) -> bool: + if path.is_dir(): + # Mirror discover_segments: any *.jsonl / *.jsonl.gz content + # (segmented or plain), sniffed on the first file in reader order. + segs = sorted( + ( + c + for c in path.iterdir() + if c.is_file() and c.name.endswith((".jsonl", ".jsonl.gz")) + ), + key=_dir_segment_sort_key, + ) + if not segs: + return False + return cls._first_record_matches(segs[0]) + if path.suffix.lower() in (".gz", ".jsonl"): + return cls._first_record_matches(path) + return False + + @classmethod + def _first_record_matches(cls, path: Path) -> bool: + """Sniff the first non-empty line; unreadable/corrupt bytes mean "not ours". + + Detection must never crash on a candidate file: a truncated final gzip + member raises EOFError, corrupt deflate data raises zlib.error, and a + non-gzip file behind a .gz name raises BadGzipFile (an OSError) -- all + of these are "this adapter does not claim the file", not errors. + """ + import gzip + import zlib + + is_gz = path.name.lower().endswith(".gz") + try: + with gzip.open(path, "rb") if is_gz else path.open("rb") as f: + for raw in f: + stripped = raw.strip() + if not stripped: + continue + rec = orjson.loads(stripped) + return isinstance(rec, dict) and _is_dynamo_trace_record(rec) + except (OSError, EOFError, zlib.error, orjson.JSONDecodeError): + return False + return False + + @classmethod + def parse( + cls, + path: Path, + ctx: GraphParseContext | None = None, + *, + direct_store: GraphSegmentUnifiedBackingStore | None = None, + ) -> ParsedGraph: + """Convert ``path`` into a :class:`ParsedGraph` via :func:`from_dynamo_trace`. + + ``ctx`` carries the run-derived knobs (seed / tokenizer / corpus / + idle-gap cap / selection), each forwarded ONLY when set so a partial + ctx never clobbers the entry's ``prompt_corpus="coding"`` default with + ``None``. ``idle_gap_cap_seconds`` follows the SAME tri-state rule as + the weka adapter: ``UNSET`` keeps the entry's shared 60s default, while + an explicit ``None`` forwards verbatim and DISABLES warping (the user's + ``synthesis.idle_gap_cap_seconds: null``) — recorded delays replay + either way, never silently zeroed. Generation is likewise always + pinned to the recorded ``output_tokens`` (weka parity; see + ``build_dynamo_llm_node``) — there is no behavior env knob left on + this entry. The run tokenizer + trust/revision publish to the loader-preload env before any callbacks + are built (:func:`publish_ctx_tokenizer_env`). + + This is the production store-build parse entry (the registry dispatches + here for ``dynamo_trace``), and it lowers freshly-read chains exactly + once, so it opts into ``release_replay=True`` to free the recorded + replay hash lists during the build. The dynamo-only knobs are set HERE, + at the adapter's own entry, rather than threaded through the + format-agnostic ``GraphParseContext`` / ``parse_graph`` dispatch. + + ``direct_store`` is the build-plane's live unified store, reaching this + entry via the ``parse_graph`` -> ``_parse_via_adapter`` ``**adapter_kwargs`` + passthrough (``GraphStoreBuilder`` -> ``parse_graph_workload(..., direct_store=...)``). + A keyword-only param -- like ``release_replay`` it is a dynamo-only knob + the uniform ``parse(path, ctx)`` registry call never supplies -- forwarded + ONLY when set so the ``direct_store=None`` protocol-default entry is + byte-identical to before (the same forward-only-when-set rule the ctx + knobs follow). + + ``ctx.num_dataset_entries`` / ``ctx.max_context_length`` forward the + filter-then-cap session-tree selection (ai-dynamo/aiperf#1106) when set; + this is the production dynamo build entry, so the run honors them on both + the serial and fused-parallel build paths. + """ + publish_ctx_tokenizer_env(ctx) + kwargs: dict[str, Any] = { + "release_replay": True, + } + if direct_store is not None: + kwargs["direct_store"] = direct_store + if ctx is not None: + if ctx.content_root_seed is not None: + kwargs["content_root_seed"] = ctx.content_root_seed + if ctx.content_tokenizer is not None: + kwargs["content_tokenizer"] = ctx.content_tokenizer + if ctx.prompt_corpus is not None: + kwargs["prompt_corpus"] = ctx.prompt_corpus + if ctx.idle_gap_cap_seconds is not UNSET: + kwargs["idle_gap_cap_seconds"] = ctx.idle_gap_cap_seconds + if ctx.num_dataset_entries is not None: + kwargs["num_dataset_entries"] = ctx.num_dataset_entries + if ctx.max_context_length is not None: + kwargs["max_context_length"] = ctx.max_context_length + return from_dynamo_trace(path, **kwargs) + + +def _is_dynamo_trace_record(rec: dict) -> bool: + """Sniff a raw record for the current ``dynamo.request.trace.v1`` schema. + + Dynamo's file sinks wrap every line in a ``{"timestamp", "event"}`` + envelope (see ``trace_reader.unwrap_sink_envelope``), so unwrap before + matching; bare records (fixtures, older captures) match directly. + ``agent_context`` is optional in the current schema (absent on replay-only + records), so a missing/``None`` context still identifies a dynamo record. + """ + rec = unwrap_sink_envelope(rec) + ac = rec.get("agent_context") + return ( + rec.get("schema") == DYNAMO_TRACE_SCHEMA_V1 + and rec.get("event_type") in _DYNAMO_TRACE_EVENT_TYPES + and (isinstance(ac, dict) or ac is None) + ) + + +# --- internal data shapes ------------------------------------------------- + + +class _Turn: + """One assistant turn = one ``request_end`` record.""" + + __slots__ = ("record",) + + def __init__(self, record: AgentTraceRecord) -> None: + self.record = record + + +class _Chain: + """Per-session chain: ordered turns + session-level identity.""" + + __slots__ = ( + "session_id", + "parent_session_id", + "turns", + ) + + def __init__( + self, + session_id: str, + *, + parent_session_id: str | None, + turns: list[_Turn], + ) -> None: + self.session_id = session_id + self.parent_session_id = parent_session_id + self.turns = turns + + +# --- record collection ---------------------------------------------------- + + +def _intern_replay_hashes(record: AgentTraceRecord, table: dict[int, int]) -> None: + """Rewrite the record's replay hash list in place with canonical int objects. + + Values never change (dict membership is by ``==``); only object identity is + shared, so every re-listed occurrence of a hash value across the whole + capture points at ONE ``int`` object instead of a fresh ~36 B ``orjson`` + allocation. Every downstream consumer of ``input_sequence_hashes`` / + ``TrieRequest.hash_ids`` is value-semantic (dict/set/list/compare), so + sharing objects is transparent -- store bytes, sidecar, and envelope are all + unchanged (the golden digest and three-way parity gate that). + + ``hashes[:] =`` mutates the SAME list object inside the (validated) pydantic + model, so no assignment-revalidation runs and the pre-intern duplicate + objects drop one record at a time (O(one record) transient), never the whole + capture at once. + + Negative-id footnote (performance, not correctness): recorded hashes are + validated non-negative in + :meth:`~aiperf.dataset.graph.adapters.dynamo.trace_reader.AgentReplayMetrics._reject_negative_hashes`, + so negatives never enter this table; the virtual negative ids are minted + later, at lowering, and never reach this code. Even hypothetically, + ``hash(-1) == hash(-2)`` only costs an extra bucket probe -- dict membership + is decided by ``==``, so equal hashes for unequal values can never alias. + """ + req = record.request + replay = req.replay if req is not None else None + if replay is None: + return + hashes = replay.input_sequence_hashes + sd = table.setdefault + hashes[:] = [sd(h, h) for h in hashes] + + +def _record_identity(record: AgentTraceRecord) -> tuple | None: + """Stable dedup key, or None when the record carries no usable identity. + + Aggregated capture dirs can hold the SAME record twice: dynamo's dual file + sinks share one output path, and the tool-event ZMQ ingress republishes + into every subscribing process's local bus. request_end dedups on the + frontend-unique request_id; tool events on (type, call id, time, status). + """ + ctx = record.agent_context + sid = ctx.session_id if ctx is not None else None + if record.event_type == "request_end": + if record.request is None: + return None + return ("request_end", sid, record.request.request_id) + if record.tool is None: + return None + return ( + record.event_type, + sid, + record.tool.tool_call_id, + record.event_time_unix_ms, + record.tool.status, + ) + + +def _collect_records( + path: str | Path, + session_id_filter: str | None, +) -> tuple[dict[str, list[AgentTraceRecord]], dict[str, str], int]: + """Group records by session; also build the parent map and a skip counter. + + Returns ``(by_session, parent_link, skipped_no_context)`` where + ``skipped_no_context`` counts records dropped for carrying no + ``agent_context`` (replay-only captures) -- the caller uses it to + distinguish "empty trace" from "trace without session identity". + + This is the SINGLE point where every dynamo record is materialized before + lowering: the streaming reader is fully drained into ``by_session`` here, so + all ``H`` recorded ``input_sequence_hashes`` slots are simultaneously live + (the "record-window plateau") before a single ``TrieRequest`` exists. Read- + time interning of each record's replay hashes (via :func:`_intern_replay_hashes`, + below) is therefore the ONLY interception that caps that plateau -- and it + also carries the canonical objects on into lowering, since the ``list()`` + copy in ``dynamo_trie_nodes`` preserves element identity. + + ``intern`` is a PER-PARSE ``dict[int, int]`` local: born before the read loop, + dead when this function returns (before lowering), never module/global state + and never crossing parses -- re-parses get a fresh table, exactly like the + lowering virtual-id counter. Its resident footprint is one entry per unique + hash value (~5.2 MB @6k, ~0.9 GB @1M); the transient dict-resize high-water + during growth is ~1.5x that resident size and is dropped at return, safely + below the downstream content-loop peak, so it never sets the parse's + high-water mark. + """ + by_session: dict[str, list[AgentTraceRecord]] = defaultdict(list) + parent_link: dict[str, str] = {} + seen: set[tuple] = set() + skipped_no_context = 0 + intern: dict[int, int] = {} + for record in iter_trace_records(path, session_id=session_id_filter): + ctx = record.agent_context + if ctx is None: + # Replay-only record with no session identity: skip grouping. + skipped_no_context += 1 + continue + identity = _record_identity(record) + if identity is not None: + if identity in seen: + continue + seen.add(identity) + _intern_replay_hashes(record, intern) + by_session[ctx.session_id].append(record) + # Subagent linkage: real captures carry it in ``parent_trajectory_id`` + # (``trajectory_id == session_id`` in every observed record, so the + # parent's trajectory id IS its session id); ``parent_session_id`` is a + # fallback for older / hand-authored traces and is used only when no + # ``parent_trajectory_id`` is present. AgentContext is stamped per + # request from headers, so records of one session may disagree (parent + # header only on later calls). First non-self parent wins. This map is + # the SINGLE parent authority: the forest guard walks it and + # _records_to_chain receives the chain's parent from it, so the two can + # never diverge. A self-parent (parent == this session) is passed + # through verbatim by dynamo's generic header mapping and means "no + # parent", never a cycle. + parent = ctx.parent_trajectory_id or ctx.parent_session_id + if parent and parent != ctx.session_id and ctx.session_id not in parent_link: + parent_link[ctx.session_id] = parent + return by_session, parent_link, skipped_no_context + + +def _records_to_chain( + recs: list[AgentTraceRecord], + *, + session_id: str, + parent_session_id: str | None, +) -> _Chain: + """Group time-ordered records into _Turn objects pivoted on request_end. + + ``parent_session_id`` comes from the caller's ``parent_link`` map (first + non-self parent across ALL of the session's records), so a session whose + earliest records lack the parent header -- e.g. a harness tool event + before any request_end -- still links to its parent. + """ + # tool_start / tool_end / tool_error records are recognized (see + # _DYNAMO_TRACE_EVENT_TYPES) but not lowered: tool time is already implicit + # in the recorded end-to-start gaps the replay honors, and no consumer + # reads a per-node tool breakdown. + turns: list[_Turn] = [ + _Turn(record=rec) for rec in recs if rec.event_type == "request_end" + ] + return _Chain( + session_id=session_id, + parent_session_id=parent_session_id, + turns=turns, + ) + + +# --- cycle / depth guard -------------------------------------------------- + + +def _guard_chain_forest( + chains: dict[str, _Chain], + parent_link: dict[str, str], + *, + max_depth: int, +) -> None: + """Raise on parent_link cycles or chains exceeding ``max_depth``. + + A cycle is any path A -> B -> ... -> A through ``parent_link``. Depth is + counted from the chain's root inclusive; a root is depth 1, its child is 2. + """ + for pid in chains: + seen: list[str] = [] + cur: str | None = pid + while cur is not None: + if cur in seen: + cycle = " -> ".join(seen + [cur]) + raise DynamoTraceAdapterError(f"parent_link cycle detected: {cycle}") + seen.append(cur) + if len(seen) > max_depth: + path = " -> ".join(seen) + raise DynamoTraceAdapterError( + f"parent_link depth exceeds AIPERF_DYNAMO_MAX_SUBAGENT_DEPTH" + f"={max_depth}: {path}" + ) + parent = parent_link.get(cur) + # Stop walking when parent isn't itself in the chain set + # (it's effectively the root for this forest). + if parent is None or parent not in chains: + break + cur = parent + + +__all__ = [ + "DynamoTraceAdapter", + "DynamoTraceAdapterError", + "EmptyDynamoTraceError", + "from_dynamo_trace", + "group_chains_into_trees", + "root_of_sessions", +] diff --git a/src/aiperf/dataset/graph/adapters/dynamo/trace_parallel.py b/src/aiperf/dataset/graph/adapters/dynamo/trace_parallel.py new file mode 100644 index 0000000000..e14408403a --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/dynamo/trace_parallel.py @@ -0,0 +1,750 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused read+build process-parallel session-tree build for the Dynamo adapter. + +This SUPERSEDES the earlier ship-collected-trees parallel path. That path read +every record in the PARENT and then shipped the raw ``input_sequence_hashes`` +(~900 ints per record, tens of millions per file) across the process boundary +to build workers -- the IPC serialization of those arrays was the wall (only +~1.35x on real captures). Here the raw hashes NEVER cross a process boundary: +they are parsed only inside the worker that builds them, and only the compact +built graph (no hashes) returns. + +Three phases, all reusing the weka pool lifecycle +(:func:`~aiperf.dataset.graph.adapters.weka.trace_parallel._run_pool_streaming` +-- forkserver context, parent-built shared-memory corpus, bounded ordered +window, graceful shutdown) with dynamo's own ``DYNAMO_GRAPH_PARALLEL_*`` tuning: + +1. **Grouping scan** (parent, :func:`_scan_grouping`): decompress every segment + ONCE and extract ONLY ``session_id`` + parent links + ``trace_block_size`` + per line via compiled regexes bounded to the prefix BEFORE the giant + ``input_sequence_hashes`` array (``agent_context`` and the block size both + precede it in the wire order), so the hash arrays are never scanned. This + yields the session -> tree-root assignment (the SAME walk + :func:`~aiperf.dataset.graph.adapters.dynamo.trace.root_of_sessions` runs for + the serial build), the pinned block size, and a per-session byte-length build + weight. + +2. **Shuffle** (parent, :func:`_shuffle_to_batch_files`): decompress every + segment a second time and append each raw record line VERBATIM (unparsed) to + a per-batch gzip temp file, routed by the line's session id -> its batch. + Streaming (bounded memory), skipping schema-less marker lines and records + with no session (dropped exactly as the serial path drops no-``agent_context`` + records). Each batch temp file is self-contained: every record of every tree + assigned to that batch, regardless of which source segment it came from, so a + subagent whose records are split across two files still lands whole in ONE + batch. + +3. **Fused build** (workers, :func:`_build_batch_file_to_blob`): each worker + READS its batch temp file (:func:`_collect_chains`) and BUILDS it + (:func:`_build_trees_sequential`, the inner build STRICTLY sequential -- no + nested pool) with the parent-pinned seed + block size + shared-memory corpus. + The shuffle makes every session-tree BATCH-LOCAL, so a worker always sees + COMPLETE trees and returns its batch's LIST of per-tree single-graph + ``ParsedGraph`` blobs (each keyed by its root id). The parent decodes each + frame in input (batch) order and FLATTENS all workers' per-tree graphs, then + the caller merges them via + :func:`~aiperf.dataset.graph.merge.merge_parsed_graphs`. + +Contiguous batching over the globally sorted-by-root tree list keeps the +parent's flattened per-tree list byte-identical to +:func:`~aiperf.dataset.graph.adapters.dynamo.trace._build_trees_sequential` over +the same capture, so the merged multi-graph workload is identical: same per-tree +node keys/order, same edge set, same content-addressed ``segment_pool``. +""" + +from __future__ import annotations + +import functools +import gzip +import os +import queue +import re +import shutil +import tempfile +import threading +from collections import defaultdict +from collections.abc import Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from pathlib import Path +from typing import IO, Any, TypeVar + +import msgspec + +from aiperf.dataset.graph.adapters.dynamo.trace_reader import ( + DynamoTraceAdapterError, + discover_segments, +) +from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + DEFAULT_VIRTUAL_BLOCK_SIZE, +) +from aiperf.dataset.graph.adapters.shared.selection import ( + SelectionStats, + log_selection_summary, +) +from aiperf.dataset.graph.codecs import ( + decode_parsed_graph_msgpack, + encode_parsed_graph_msgpack, +) +from aiperf.dataset.graph.models import ParsedGraph + +_T = TypeVar("_T") + +# Batch-result frame codecs: ``[encode_parsed_graph_msgpack(pg) bytes, tags]``. +# The nested typed ``ParsedGraph`` blob keeps the existing cross-process codec +# for the graph + pool; the outer plain frame carries the trace tags alongside. +_FRAME_ENCODER = msgspec.msgpack.Encoder() +_FRAME_DECODER = msgspec.msgpack.Decoder() + +# Regexes for the grouping scan. Every field they capture (``session_id``, the +# parent links, ``trace_block_size``) precedes ``input_sequence_hashes`` in the +# wire order, so the caller bounds every search to the prefix before that array +# and the giant hash lists are never scanned. ``"session_id":"`` cannot match +# inside ``"parent_session_id":"`` (the char before ``session_id`` there is ``_``, +# not ``"``), so the anchored quote reliably isolates the real session id. +_SESSION_RE = re.compile(rb'"session_id":"([^"]+)"') +_PARENT_TRAJ_RE = re.compile(rb'"parent_trajectory_id":"([^"]+)"') +_PARENT_SESSION_RE = re.compile(rb'"parent_session_id":"([^"]+)"') +_BLOCK_SIZE_RE = re.compile(rb'"trace_block_size":\s*(\d+)') +_REQUEST_END_MARK = b'"event_type":"request_end"' +_HASHES_KEY = b'"input_sequence_hashes"' + +# Hash-free scalar reads for the OPTIONAL peak-context selection scan. Every +# field precedes ``input_sequence_hashes`` in the wire order (``input_length`` +# lives inside ``replay`` before its hash array; ``input_tokens`` / +# ``output_tokens`` are ``request`` fields before ``replay``), so the caller's +# prefix bound keeps them hash-free. ``input_length`` appears ONLY in replay, so +# the regex reliably isolates it. These MIRROR +# :func:`~aiperf.dataset.graph.adapters.shared.peak_context.dynamo_tree_peak_context` +# so the parent's scan-based tree selection matches the serial helper. +_INPUT_LENGTH_RE = re.compile(rb'"input_length":\s*(\d+)') +_INPUT_TOKENS_RE = re.compile(rb'"input_tokens":\s*(\d+)') +_OUTPUT_TOKENS_RE = re.compile(rb'"output_tokens":\s*(\d+)') + +_TIMEOUT_HINT = ( + "dynamo graph build worker produced no result for one session-tree batch: " + "the worker process was most likely killed mid-build (OOM kill / external " + "SIGKILL) -- a raw multiprocessing Pool cannot complete a killed worker's " + "in-flight task. Reduce worker count " + "(AIPERF_DATASET_DYNAMO_GRAPH_PARALLEL_WORKERS) to lower peak memory, or " + "raise AIPERF_DATASET_DYNAMO_GRAPH_PARALLEL_ITEM_TIMEOUT_SECONDS if a single " + "batch legitimately builds slower than the timeout." +) + + +def _dynamo_threshold() -> int: + from aiperf.common.environment import Environment + + return max(0, Environment.DATASET.DYNAMO_GRAPH_PARALLEL_THRESHOLD) + + +def _dynamo_workers(*, item_count: int) -> int: + """Resolve the worker count from ``DYNAMO_GRAPH_PARALLEL_*``, capped at trees. + + 0 = auto (``min(cpu_count - 1, DYNAMO_GRAPH_PARALLEL_AUTO_MAX_WORKERS)``); + a positive value pins it. Always capped at ``item_count`` (the tree count) + so a two-tree capture never spawns sixteen idle workers. + """ + from aiperf.common.environment import Environment + + configured = Environment.DATASET.DYNAMO_GRAPH_PARALLEL_WORKERS + if configured > 0: + resolved = configured + else: + cpu = os.cpu_count() or 1 + resolved = min( + max(cpu - 1, 1), + Environment.DATASET.DYNAMO_GRAPH_PARALLEL_AUTO_MAX_WORKERS, + ) + return max(1, min(resolved, item_count)) + + +def _io_threads() -> int: + """Thread count for the parallel decompression rounds (scan + shuffle). + + gzip decompression releases the GIL, so threads parallelize the segment + reads without a process pool. Capped at the CPU count (and, inside each + round, at the segment count) -- more decompress threads than cores buys + nothing once every segment already owns one. + """ + return max(1, os.cpu_count() or 1) + + +def _handle_budget() -> int: + """Max simultaneously-open temp-file writers, from the fd soft limit. + + Phase 2 keeps one gzip writer open per batch; capping the batch count at a + fraction of ``RLIMIT_NOFILE`` (leaving headroom for the input read, the pool, + stdio and worker fds) keeps a capture with tens of thousands of trees from + exhausting file descriptors. ``resource`` is Unix-only (the forkserver pool + is Linux/macOS anyway); a conservative constant is used where it is absent. + """ + try: + import resource + + soft, _hard = resource.getrlimit(resource.RLIMIT_NOFILE) + except (ImportError, ValueError, OSError): + return 256 + if soft <= 0: + return 256 + return max(8, soft - 256) + + +def _contiguous_weight_batches( + items: Sequence[_T], weights: Sequence[int], *, num_batches: int +) -> list[list[_T]]: + """Split ``items`` into <= ``num_batches`` CONTIGUOUS weight-balanced batches. + + Contiguous (order-preserving) so a downstream in-order union reproduces the + global item order exactly. Each batch accumulates items until its cumulative + weight crosses the next ``k * total / num_batches`` boundary, always + reserving at least one item for every still-unopened batch, so batches carry + roughly equal total weight; a single item heavier than the target closes its + batch alone (items are atomic). ``weights`` must be positive (the caller + floors them at 1) so a zero-weight run cannot collapse every item into one + batch. + """ + if not items: + return [] + n = max(1, min(num_batches, len(items))) + if n == 1: + return [list(items)] + + total = sum(weights) + batches: list[list[_T]] = [] + current: list[_T] = [] + cumulative = 0 + for index, (item, weight) in enumerate(zip(items, weights, strict=True)): + current.append(item) + cumulative += weight + batches_unopened = n - len(batches) - 1 + items_left = len(items) - (index + 1) + boundary = (len(batches) + 1) * total / n + crossed = cumulative >= boundary + room_for_rest = items_left >= batches_unopened + if batches_unopened >= 1 and crossed and room_for_rest: + batches.append(current) + current = [] + if current: + batches.append(current) + return batches + + +# --- Phase 1: cheap grouping scan (no hash parse) -------------------------- + + +@dataclass(slots=True) +class _GroupingScan: + """Result of the parent's hash-free grouping scan over every segment. + + ``request_end_sessions`` are the sessions that carry at least one + ``request_end`` (the serial build's ``chains`` keys); ``parent_link`` maps a + session to its first non-self parent (``parent_trajectory_id`` preferred over + ``parent_session_id``, first occurrence wins -- identical to + :func:`~aiperf.dataset.graph.adapters.dynamo.trace._collect_records`); + ``session_weight`` is the summed record-line byte length per session (a cheap + proxy for the recorded hash volume the worker will parse, needing no hash + parse); ``block_size`` is the single recorded ``trace_block_size`` across the + whole capture (fail-loud on a mix), defaulting to the virtual size when no + replay is present. + """ + + request_end_sessions: set[str] = field(default_factory=set) + parent_link: dict[str, str] = field(default_factory=dict) + session_weight: dict[str, int] = field(default_factory=dict) + session_peak: dict[str, int] = field(default_factory=dict) + block_sizes: set[int] = field(default_factory=set) + block_size: int = DEFAULT_VIRTUAL_BLOCK_SIZE + + +def _open_raw(path: Path) -> IO[bytes]: + """Open a segment for raw BYTE line iteration (gzip transparently).""" + if path.suffix.lower() == ".gz": + return gzip.open(path, "rb") + return path.open("rb") + + +def _line_peak_context(line: bytes, end: int) -> int: + """Hash-free ``input_length(+output)`` peak for ONE request_end line prefix. + + Mirrors + :func:`~aiperf.dataset.graph.adapters.shared.peak_context.dynamo_tree_peak_context` + exactly: ``replay.input_length`` when present, else ``input_tokens`` (a + recorded 0 is treated as absent -> 1, matching the helper's truthiness + fallback), else 1; plus ``output_tokens`` or 0. Searches only the prefix + before ``input_sequence_hashes`` so no hash array is scanned. + """ + il_match = _INPUT_LENGTH_RE.search(line, 0, end) + if il_match is not None: + input_length = int(il_match.group(1)) + else: + it_match = _INPUT_TOKENS_RE.search(line, 0, end) + input_length = int(it_match.group(1)) if it_match is not None else 0 + if input_length == 0: + input_length = 1 + ot_match = _OUTPUT_TOKENS_RE.search(line, 0, end) + output_tokens = int(ot_match.group(1)) if ot_match is not None else 0 + return input_length + output_tokens + + +def _scan_one_file(segment: Path, *, capture_peak: bool = False) -> _GroupingScan: + """Scan ONE segment for grouping fields, NO hash parse (thread worker). + + For each raw line the search is bounded to the bytes BEFORE + ``"input_sequence_hashes"`` (found with one cheap ``bytes.find``), so the + multi-kilobyte hash array is never scanned by a regex. A line with no + ``session_id`` in that prefix has no ``agent_context`` and is ignored -- the + serial path drops those records the same way. Runs under a thread: gzip + decompression releases the GIL, so N segments scan concurrently. + + ``capture_peak`` (only when a selection knob is set) additionally records the + per-session peak context over its ``request_end`` records via + :func:`_line_peak_context`, so the parent can screen trees by + ``--max-context-length`` / cap them by ``--num-dataset-entries`` before the + build. It is OFF by default so the knob-less scan pays no extra regex cost. + """ + partial = _GroupingScan() + session_weight: dict[str, int] = defaultdict(int) + session_peak: dict[str, int] = defaultdict(int) + with _open_raw(segment) as handle: + for line in handle: + cut = line.find(_HASHES_KEY) + end = len(line) if cut < 0 else cut + match = _SESSION_RE.search(line, 0, end) + if match is None: + continue + sid = match.group(1).decode() + session_weight[sid] += len(line) + if line.find(_REQUEST_END_MARK, 0, end) != -1: + partial.request_end_sessions.add(sid) + block_match = _BLOCK_SIZE_RE.search(line, 0, end) + if block_match is not None: + partial.block_sizes.add(int(block_match.group(1))) + if capture_peak: + peak = _line_peak_context(line, end) + if peak > session_peak[sid]: + session_peak[sid] = peak + if sid not in partial.parent_link: + parent_match = _PARENT_TRAJ_RE.search( + line, 0, end + ) or _PARENT_SESSION_RE.search(line, 0, end) + if parent_match is not None: + parent = parent_match.group(1).decode() + if parent and parent != sid: + partial.parent_link[sid] = parent + partial.session_weight = dict(session_weight) + partial.session_peak = dict(session_peak) + return partial + + +def _scan_grouping( + path: str | Path, *, threads: int, capture_peak: bool = False +) -> _GroupingScan: + """Phase 1: extract session ids + parent links + block size, NO hash parse. + + Scans every segment in parallel across a thread pool (gzip decompression + releases the GIL) and merges the per-segment partials IN SEGMENT ORDER, so + the ``parent_link`` "first non-self parent wins" resolution is identical to + the serial reader's global left-to-right scan. Block sizes are collected only + from ``request_end`` records carrying a session (matching ``_resolve_block_size`` + over the built chains); a mix fails loud. + + ``capture_peak`` propagates to each :func:`_scan_one_file` so the per-session + peak context is aggregated (max across segments -- a session's records may + span files) for the optional filter-then-cap tree selection. + """ + segments = discover_segments(Path(path)) + workers = max(1, min(threads, len(segments))) + scan_fn = functools.partial(_scan_one_file, capture_peak=capture_peak) + if workers == 1: + partials = [scan_fn(seg) for seg in segments] + else: + with ThreadPoolExecutor(max_workers=workers) as pool: + partials = list(pool.map(scan_fn, segments)) + + merged = _GroupingScan() + session_weight: dict[str, int] = defaultdict(int) + session_peak: dict[str, int] = defaultdict(int) + block_sizes: set[int] = set() + # Segment order (ThreadPoolExecutor.map preserves input order) -> first + # non-self parent from the earliest segment wins, as the serial scan does. + for partial in partials: + merged.request_end_sessions |= partial.request_end_sessions + for sid, parent in partial.parent_link.items(): + if sid not in merged.parent_link: + merged.parent_link[sid] = parent + for sid, weight in partial.session_weight.items(): + session_weight[sid] += weight + for sid, peak in partial.session_peak.items(): + if peak > session_peak[sid]: + session_peak[sid] = peak + block_sizes |= partial.block_sizes + + if len(block_sizes) > 1: + raise DynamoTraceAdapterError( + f"mixed replay trace_block_size values are not supported: " + f"{sorted(block_sizes)}" + ) + merged.session_weight = dict(session_weight) + merged.session_peak = dict(session_peak) + merged.block_size = next(iter(block_sizes), 0) or DEFAULT_VIRTUAL_BLOCK_SIZE + return merged + + +# --- Phase 2: shuffle raw lines to per-batch temp files -------------------- + + +_SHUFFLE_QUEUE_MAXSIZE = 8192 +_SHUFFLE_SENTINEL = object() + + +def _shuffle_to_batch_files( + path: str | Path, + session_to_batch: dict[str, int], + tmpdir: Path, + *, + threads: int, +) -> list[Path]: + """Phase 2: append each record line VERBATIM to its batch's gzip temp file. + + Producer threads decompress the segments in parallel (gzip releases the GIL) + and hand each routed line ``(batch, bytes)`` to ONE consumer thread that owns + every batch writer -- so the open-file-descriptor count is bounded by the + batch count (not multiplied by the thread count) and no two threads ever + write the same file. Streaming through a bounded queue keeps resident memory + flat (never the whole re-sharded corpus). A line's session id (same + prefix-bounded regex as the scan) selects its batch; lines with no session id + (schema-less markers, replay-only records) or a session with no request_end + are dropped -- exactly the records the serial path never lowers. Verbatim + bytes mean the worker re-reads them through the unchanged reader (envelope + unwrap, dedup, interning all happen there), so hash parsing stays entirely + worker-side. Returns the batch temp-file paths in batch-index order + (contiguous global tree order). + """ + segments = discover_segments(Path(path)) + line_queue: queue.Queue[Any] = queue.Queue(maxsize=_SHUFFLE_QUEUE_MAXSIZE) + files: dict[int, Path] = {} + + def _consume() -> None: + writers: dict[int, IO[bytes]] = {} + try: + while True: + item = line_queue.get() + if item is _SHUFFLE_SENTINEL: + return + batch, data = item + writer = writers.get(batch) + if writer is None: + fpath = tmpdir / f"batch_{batch:05d}.jsonl.gz" + # Long-lived per-batch writer, closed in the finally below; + # a per-line context manager would reopen (and re-header) + # the gzip stream on every record. + writer = gzip.open(fpath, "wb", compresslevel=1) # noqa: SIM115 + writers[batch] = writer + files[batch] = fpath + writer.write(data) + finally: + for writer in writers.values(): + writer.close() + + def _produce(segment: Path) -> None: + with _open_raw(segment) as handle: + for line in handle: + cut = line.find(_HASHES_KEY) + end = len(line) if cut < 0 else cut + match = _SESSION_RE.search(line, 0, end) + if match is None: + continue + batch = session_to_batch.get(match.group(1).decode()) + if batch is None: + continue + line_queue.put((batch, line if line.endswith(b"\n") else line + b"\n")) + + consumer = threading.Thread(target=_consume) + consumer.start() + try: + workers = max(1, min(threads, len(segments))) + with ThreadPoolExecutor(max_workers=workers) as pool: + list(pool.map(_produce, segments)) + finally: + line_queue.put(_SHUFFLE_SENTINEL) + consumer.join() + return [files[batch] for batch in sorted(files)] + + +# --- Phase 3: fused per-batch read+build in workers ------------------------ + + +def _build_batch_file_to_blob(task: tuple[str, dict[str, Any], int]) -> bytes: + """Pool worker entry: READ one batch temp file and BUILD it, return a frame. + + The FUSED step: the recorded hashes are parsed here (worker-side), never in + the parent. Regroups the file's chains into trees and calls the SERIAL + per-tree loop directly (never ``from_dynamo_trace`` / ``_build_trees_flat``, + which would re-enter the parallel dispatch and spawn a NESTED pool), so the + inner build is strictly single-threaded. The shuffle keeps every session-tree + BATCH-LOCAL, so this worker always sees COMPLETE trees and returns its batch's + LIST of per-tree single-graph ``ParsedGraph`` blobs (the parent flattens all + workers' per-tree graphs and merges). ``direct_store=None`` -- the parallel + path never carries a live write-through store. + """ + from aiperf.dataset.graph.adapters.dynamo.trace import ( + _build_trees_sequential, + _collect_chains, + group_chains_into_trees, + ) + + file_path, build_kwargs, max_depth = task + chains = _collect_chains(file_path, None, max_depth=max_depth) + parent_link = { + sid: chain.parent_session_id + for sid, chain in chains.items() + if chain.parent_session_id is not None + } + trees = group_chains_into_trees(chains, parent_link) + per_tree = _build_trees_sequential(trees, direct_store=None, **build_kwargs) + return _encode_batch_result(per_tree) + + +def _encode_batch_result(per_tree: list[ParsedGraph]) -> bytes: + """Encode a batch's per-tree ``ParsedGraph``s to a cross-process list frame. + + Each per-tree ``ParsedGraph`` (one root ``TraceRecord`` + that tree's graph + and pool) is msgpack-encoded through the existing typed codec; the outer + frame is the plain list of those blobs in the batch's tree order. + """ + return _FRAME_ENCODER.encode([encode_parsed_graph_msgpack(pg) for pg in per_tree]) + + +def _decode_batch_result(blob: bytes) -> list[ParsedGraph]: + """Decode a per-tree ``ParsedGraph`` list frame back into ``ParsedGraph``s.""" + pg_blobs = _FRAME_DECODER.decode(blob) + return [decode_parsed_graph_msgpack(pg_bytes) for pg_bytes in pg_blobs] + + +# --- orchestration --------------------------------------------------------- + + +def maybe_build_fused_parallel( + path: str | Path, + *, + content_root_seed: int, + idle_gap_cap_seconds: float | None, + content_tokenizer: str | None, + prompt_corpus: str, + release_replay: bool, + max_depth: int, + num_dataset_entries: int | None = None, + max_context_length: int | None = None, + selection_out: list[SelectionStats] | None = None, +) -> list[ParsedGraph] | None: + """Fuse read+build across a process pool, or ``None`` to stay serial. + + Runs the Phase-1 grouping scan (cheap, no hash parse) to learn the tree + count, then returns ``None`` -- caller runs the serial read-then-build, NO + pool spawn -- when the tree count is at or below + ``DYNAMO_GRAPH_PARALLEL_THRESHOLD`` or the resolved worker count collapses to + 1 (a single tree, or ``DYNAMO_GRAPH_PARALLEL_WORKERS=1``). Otherwise shuffles + the raw lines to per-batch temp files and builds every batch on the pool, + returning the flattened LIST of per-tree single-graph ``ParsedGraph``s (in + contiguous global tree order) -- byte-identical to + :func:`~aiperf.dataset.graph.adapters.dynamo.trace._build_trees_sequential` + over the same capture. The caller folds the base tag and merges them via + :func:`~aiperf.dataset.graph.merge.merge_parsed_graphs`. + + ``num_dataset_entries`` / ``max_context_length`` drive the SAME filter-then-cap + tree selection the serial path applies (ai-dynamo/aiperf#1106): the scan + additionally records per-session peak context (hash-free), trees are screened + by ``--max-context-length`` and capped at ``--num-dataset-entries`` (root-sorted + order), and ONLY the selected trees are shuffled + built. When both are ``None`` + no selection runs and the output stays byte-identical. ``selection_out`` + receives the :class:`SelectionStats` only when the pool path actually builds + (a decline hands selection to the serial path, which appends its own). + """ + select = num_dataset_entries is not None or max_context_length is not None + io_threads = _io_threads() + scan = _scan_grouping(path, threads=io_threads, capture_peak=select) + if not scan.request_end_sessions: + # No lowerable records: let the serial path raise the precise + # EmptyDynamoTraceError (it distinguishes "empty" from "no session + # identity" via its skip counter, which this scan does not track). + return None + + from aiperf.dataset.graph.adapters.dynamo.trace import root_of_sessions + + root_of = root_of_sessions(scan.request_end_sessions, scan.parent_link) + roots = sorted(set(root_of.values())) + + stats: SelectionStats | None = None + if select: + roots, root_of, stats = _select_roots_filter_then_cap( + scan, + root_of=root_of, + roots=roots, + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + ) + if not roots: + # Every tree filtered out: defer to the serial path so it raises the + # precise EmptyDynamoTraceError (and appends the authoritative stats). + return None + + if len(roots) <= _dynamo_threshold(): + return None + workers = _dynamo_workers(item_count=len(roots)) + if workers <= 1: + return None + + if stats is not None: + # Parent-side finalize point for the fused BUILD path: the scan ran once + # in this parent process, so log the summary once here (a fused DECLINE + # returns above and hands logging to the serial fallback). + log_selection_summary( + stats, + source=str(path), + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + ) + if selection_out is not None: + selection_out.append(stats) + return _build_fused_parallel( + path, + scan=scan, + root_of=root_of, + roots=roots, + workers=workers, + io_threads=io_threads, + content_root_seed=content_root_seed, + idle_gap_cap_seconds=idle_gap_cap_seconds, + content_tokenizer=content_tokenizer, + prompt_corpus=prompt_corpus, + release_replay=release_replay, + max_depth=max_depth, + ) + + +def _select_roots_filter_then_cap( + scan: _GroupingScan, + *, + root_of: dict[str, str], + roots: list[str], + num_dataset_entries: int | None, + max_context_length: int | None, +) -> tuple[list[str], dict[str, str], SelectionStats]: + """Filter-then-cap the tree roots by per-tree peak context (parent-side). + + A tree's peak is the max over its sessions' scan-recorded peaks; the roots + are screened in sorted order (deterministic), so the selected set matches the + serial path's :func:`dynamo_tree_peak_context` selection. Returns the + selected roots (sorted), the ``root_of`` map restricted to the selected + trees' sessions, and the :class:`SelectionStats`. + """ + from aiperf.dataset.graph.adapters.shared.selection import filter_then_cap + + tree_peak: dict[str, int] = defaultdict(int) + for sid, root in root_of.items(): + peak = scan.session_peak.get(sid, 0) + if peak > tree_peak[root]: + tree_peak[root] = peak + + selected, stats = filter_then_cap( + ((root, tree_peak.get(root, 0)) for root in roots), + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + ) + selected_set = set(selected) + restricted = {sid: root for sid, root in root_of.items() if root in selected_set} + return selected, restricted, stats + + +def _build_fused_parallel( + path: str | Path, + *, + scan: _GroupingScan, + root_of: dict[str, str], + roots: list[str], + workers: int, + io_threads: int, + content_root_seed: int, + idle_gap_cap_seconds: float | None, + content_tokenizer: str | None, + prompt_corpus: str, + release_replay: bool, + max_depth: int, +) -> list[ParsedGraph]: + """Shuffle to per-batch temp files, fuse-build on the pool, flatten per-tree. + + ``content_root_seed`` and ``block_size`` are PINNED by the parent and + threaded to every worker (the seed also seeds ``_run_pool_streaming``'s + parent-built shared-memory corpus), so the fused build is byte-identical to + the serial loop. Trees are batched CONTIGUOUSLY over the sorted-by-root list + (weighted by the Phase-1 byte proxy), so batch results arriving in input + order flatten to the same global tree order. Each worker returns its batch's + LIST of per-tree ``ParsedGraph`` blobs; this concatenates them across batches + (arrival order) and hands the flat list back for the caller to merge. The + temp dir is always removed. + """ + from aiperf.common.environment import Environment + from aiperf.dataset.graph.adapters.weka.trace_parallel import _run_pool_streaming + + prefetch = Environment.DATASET.DYNAMO_GRAPH_PARALLEL_PREFETCH_MULTIPLIER + item_timeout_s = Environment.DATASET.DYNAMO_GRAPH_PARALLEL_ITEM_TIMEOUT_SECONDS + + # Iterate ``root_of`` (not ``scan.request_end_sessions``): under selection it + # is restricted to the kept trees' sessions, so a dropped session never + # KeyErrors here. Unfiltered, ``root_of`` keys == ``request_end_sessions``, + # so this is byte-identical to the prior weighting. + tree_weight: dict[str, int] = defaultdict(int) + for sid, root in root_of.items(): + tree_weight[root] += scan.session_weight.get(sid, 0) + weights = [max(1, tree_weight[root]) for root in roots] + + num_batches = max(1, min(len(roots), workers * prefetch, _handle_budget())) + root_batches = _contiguous_weight_batches(roots, weights, num_batches=num_batches) + root_to_batch = { + root: index for index, group in enumerate(root_batches) for root in group + } + session_to_batch = {sid: root_to_batch[root_of[sid]] for sid in root_of} + + build_kwargs: dict[str, Any] = { + "block_size": scan.block_size, + "content_root_seed": content_root_seed, + "idle_gap_cap_seconds": idle_gap_cap_seconds, + "content_tokenizer": content_tokenizer, + "prompt_corpus": prompt_corpus, + "release_replay": release_replay, + } + + tmpdir = Path(tempfile.mkdtemp(prefix="aiperf-dynamo-fused-")) + try: + batch_files = _shuffle_to_batch_files( + path, session_to_batch, tmpdir, threads=io_threads + ) + tasks = ((str(bf), build_kwargs, max_depth) for bf in batch_files) + + per_tree: list[ParsedGraph] = [] + # Results arrive in INPUT (batch) order -> contiguous global tree order, + # so the flattened per-tree list matches the serial per-tree loop and the + # merge is byte-identical (traces are id-sorted, pools content-unioned). + for blob in _run_pool_streaming( + _build_batch_file_to_blob, + tasks, + workers=workers, + root_seed=content_root_seed, + content_tokenizer=content_tokenizer, + prompt_corpus=prompt_corpus, + prefetch_multiplier=prefetch, + item_timeout_s=item_timeout_s, + timeout_hint=_TIMEOUT_HINT, + ): + per_tree.extend(_decode_batch_result(blob)) + + return per_tree + finally: + shutil.rmtree(tmpdir, ignore_errors=True) + + +__all__ = [ + "maybe_build_fused_parallel", +] diff --git a/src/aiperf/dataset/graph/adapters/dynamo/trace_reader.py b/src/aiperf/dataset/graph/adapters/dynamo/trace_reader.py new file mode 100644 index 0000000000..f921bc71f7 --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/dynamo/trace_reader.py @@ -0,0 +1,444 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Typed reader for Dynamo request-trace files (`*.jsonl`, `*.jsonl.gz`, segmented). + +Schema source (dynamo repo): lib/llm/src/request_trace/types.rs + and lib/llm/src/protocols/common/extensions.rs (AgentContext) + +Dynamo's file sinks (`jsonl` and `jsonl_gz`, selected by +`DYN_REQUEST_TRACE_SINKS`, gated by `DYN_REQUEST_TRACE`) write ONE envelope +per line: `{"timestamp": , "event": {}}` +(`telemetry/jsonl_gz.rs::GzipEntry`); only the stderr sink emits the bare +record. `iter_raw_records` unwraps the envelope and also accepts bare records +(hand-authored fixtures). With the jsonl_gz sink, segments roll into files +like `prefix.000000.jsonl.gz`, `prefix.000001.jsonl.gz`; each gzip member is +a complete batch. + +Usage: + for record in iter_trace_records("/tmp/dynamo-trace"): + # iterates plain JSONL OR all .000000-.NNNNNN.jsonl.gz segments + ... + + for record in iter_trace_records( + "/tmp/dynamo-trace", + event_types={"request_end"}, + session_id="session-42", + ): + ... + +Note: in the current `dynamo.request.trace.v1` schema, `event_source`, +`agent_context`, and `request.model` are all optional (Option in Rust) and may +be absent on replay-only records, so the models below make them optional too. +""" + +from __future__ import annotations + +import gzip +import io +import re +import zlib +from collections.abc import Iterable, Iterator +from pathlib import Path +from typing import Any, Literal + +import orjson +from pydantic import BaseModel, ConfigDict, Field, field_validator + +from aiperf.common.finite import FiniteFloat + +# Wire discriminator for the current Dynamo trace schema. Imported by the sniff +# in ``dynamo/trace.py``; the ``AgentTraceRecord.schema_`` Literal below carries +# the same value for typed decode (a Literal cannot reference a variable). +DYNAMO_TRACE_SCHEMA_V1 = "dynamo.request.trace.v1" + + +class DynamoTraceAdapterError(ValueError): + """Raised when a Dynamo trace cannot be converted into a ParsedGraph.""" + + +class EmptyDynamoTraceError(DynamoTraceAdapterError): + """Raised when the trace file/dir contains no usable records.""" + + +class AgentContext(BaseModel): + model_config = ConfigDict(extra="ignore") + session_id: str = Field(description="Stable reasoning/tool session identifier.") + parent_session_id: str | None = Field( + default=None, + description="Legacy parent-session link for subagents. Real captures " + "leave this unset and carry the linkage in parent_trajectory_id " + "instead; kept as a fallback for older/hand-authored traces.", + ) + trajectory_id: str | None = Field( + default=None, + description="Per-trajectory identifier; equals session_id in every " + "observed capture, so a subagent's parent_trajectory_id resolves to the " + "parent's session_id.", + ) + parent_trajectory_id: str | None = Field( + default=None, + description="Parent trajectory for subagents (the parent's " + "trajectory_id == its session_id). The AUTHORITATIVE subagent linkage in " + "real captures; parent_session_id is only a fallback.", + ) + session_type_id: str | None = Field( + default=None, + description="Coarse session kind recorded by the harness (e.g. " + "'opencode'); carried through for provenance, not used for linkage.", + ) + + +class WorkerInfo(BaseModel): + model_config = ConfigDict(extra="ignore") + prefill_worker_id: int | None = Field( + default=None, description="Prefill worker id when recorded." + ) + prefill_dp_rank: int | None = Field( + default=None, description="Prefill DP rank when recorded." + ) + decode_worker_id: int | None = Field( + default=None, description="Decode worker id when recorded." + ) + decode_dp_rank: int | None = Field( + default=None, description="Decode DP rank when recorded." + ) + + +class AgentReplayMetrics(BaseModel): + """KV-cache block-hash provenance for replay. + + Emitted unconditionally on every dynamo ``request_end`` at the current + schema (the record is skipped entirely when the model's KV block size is + unavailable, ``request_trace/integration.rs``); absent only on older + captures or hand-authored traces. + """ + + model_config = ConfigDict(extra="ignore") + trace_block_size: int = Field( + ge=1, + description="KV cache block size used to derive replay hashes. Dynamo " + "never emits a replay block with size 0 (it skips the record instead), " + "so a non-positive value is a corrupt or hand-mangled trace.", + ) + input_length: int = Field( + description="Prompt/input token count represented by the replay hashes." + ) + input_sequence_hashes: list[int] = Field( + description="Stable sequence-aware prompt block hashes (u64 in Rust, may be > 2^63)." + ) + + @field_validator("input_sequence_hashes") + @classmethod + def _reject_negative_hashes(cls, hashes: list[int]) -> list[int]: + """Reject negative recorded hashes: they collide with the virtual negative-id namespace. + + Dynamo records these as u64 (values may exceed 2^63, so they stay + positive Python ints). A negative value would clash with the virtual + negative ids the trie lowering mints for non-replay turns, silently + aliasing distinct content. + """ + # `min()` iterates the list in C (u64 values may exceed 2^63, so an + # array('q')/numpy int64 scan would overflow -- keep Python ints). The + # empty guard preserves the prior `any([]) is False` no-raise behavior. + if hashes and min(hashes) < 0: + raise ValueError( + "input_sequence_hashes must be non-negative (u64 in Rust); " + "negative values collide with the virtual negative-id namespace." + ) + return hashes + + +class AgentRequestMetrics(BaseModel): + model_config = ConfigDict(extra="ignore") + request_id: str = Field(description="Dynamo request ID for the LLM call.") + x_request_id: str | None = Field( + default=None, description="Caller-provided logical request ID when present." + ) + model: str | None = Field( + default=None, description="Model name (optional in the current schema)." + ) + input_tokens: int | None = Field( + default=None, description="Prompt/input token count when known." + ) + output_tokens: int | None = Field( + default=None, description="Final output token count when known." + ) + cached_tokens: int | None = Field( + default=None, + description="Prompt tokens served from prefix/KV cache when known.", + ) + request_received_ms: int | None = Field( + default=None, description="Request receive time in Unix epoch milliseconds." + ) + prefill_wait_time_ms: float | None = Field( + default=None, description="Time from request receipt to prefill start." + ) + prefill_time_ms: float | None = Field( + default=None, description="Time from prefill start to first token." + ) + ttft_ms: FiniteFloat | None = Field( + default=None, description="Time from request receipt to first token." + ) + total_time_ms: float | None = Field( + default=None, description="Time from request receipt to request completion." + ) + avg_itl_ms: FiniteFloat | None = Field( + default=None, description="Average inter-token latency after first token." + ) + kv_hit_rate: float | None = Field( + default=None, description="Effective KV-cache hit rate observed by the router." + ) + kv_transfer_estimated_latency_ms: FiniteFloat | None = Field( + default=None, + description="Upper-bound estimated disaggregated KV transfer latency.", + ) + queue_depth: int | None = Field( + default=None, + description="Router queue depth observed when routing the request.", + ) + worker: WorkerInfo | None = Field( + default=None, + description="Prefill/decode worker IDs and DP ranks when recorded.", + ) + replay: AgentReplayMetrics | None = Field( + default=None, description="Text-free replay metadata for Mooncake/mocker." + ) + + +class AgentToolEvent(BaseModel): + model_config = ConfigDict(extra="ignore") + tool_call_id: str = Field(description="Harness-provided tool call identifier.") + tool_class: str = Field(description="Tool class/category name.") + status: ( + Literal[ + "running", + "succeeded", + "ok", + "success", + "error", + "failed", + "cancelled", + "canceled", + "timeout", + ] + | None + ) = Field( + default=None, + description="Terminal status when present. Dynamo-written files only " + "ever contain the canonical serde names (running/succeeded/error/" + "cancelled) -- the Rust aliases ('ok', 'success', 'failed', " + "'canceled', 'timeout'; request_trace/types.rs) are deserialize-only " + "leniency for hand-authored traces and are NOT canonicalized here.", + ) + duration_ms: float | None = Field( + default=None, description="Wall-clock duration in milliseconds." + ) + + +class AgentTraceRecord(BaseModel): + model_config = ConfigDict(extra="ignore", populate_by_name=True) + schema_: Literal["dynamo.request.trace.v1"] = Field( + alias="schema", + description="Trace schema discriminator. Aliased to the wire field 'schema' " + "because it shadows pydantic.BaseModel.schema. Use model_dump(by_alias=True) " + "to round-trip.", + ) + event_type: Literal["request_end", "tool_start", "tool_end", "tool_error"] = Field( + description="Trace event kind." + ) + event_time_unix_ms: int = Field( + description="Event wall-clock time in Unix epoch milliseconds." + ) + event_source: Literal["dynamo", "harness"] | None = Field( + default=None, + description="Producer identity. Optional in the current schema and absent " + "on replay-only records.", + ) + agent_context: AgentContext | None = Field( + default=None, + description="Session identity attached to the event. Optional in the current " + "schema and absent on replay-only records.", + ) + request: AgentRequestMetrics | None = Field( + default=None, description="Populated for request_end events." + ) + tool: AgentToolEvent | None = Field( + default=None, description="Populated for tool_* events." + ) + + +class DynamoTraceReadError(ValueError): + """Raised when a trace file/directory can't be parsed.""" + + +_SEGMENT_PATTERN = re.compile(r"^(.+?)\.(\d{6,})\.jsonl\.gz$") + + +def _dir_segment_sort_key(p: Path) -> tuple[str, int]: + """Numeric segment order within a shared prefix; plain names sort by name. + + Lexicographic name order breaks at the 7-digit rollover + (``1000000`` < ``999999``); dynamo's writer widens the index naturally, so + sort segment-shaped names by ``(prefix, int(index))``. + """ + m = _SEGMENT_PATTERN.match(p.name) + if m is None: + return (p.name, -1) + return (m.group(1), int(m.group(2))) + + +def discover_segments(path: Path) -> list[Path]: + """Return ordered list of `.NNNNNN.jsonl.gz` segments sharing the prefix. + + If `path` is a plain `.jsonl` or `.jsonl.gz` file, returns `[path]`. + If `path` is a directory, returns all sorted `*.jsonl` and `*.jsonl.gz` files + inside (segments share a prefix and sort numerically by segment number). + If `path` matches a segmented prefix (e.g. `/tmp/dynamo-trace`), returns all + segments whose names match `.NNNNNN.jsonl.gz`. Mirroring dynamo's + own segment naming (`telemetry/jsonl_gz.rs::segment_path`), a trailing + `.jsonl.gz` / `.jsonl` on a non-existent prefix path is stripped first, so + the configured `DYN_REQUEST_TRACE_OUTPUT_PATH` value works verbatim. + """ + p = path + if p.is_file(): + return [p] + if p.is_dir(): + out = sorted( + list(p.glob("*.jsonl")) + list(p.glob("*.jsonl.gz")), + key=_dir_segment_sort_key, + ) + if not out: + raise DynamoTraceReadError( + f"{p}: no .jsonl or .jsonl.gz files in directory" + ) + return out + parent = p.parent + if not parent.is_dir(): + raise DynamoTraceReadError(f"{p}: not a file/dir, and parent isn't a directory") + prefix = p.name + for suffix in (".jsonl.gz", ".jsonl"): + if prefix.endswith(suffix): + prefix = prefix[: -len(suffix)] + break + candidates = [ + c for c in parent.glob(f"{prefix}.*.jsonl.gz") if _SEGMENT_PATTERN.match(c.name) + ] + out = sorted( + candidates, + key=lambda x: int(_SEGMENT_PATTERN.match(x.name).group(2)), # type: ignore[union-attr] + ) + if not out: + raise DynamoTraceReadError( + f"{p}: no matching segments found ({prefix}.*.jsonl.gz)" + ) + return out + + +def _open_segment(path: Path) -> io.TextIOBase: + # .lower(): detection sniffing is case-insensitive, so a .JSONL.GZ capture + # that passed can_load must not be opened as utf-8 text here. + if path.suffix.lower() == ".gz": + return io.TextIOWrapper(gzip.open(path, "rb"), encoding="utf-8") + return path.open("rt", encoding="utf-8") + + +def unwrap_sink_envelope(raw: dict[str, Any]) -> dict[str, Any]: + """Unwrap dynamo's file-sink line envelope, passing bare records through. + + Both dynamo file sinks (``jsonl`` and ``jsonl_gz``) wrap every line as + ``{"timestamp": , "event": {}}`` + (``telemetry/jsonl_gz.rs::GzipEntry``, ``recorder.rs``); only the stderr + sink emits the bare record. Bare records (hand-authored fixtures, older + captures) are detected by their required top-level ``schema`` key. + """ + if "schema" not in raw and isinstance(raw.get("event"), dict): + return raw["event"] + return raw + + +def iter_raw_records(path: Path | str) -> Iterator[dict[str, Any]]: + """Stream raw JSON dicts (sink envelope unwrapped) from a file/prefix/dir.""" + p = Path(path) + for segment in discover_segments(p): + with _open_segment(segment) as f: + try: + for raw_line in f: + line = raw_line.strip() + if not line: + continue + try: + rec = unwrap_sink_envelope(orjson.loads(line)) + except orjson.JSONDecodeError as e: + raise DynamoTraceReadError( + f"{segment}: invalid JSON line: {e}" + ) from e + # Skip non-record marker lines the S3 uploader sidecar appends + # (e.g. a trailing {"verification": "trace-s3-uploader"}): a real + # dynamo record always carries the schema discriminator, so a + # schema-less line is a control marker, not a dropped record. + if "schema" not in rec: + continue + yield rec + except (EOFError, gzip.BadGzipFile, zlib.error) as e: + # A SIGKILL mid-append leaves a truncated final gzip member + # (EOFError); bit rot or a partial copy corrupts the deflate + # stream (zlib.error); a non-gzip file behind a .gz name fails + # the header check (BadGzipFile). + raise DynamoTraceReadError( + f"{segment}: truncated or corrupt gzip stream (capture " + f"interrupted mid-flush, partial copy, or not gzip?): {e}" + ) from e + except UnicodeDecodeError as e: + raise DynamoTraceReadError( + f"{segment}: not valid UTF-8 JSONL (binary or gzip bytes " + f"behind a .jsonl name?): {e}" + ) from e + + +def iter_trace_records( + path: Path | str, + *, + event_types: Iterable[str] | None = None, + session_id: str | None = None, + time_range_ms: tuple[int, int] | None = None, +) -> Iterator[AgentTraceRecord]: + """Stream typed `AgentTraceRecord` from a trace file/dir/prefix. + + Filters apply at parse time so unwanted records skip Pydantic validation. + `time_range_ms` is inclusive on both bounds. + """ + et_set = set(event_types) if event_types is not None else None + lo, hi = time_range_ms if time_range_ms is not None else (None, None) + for raw in iter_raw_records(path): + if et_set is not None and raw.get("event_type") not in et_set: + continue + ts = raw.get("event_time_unix_ms") + if lo is not None and (not isinstance(ts, int) or ts < lo): + continue + if hi is not None and (not isinstance(ts, int) or ts > hi): + continue + ac = raw.get("agent_context") or {} + if session_id is not None and ac.get("session_id") != session_id: + continue + try: + yield AgentTraceRecord.model_validate(raw) + except Exception as e: + raise DynamoTraceReadError( + f"failed to parse trace record: {e!s} (raw keys: {sorted(raw.keys())})" + ) from e + + +__all__ = [ + "AgentContext", + "AgentReplayMetrics", + "AgentRequestMetrics", + "AgentToolEvent", + "AgentTraceRecord", + "DynamoTraceAdapterError", + "DynamoTraceReadError", + "EmptyDynamoTraceError", + "WorkerInfo", + "discover_segments", + "iter_raw_records", + "iter_trace_records", +] diff --git a/src/aiperf/dataset/graph/adapters/dynamo/trie_lowering.py b/src/aiperf/dataset/graph/adapters/dynamo/trie_lowering.py new file mode 100644 index 0000000000..da38024016 --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/dynamo/trie_lowering.py @@ -0,0 +1,465 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Normalize dynamo trace chains into the shared segment-trie IR shape. + +One ``TrieNode`` per ``request_end``; the hash source is the recorded +``input_sequence_hashes`` when a record carries replay metadata, else +per-session VIRTUAL negative ids (tagged ``virtual-hash-fallback``) -- the +same recorded-when-present rule as the weka path. Content sharing, frozen +roles, edges, and the segment pool are all owned by +:func:`segment_ir.trie_content.build_trie_ir` -- this module only maps dynamo +record fields onto :class:`TrieRequest` and assembles the dynamo-flavored +``LlmNode``. +""" + +from __future__ import annotations + +import itertools +from collections.abc import Iterable +from typing import TYPE_CHECKING, Any + +from aiperf.common.constants import MICROS_PER_SECOND +from aiperf.dataset.graph.adapters.dynamo.trace_reader import ( + DynamoTraceAdapterError, +) +from aiperf.dataset.graph.adapters.shared.output_cap import wire_output_cap +from aiperf.dataset.graph.models import ExpectedTokens, LlmNode +from aiperf.dataset.graph.segment_ir.envelope import stamp_prompt_segment_ids +from aiperf.dataset.graph.segment_ir.trie_content import ( + ReconCallbacks, + TrieNode, + TrieNodeBuild, + TrieRequest, +) + +if TYPE_CHECKING: + from aiperf.dataset.graph.adapters.dynamo.trace import _Chain, _Turn + +DEFAULT_VIRTUAL_BLOCK_SIZE = 16 +_VIRTUAL_HASH_FALLBACK_TAG = "virtual-hash-fallback" + + +class DynamoISLMismatchError(ValueError): + """Raised when a recorded ``input_length`` disagrees with its replay hashes. + + Dynamo records ``input_length`` (the full prompt token count) alongside + ``input_sequence_hashes`` (full-block hashes plus one partial tail over that + same prompt). A block-consistent record satisfies + ``(n_hashes - 1) * block_size < input_length <= n_hashes * block_size``; a + divergence means the trace metadata is internally inconsistent and no + reconstruction can honor both fields -- the parse aborts fail-loud rather + than materialize a wrong-sized prompt. + """ + + +def _guard_session_scopes(session_ids: Iterable[str]) -> None: + """Fail loud on session ids the ``{scope}:{turn}`` grammar cannot carry. + + Node ids are ``{session_id}:{k}`` with the RECORDED session id verbatim + (no hashing -- the id is the data), so ``:`` -- the one reserved identity + separator (runtime ids append ``::{nonce}``) -- must not appear inside a + recorded session id. + """ + for sid in session_ids: + if not sid or ":" in sid: + raise ValueError( + f"session id {sid!r} cannot form node ids: session ids must " + "be non-empty and contain no ':' (the reserved identity " + "separator)" + ) + + +def _turn_times_ms(turn: _Turn) -> tuple[int, int]: + """(start_ms, total_ms) mirroring upstream agent_trace_to_mooncake.""" + rec = turn.record + req = rec.request + event = rec.event_time_unix_ms + received = req.request_received_ms if req is not None else None + total_f = req.total_time_ms if req is not None else None + if total_f is not None: + total = max(0, round(total_f)) + elif received is not None: + total = max(0, event - received) + else: + total = 0 + start = received if received is not None else event - total + return start, total + + +def _assert_block_aligned( + node_id: str, hashes: list[int], input_length: int, bs: int +) -> None: + """Fail-loud gate: ``input_length`` must be spanned by its replay hashes. + + Dynamo generates full-block hashes plus one partial tail + (``request_trace/replay.rs``), so a block-consistent record satisfies + ``(n-1)*bs < input_length <= n*bs``. + """ + n = len(hashes) + lo, hi = (n - 1) * bs, n * bs + if not (lo < input_length <= hi): + raise DynamoISLMismatchError( + f"dynamo node {node_id!r}: recorded input_length={input_length} is not " + f"block-aligned to its {n} replay hashes at block_size={bs} " + f"(expected {lo} < input_length <= {hi})." + ) + + +def _resolve_block_size(chains: dict[str, _Chain]) -> int | None: + """Single recorded ``trace_block_size`` across all replay turns; fail-loud on mix.""" + block_size: int | None = None + for chain in chains.values(): + for turn in chain.turns: + req = turn.record.request + if req is None or req.replay is None: + continue + tbs = req.replay.trace_block_size + if block_size is None: + block_size = tbs + elif tbs != block_size: + raise DynamoTraceAdapterError( + f"mixed replay trace_block_size values are not " + f"supported: {block_size} and {tbs}" + ) + return block_size + + +def dynamo_trie_nodes( + chains: dict[str, _Chain], + *, + release_replay: bool = False, + block_size: int | None = None, +) -> tuple[list[TrieNode], int, list[str]]: + """Flatten every session's turns into shared TrieNodes. + + Returns ``(nodes, block_size, extra_trace_tags)``. Nodes are ``{session_id}:{k}`` + (recorded session id verbatim + 0-based turn), ordered by + recorded start time (ties by ``(session_id, turn_index)``) -- the + deterministic ``order`` the content-parent trie resolves in. + ``block_size`` is the single recorded ``trace_block_size`` when any record + carries replay metadata, else :data:`DEFAULT_VIRTUAL_BLOCK_SIZE`. + + ``block_size`` (keyword) pins the block size instead of resolving it from + these chains. The per-session-tree build (:func:`from_dynamo_trace`) + resolves ONE size across the WHOLE capture (fail-loud on a mix) and pins it + into every tree so a tree that happens to carry no replay never silently + drops to the virtual default while a sibling tree carries a recorded size -- + keeping the flat graph byte-identical to a single global build. ``None`` + (every other caller) resolves it from ``chains`` exactly as before. + + Recorded hash objects arrive already interned to canonical ints from + :func:`~aiperf.dataset.graph.adapters.dynamo.trace._collect_records`, so + the recorded branch copies them with a plain ``list()`` (not ``int(h)``), + preserving that shared identity into every ``TrieRequest.hash_ids``. + + ``release_replay`` (opt-in) nulls each record's ``request.replay`` right + after its ``input_sequence_hashes`` / ``input_length`` have been copied into + the emitted :class:`TrieRequest`, freeing the pydantic replay hash lists the + ``chains`` otherwise pin for the whole build. It is OFF by default because + the ``recorded is None`` branch below is the LEGITIMATE virtual-hash + fallback, not an error: re-lowering the SAME in-memory ``chains`` after a + release would silently degrade every previously-recorded turn to negative + virtual ids (a DIFFERENT, wrong trie), so only a caller that lowers a fresh + ``chains`` exactly once (the production store build, via + :meth:`DynamoTraceAdapter.parse`) may opt in. :func:`_resolve_block_size` + has already consumed every replay before this loop runs. + """ + _guard_session_scopes(chains) + if block_size is None: + block_size = _resolve_block_size(chains) + bs = block_size or DEFAULT_VIRTUAL_BLOCK_SIZE + + # Flatten turns with timing, sorted globally by recorded start. The + # tie-break is the NUMERIC (sid, k) pair, never the node-id string: + # ":10" sorts before ":2" lexicographically, which would misorder + # same-millisecond turns past index 9 (order is the trie's ground truth). + flat: list[ + tuple[int, str, int, int, _Chain, _Turn] + ] = [] # (start_ms, sid, k, total_ms, ...) + for sid, chain in sorted(chains.items()): + for k, turn in enumerate(chain.turns): + start_ms, total_ms = _turn_times_ms(turn) + flat.append((start_ms, sid, k, total_ms, chain, turn)) + flat.sort(key=lambda item: (item[0], item[1], item[2])) + t0 = flat[0][0] if flat else 0 + + # Per-session (start_ms, node_id) in start order, for causal-parent + # resolution of subagent first turns. flat is globally start-sorted, so + # per-session sublists inherit that order. + session_turns: dict[str, list[tuple[int, str]]] = {} + for start_ms, sid, k, _total, chain, _turn in flat: + session_turns.setdefault(chain.session_id, []).append((start_ms, f"{sid}:{k}")) + + virtual = itertools.count(-1, -1) + virtual_prev: dict[str, list[int]] = {} + tags: set[str] = set() + nodes: list[TrieNode] = [] + + for order, (start_ms, sid, k, total_ms, chain, turn) in enumerate(flat): + node_id = f"{sid}:{k}" + req = turn.record.request + input_tokens = int(req.input_tokens) if req and req.input_tokens else 1 + output_tokens = int(req.output_tokens) if req and req.output_tokens else 0 + + recorded = req.replay if req is not None else None + if recorded is not None: + # ``list()`` (not ``[int(h) for h in ...]``): the hash objects arrive + # already interned to canonical ints from ``_collect_records``, so a + # plain copy preserves that shared identity (and drops ~H int() C-calls + # per parse). Value-identical -- pydantic guarantees exact ints post- + # validation -- and it still copies, so the release_replay / virtual_prev + # aliasing below is unchanged. + hashes = list(recorded.input_sequence_hashes) + input_length = int(recorded.input_length) + _assert_block_aligned(node_id, hashes, input_length, bs) + # A non-block-aligned input records ONE trailing partial-tail + # hash. Drop it from the lowered hash list: engines cache/share + # FULL blocks only (a partial block is never a prefix-cache hit), + # its tail content is seed-sampled rather than decoded, and + # keeping it would skew block-tag segmentation and theoretical + # block totals against the same recording in other formats. + # ``input_length`` still carries the tail tokens for the sampled + # sub-block remainder. + if input_length < len(hashes) * bs: + hashes.pop() + virtual_prev[chain.session_id] = hashes + if release_replay: + # hashes/input_length are now copied into the TrieRequest below; + # free the recorded replay list (mutable BaseModel) so the whole + # capture's hash lists are not pinned for the rest of the build. + req.replay = None + else: + tags.add(_VIRTUAL_HASH_FALLBACK_TAG) + input_length = input_tokens + prev = virtual_prev.get(chain.session_id, []) + m = input_length // bs + hashes = ( + prev[:m] + if m <= len(prev) + else prev + [next(virtual) for _ in range(m - len(prev))] + ) + virtual_prev[chain.session_id] = hashes + + if k > 0: + causal = f"{sid}:{k - 1}" + elif chain.parent_session_id is not None: + causal = None + for p_start, pid in session_turns.get(chain.parent_session_id, []): + if p_start <= start_ms: + causal = pid + else: + break + else: + causal = None + + node = TrieNode( + node_id=node_id, + request=TrieRequest( + hash_ids=hashes, + input_length=input_length, + output_length=output_tokens, + t=(start_ms - t0) / 1000.0, + api_time=total_ms / 1000.0, + model=req.model if req is not None else None, + ttft=( + (req.ttft_ms / 1000.0) + if (req is not None and req.ttft_ms is not None) + else None + ), + ), + order=order, + causal_parent_id=causal, + ) + node.dynamo_meta = _node_meta(chain, turn, k=k) + nodes.append(node) + + return nodes, bs, sorted(tags) + + +def _node_meta(chain: _Chain, turn: _Turn, *, k: int) -> dict[str, Any]: + """Per-node dynamo payload: session headers, identity breadcrumbs, expected tokens.""" + rec = turn.record + req = rec.request + n = len(chain.turns) + + # Session identity rides in HTTP HEADERS, never the body: dynamo's NvExt + # is deny_unknown_fields and explicitly rejects body-level agent_context + # (extensions.rs::nvext_agent_context_is_rejected); no session_control + # object exists in the protocol. The frontend ingests x-dynamo-session-id + # / x-dynamo-parent-session-id on every request and x-dynamo-session-final + # on the session's LAST request (protocols/agents.rs), which drives + # session-affinity routing and end-of-session KV eviction. A single-turn + # session (n == 1, k == 0) stamps BOTH the session id and session-final on + # its only request, so the session it opens is also evicted. The RECORDED + # ids are stamped verbatim here (build time knows nothing of replay + # concurrency); the worker suffixes the two identity headers per replay + # instance at dispatch (worker_materialize.uniquify_dynamo_session_headers) + # so concurrent instances of one trace never share a server session. + extra_headers: dict[str, str] = {"x-dynamo-session-id": chain.session_id} + if chain.parent_session_id is not None: + extra_headers["x-dynamo-parent-session-id"] = chain.parent_session_id + if k == n - 1: + extra_headers["x-dynamo-session-final"] = "true" + + # No recorded-scalar round-trip is stamped: everything metadata used to + # duplicate lives on native fields (model / max_tokens / expected) or in + # the capture file itself, and node metadata survives the graph_meta + # sidecar strip (which clears only prompt and metadata["trie"]), so every + # extra key here bloats the content-free structural plane at corpus scale. + return { + "extra_headers": extra_headers, + "session_id": chain.session_id, + "parent_session_id": chain.parent_session_id, + "turn_index": k, + "expected": ExpectedTokens( + input_tokens=req.input_tokens if req else None, + output_tokens=req.output_tokens if req else None, + cache_read_tokens=req.cached_tokens if req else None, + cache_creation_tokens=None, + ), + } + + +def dynamo_recon_callbacks( + tokenizer_name: str, + prompt_corpus: str, + root_seed: int | None, + *, + block_size: int, + trace_scope: str, +) -> ReconCallbacks: + """Byte-faithful reconstructor callbacks at the DYNAMO trie block size. + + Reuses the shared (cached) ``CorpusContentSynthesizer`` for the corpus / + tokenizer / partial-tail machinery, but decodes each hash to ``block_size`` + tokens instead of the synthesizer's weka-fixed 64: dynamo records blocks + at the model's KV cache block size (default 16, backend-overridable; or + :data:`DEFAULT_VIRTUAL_BLOCK_SIZE` for virtual fallback ids), and the + covered-count ISL contract requires exactly ``block_size`` tokens per + covered block. A PRIVATE per-parse cache is used because the synthesizer + instance is shared across adapters within a process and its ``pg._cache`` + is keyed by bare hash id -- mixing block sizes in it would hand a 64-token + weka block to a 16-token dynamo decode (or vice versa). The cache stores one + ``int`` corpus OFFSET per hash id (not the decoded token list) via + :meth:`~aiperf.dataset.graph.adapters.shared.content.CorpusContentSynthesizer._decode_block_tokens_offset_cached`, + re-slicing the block on each hit -- the decode-cache tier's largest on-peak + memory shave at corpus scale, byte-identical to the list cache by that + method's full-reseed contract. + + Hash-id scope is deliberately GLOBAL (no ``trace_id`` threading, unlike + the weka path's per-trace ``hash_id_scope: "local"`` namespace): dynamo's + recorded ``input_sequence_hashes`` are chained sequence hashes over the + actual prompt tokens -- equal hash means equal upstream content by + construction, so a single namespace reproduces the recorded sharing + exactly. Virtual ids (negative, per-parse counter) never recur across + sessions within a parse. + """ + from aiperf.dataset.graph.adapters.shared.content import ( + get_or_build_synthesizer, + ) + + synth = get_or_build_synthesizer( + tokenizer_name, prompt_corpus=prompt_corpus, root_seed=root_seed + ) + offsets: dict[int, int] = {} + + # Partial-tail / response seeds are trace-prefixed (weka parity): node ids + # are trajectory-local (``{session}:{k}``) and would otherwise synthesize + # identical tiny-prompt/response bytes for same-shaped nodes of DIFFERENT + # trees -- manufactured cross-trace content sharing. Block decode stays + # bare-hash-id global (equal recorded hash == equal content). + def _tail_scoped(n_tokens: int, seed: str) -> list[int]: + return synth._sample_partial_tail_tokens(n_tokens, f"{trace_scope}:{seed}") + + return ReconCallbacks( + decode_block_tokens=lambda hash_ids: synth._decode_block_tokens_offset_cached( + hash_ids, block_size=block_size, offset_cache=offsets + ), + sample_partial_tail_tokens=_tail_scoped, + decode_tokens_to_text=synth._decode_tokens_to_text, + ) + + +def build_dynamo_llm_node( + node: TrieNode, + *, + build: TrieNodeBuild, +) -> LlmNode: + """Assemble the dynamo-flavored LlmNode from the shared build artifacts. + + The node carries NO inline prompt (``prompt=[]``): on the trie route the + prompt content lives ONLY in the run's ``SegmentPool``, addressed by + ``metadata["trie"]["prompt_segment_ids"]`` (stamped below). The store build, + ``graph_meta`` sidecar (``strip_replay_text`` forces ``prompt=[]``), and the + worker all materialize from the segment pool / mmap store -- none reads + ``node.prompt`` -- matching the dag_jsonl adapter's ``prompt=[]`` convention. + For in-process debugging, recover the content with + ``segment_pool.materialize(read_prompt_segment_ids(node))``. + + Body params ride the NATIVE node fields (``model`` / ``streaming`` / + ``max_tokens``); the envelope builder folds them into the wire body (the + graph-credit path bypasses the endpoint's ``format_payload``). + + Must run AFTER :func:`~aiperf.dataset.graph.segment_ir.trie_content.build_trie_ir`: + ``node.start`` (= ``warped_start``) is only stamped by the driver's idle-warp + pass, so calling this earlier would freeze every ``arrival_offset_us`` at 0. + + Generation is ALWAYS pinned to the recording (weka parity): the native + ``LlmNode.max_tokens`` carries the recorded ``output_tokens`` + (``Turn.max_tokens`` naming), folded into the wire body by the envelope + builder and endpoint-mapped to the wire token field by the worker. + A recorded 0 (zero-output/aborted turn, or a capture without + ``output_tokens``) upgrades to 1 with a warning (:func:`wire_output_cap`) + -- a 0 cap is a meaningless (and for some servers harmful) wire value. + """ + meta = node.dynamo_meta + req = node.request + + # Best-available streaming proxy: dynamo does NOT record the client's + # stream flag; ttft_ms presence only means the request was routed through + # a first-token-recording path (the push router records it regardless of + # client streaming). The native ``streaming`` field reaches the envelope's + # top-level ``stream``, which the worker's per-request stream override + # (apply_run_level_payload_options) honors. + streaming = req.ttft is not None + + llm = LlmNode( + prompt=[], + output=f"{node.node_id}_out", + streaming=streaming, + model=req.model, + max_tokens=wire_output_cap(int(req.output_length), node_id=node.node_id), + # Session identity is header-borne (see _node_meta); the native field + # reaches the envelope and the worker attaches + uniquifies it. + extra_headers=meta["extra_headers"], + arrival_offset_us=int(round(node.start * MICROS_PER_SECOND)), + expected=meta["expected"], + metadata={ + # Session identity breadcrumbs: node ids are HASHES of session ids, + # so this is the only way to map a lowered node back to its + # recorded session. small_prompt marks the covered-count-0 + # reconstruction fallback (not byte-faithful) for ISL audits. + "dynamo": { + "session_id": meta["session_id"], + "parent_session_id": meta["parent_session_id"], + "turn_index": meta["turn_index"], + "small_prompt": build.small_prompt, + }, + }, + ) + # Only ``prompt_segment_ids`` is stamped: the build-synthesis + # ``response_id`` / ``hash_ids`` companions (a duplicate per-node hash-list + # copy) reach no build, sidecar, store, or dispatch consumer -- the sidecar + # empties ``metadata["trie"]`` and the store envelope carries + # ``prompt_segment_ids`` only -- so dynamo persists the same trie shape as + # the dag_jsonl/native adapters. ``extra`` stays available for weka. + return stamp_prompt_segment_ids(llm, build.prompt_path) + + +__all__ = [ + "DEFAULT_VIRTUAL_BLOCK_SIZE", + "DynamoISLMismatchError", + "build_dynamo_llm_node", + "dynamo_recon_callbacks", + "dynamo_trie_nodes", +] diff --git a/src/aiperf/dataset/graph/adapters/native.py b/src/aiperf/dataset/graph/adapters/native.py new file mode 100644 index 0000000000..6ac6f57e46 --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/native.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Plugin-registry facade for the canonical native graph format. + +Native graph files (`.yaml` / `.yml` / `.jsonl` conforming to the graph +schema) are not third-party trace formats but are registered as a +plugin here so that detection is uniform: every recognized format +goes through the same `graph_adapter` registry walk. This adapter has +the lowest detection_priority (1) so it acts as the catch-all fallback. +""" + +from __future__ import annotations + +from pathlib import Path + +from aiperf.dataset.graph.models import ParsedGraph +from aiperf.dataset.graph.parse_context import GraphParseContext + + +class NativeGraphAdapter: + """Canonical AIPerf graph workload format (.yaml / .yml / .jsonl).""" + + @classmethod + def can_load(cls, path: Path) -> bool: + return path.suffix.lower() in (".yaml", ".yml", ".jsonl") + + @classmethod + def parse(cls, path: Path, ctx: GraphParseContext | None = None) -> ParsedGraph: + # ctx is accepted for protocol uniformity and ignored: native graph + # files are fully self-describing, with no run-derived parse knobs. + del ctx + # `parse_native` lives in `parser`. Import it lazily so a cold + # `import aiperf.dataset.graph.adapters` (which pulls in this module at + # `adapters/__init__` time, e.g. via the worker's unified store client) + # does not eagerly import the full `parser` module. `parser` imports + # nothing under `adapters` (its format-detection helpers live inside + # `parser` itself), so this lazy import keeps the two decoupled rather + # than breaking a hard cycle. + from aiperf.dataset.graph.parser import parse_native + + return parse_native(path) diff --git a/src/aiperf/dataset/graph/adapters/protocols.py b/src/aiperf/dataset/graph/adapters/protocols.py new file mode 100644 index 0000000000..9e605c1a89 --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/protocols.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Plugin-registry protocol for graph workload-format adapters. + +Adapters are registered under the `graph_adapter` plugin category and +selected by name (`--graph-format`) or by priority-ordered `can_load()` +auto-detection. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Protocol + +from aiperf.dataset.graph.models import ParsedGraph +from aiperf.dataset.graph.parse_context import GraphParseContext + + +class GraphAdapterProtocol(Protocol): + """Convert a third-party trace/log file into a ParsedGraph. + + Detection is two-stage: + * `can_load(path)` is a cheap predicate over file extension + content + sniff. Implementations should read only enough bytes to disambiguate. + * `parse(path, ctx)` does the full conversion. `ctx` carries the + run-derived knobs (see :class:`GraphParseContext`); adapters map the + fields they consume onto their entry function, forwarding a field + ONLY when it is set so `parse(path)` stays byte-equal to the + protocol-default entry. Adapters may raise an adapter-specific + ValueError subclass on failure; callers wrap into `GraphParseError` + at the parser layer. + + Multiple adapters may return True for the same file; the registry's + `detection_priority` metadata breaks ties (higher wins). + """ + + @classmethod + def can_load(cls, path: Path) -> bool: ... + + @classmethod + def parse(cls, path: Path, ctx: GraphParseContext | None = None) -> ParsedGraph: ... diff --git a/src/aiperf/dataset/graph/adapters/shared/__init__.py b/src/aiperf/dataset/graph/adapters/shared/__init__.py new file mode 100644 index 0000000000..0396803bb6 --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/shared/__init__.py @@ -0,0 +1,12 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Format-agnostic helpers shared across graph adapters. + +Cross-cutting utilities (corpus-backed content synthesis) used by more than +one adapter family. Import the concrete module directly, e.g. +``from aiperf.dataset.graph.adapters.shared.content import ...``. + +Workload-format detection lives in ``aiperf.dataset.graph.parser`` +(``detect_format``): ``parser`` was its only importer, so it was never +adapter-shared. +""" diff --git a/src/aiperf/dataset/graph/adapters/shared/content.py b/src/aiperf/dataset/graph/adapters/shared/content.py new file mode 100644 index 0000000000..067b47fed3 --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/shared/content.py @@ -0,0 +1,557 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared corpus-backed content synthesis for the recorded-trace adapters. + +The weka AND dynamo adapters both +consume it, and the trie core's content callbacks are built on it. + +Weka traces carry ``hash_ids`` + token counts (``in``/``out``), NOT real text. +The block-decode callback is a plain ``corpus[start:start+bs]`` window seeded +per ``(trace_id, +hash_id)`` (NO block-separation token, full-size blocks). Partial tails use +sha256-keyed corpus offsets so two processes produce identical bytes. The result +is byte-deterministic given ``(trace, root_seed, tokenizer)``. + +:class:`CorpusContentSynthesizer` owns a corpus-backed +:class:`~aiperf.dataset.generator.prompt.PromptGenerator` and exposes the +block-decode / partial-tail / decode-to-text callbacks the trie builders +(:mod:`~aiperf.dataset.graph.adapters.weka.trie_build`, +:mod:`~aiperf.dataset.graph.adapters.dynamo.trie_lowering`) consume to emit +message-unit prompt segments. This module carries no +conversation-reconstruction state of its own. +""" + +from __future__ import annotations + +import contextlib +import hashlib +import secrets +from collections.abc import Iterator +from typing import TYPE_CHECKING, Literal + +from aiperf.common import random_generator as rng + +if TYPE_CHECKING: + import numpy as np + + # The shm-backed int32 token array a parallel-parse worker attaches, or a + # plain token-id list; both expose the ``__getitem__``/``__len__`` surface + # the block-decode and partial-tail callbacks read. + SharedCorpus = np.ndarray | list[int] +from aiperf.common.hash_id_random_generator import HashIdRandomGenerator +from aiperf.common.tokenizer import Tokenizer +from aiperf.config import PromptConfig +from aiperf.dataset.generator.prompt import PromptGenerator + +# Last-resort seed for offline weka synthesis when no ambient RNG manager exists +# and no explicit run seed was threaded in (offline tooling calling +# ``parse_graph`` directly, which never runs ``bootstrap``). A fixed value +# keeps offline content byte-deterministic across runs and processes. +_DEFAULT_OFFLINE_SEED = 42 + + +@contextlib.contextmanager +def _seeded_global_rng(root_seed: int | None) -> Iterator[None]: + """Pin the global RNG ``root_seed`` for the duration, then restore it. + + Weka content synthesis derives its corpus pool + per-hash reseed key from the + GLOBAL ``rng`` manager's ``root_seed`` (see + :meth:`CorpusContentSynthesizer._build_generator`). In the main benchmark path + ``bootstrap`` seeds that manager with the run ``--random-seed`` before the + dataset is parsed, so the global state is already correct. But the directory + parse fans the per-file build across a multiprocessing pool; a ``fork`` + worker inherits the seeded ``_manager``, while a forkserver/``spawn`` worker + starts with ``_manager is None`` and ``rng.derive`` raises + ``InvalidStateError`` (no seed inherited). Threading the run seed in and + pinning it here makes the synthesized bytes identical to the in-process + build at ANY seed, regardless of the pool start method. + + ``root_seed is None`` defers to the ambient manager when one exists (the + bootstrap-seeded in-process path + the unit-test auto-fixture seed). + Production offline routes -- any direct ``parse_graph`` call from tooling + with no run config -- resolve a concrete seed via + :func:`resolve_effective_root_seed` BEFORE reaching this pin, so they pass a + non-``None`` ``root_seed`` and never take the ``None`` branch. The pin + survives only as a last resort: a caller that passes ``root_seed=None`` with + no ambient manager would otherwise trip the ``rng.derive`` in + :meth:`CorpusContentSynthesizer._build_generator` (``InvalidStateError``), so + :data:`_DEFAULT_OFFLINE_SEED` is pinned for the duration to keep offline + content deterministic. Either way the prior manager (including an unset + ``None``) is restored on exit so a shared in-process caller (tests, the + serial parse path) is not left with a mutated global. + """ + if root_seed is None: + if rng._manager is not None: + yield + return + prior = rng._manager + rng.init(_DEFAULT_OFFLINE_SEED) + try: + yield + finally: + rng._manager = prior + return + prior = rng._manager + rng.reset() + rng.init(root_seed) + try: + yield + finally: + rng._manager = prior + + +def resolve_effective_root_seed(root_seed: int | None) -> int: + """Resolve the concrete content seed threaded into every recorded-trace parse route (weka and dynamo). + + Explicit run seed wins; otherwise the ambient bootstrap-seeded manager's + root seed; otherwise a fresh OS-entropy seed generated ONCE here. The + resolved ``int`` is threaded (via the resolved parse kwargs) into the + serial path and every pool worker, so one run's synthesized content is + internally consistent at any parallel threshold while distinct unseeded + runs differ — the unseeded fallback is deliberately NOT a hardcoded + constant. No other process re-derives this seed: the DatasetManager parses + once and the TimingManager ingests the sidecar it writes rather than + re-parsing, so a per-process random fallback cannot diverge content within a + run. + """ + if root_seed is not None: + return root_seed + ambient = rng.root_seed() + if ambient is not None: + return ambient + return secrets.randbits(64) + + +# Weka content corpus selector. A local string literal (only this module needs +# it, so no shared enum). ``"coding"`` matches the weka loader's +# ``default_prompt_corpus: coding`` -- the corpus the recorded weka workloads +# were captured against. +PromptCorpus = Literal["coding", "sonnet"] +PROMPT_CORPUS_CODING: PromptCorpus = "coding" + + +# --------------------------------------------------------------------------- +# dag-v3 integration: corpus-backed synthesizer exposing decode/sample callbacks. +# --------------------------------------------------------------------------- + + +class CorpusContentSynthesizer: + """Owns a corpus-backed generator and exposes recorded-trace content callbacks. + + Determinism: given the same ``(trace, root_seed, tokenizer, corpus)`` the + decoded block tokens are byte-identical to agentx's ``WekaTraceLoader`` wire + payload. Block tokens are seeded per ``(trace_id, hash_id)`` via the + generator's ``_hash_id_corpus_rng``; partial tails via ``sha256(seed_string)`` + corpus offsets. + + Corpus selection MUST match agentx's live weka path for byte-exact payload + parity. agentx's ``CustomDatasetComposer`` honors the weka loader's + ``default_prompt_corpus: coding`` (``plugins.yaml``) and injects a + :class:`~aiperf.dataset.generator.coding_content.CodingContentGenerator` + (a procedurally-built coding-text tool pool, RNG-keyed + ``dataset.coding_content.corpus``) rather than the Shakespeare + :class:`~aiperf.dataset.generator.prompt.PromptGenerator` + (``dataset.prompt.corpus``). dag-v3 therefore defaults to ``"coding"`` so the + graph weka synthesis draws from the SAME corpus and reseed key; the wire + bytes then match agentx turn-for-turn (a plain ``PromptGenerator`` produces + Shakespeare text with matching token COUNTS but entirely different BYTES, + which is why a Shakespeare-backed synthesizer diverges from the live agentx + A/B). ``"sonnet"`` selects the Shakespeare path for callers that want it. + + Both backends expose the same ``_tokenized_corpus`` / ``_corpus_size`` / + ``_cache`` / ``_hash_id_corpus_rng`` / ``tokenizer`` surface the block-decode + and partial-tail callbacks read, so the content path is corpus-agnostic. + """ + + def __init__( + self, + tokenizer_name: str, + *, + prompt_corpus: PromptCorpus = PROMPT_CORPUS_CODING, + root_seed: int | None = None, + shared_corpus: SharedCorpus | None = None, + ) -> None: + """Build the corpus-backed generator (or adopt a pre-built corpus). + + ``shared_corpus`` supplies a pre-built token array (the parallel + parser's parent-built shared-memory block): the generator is then + constructed WITHOUT paying its own corpus build and reads the supplied + array directly. The array MUST be byte-identical to what this + synthesizer's ``(tokenizer, prompt_corpus, root_seed)`` would build + (the parent builds it with the SAME key), so decoded content is + byte-identical to a self-built synthesizer. + """ + self._tokenizer_name = tokenizer_name + self._prompt_corpus = prompt_corpus + self._root_seed = root_seed + with _seeded_global_rng(root_seed): + self._pg, self._hash_id_corpus_rng = self._build_generator( + tokenizer_name, prompt_corpus, shared_corpus=shared_corpus + ) + self._block_size = 64 + # First ``_corpus_size`` seen by :meth:`_decode_block_tokens_offset_cached`; + # asserted per call so a corpus rebind that would invalidate the cached + # offsets fails loud instead of silently reproducing wrong bytes. + self._offset_cache_corpus_size: int | None = None + + @staticmethod + def reset_worker_cache() -> None: + """Drop the process-level synthesizer cache (test isolation).""" + _WORKER_SYNTH_CACHE.clear() + + def corpus_tokens(self) -> list[int]: + """Return the generator's built corpus token ids (``_tokenized_corpus``). + + Used by the parallel parser + (:mod:`~aiperf.dataset.graph.adapters.weka.trace_parallel`) to hoist the + deterministic corpus into a single ``shared_memory`` block built ONCE in + the parent process; every worker then attaches that block (CoW) instead + of rebuilding the ~600K-token coding pool per worker -- the per-worker + rebuild is what ballooned RSS into tens of GB on the real corpus. + """ + return list(self._pg._tokenized_corpus) + + def attach_shared_corpus(self, corpus: SharedCorpus) -> None: + """Point the generator's corpus at a pre-built (shared) token array. + + ``corpus`` must be the byte-identical token sequence + :meth:`corpus_tokens` would have produced for this synthesizer's + ``(tokenizer, prompt_corpus, root_seed)`` -- the parent builds it once + with the SAME seed, so determinism (and the byte-exact serial/parallel + parity contract) is preserved. Replaces ``_tokenized_corpus`` / + ``_corpus_size`` in place; the block-decode and partial-tail callbacks + read those attributes, so no other wiring changes. Prefer constructing + with ``shared_corpus=`` instead, which skips the private corpus build + entirely; this rebind is for an already-built synthesizer. + """ + self._pg._tokenized_corpus = corpus + self._pg._corpus_size = len(corpus) + + @staticmethod + def _build_generator( + tokenizer_name: str, + prompt_corpus: PromptCorpus, + shared_corpus: SharedCorpus | None = None, + ) -> tuple[PromptGenerator, HashIdRandomGenerator]: + """Build the corpus backend + its per-``(trace_id, hash_id)`` reseed RNG. + + Returns ``(generator, hash_id_corpus_rng)``. For ``"coding"`` the + ``CodingContentGenerator`` already builds its own ``_hash_id_corpus_rng`` + in ``__init__`` (off ``rng.derive("dataset.coding_content.corpus")``), + so it is used directly -- matching agentx's live weka path exactly. For + ``"sonnet"`` the plain ``PromptGenerator`` does NOT build one (only + ``CodingContentGenerator`` does), so it is derived here from the + generator's ``_corpus_rng`` (``rng.derive("dataset.prompt.corpus")``). + ``from_base_rng`` reads ``base_rng.seed`` without consuming RNG state, so + the reseed key is stable. + + ``shared_corpus`` short-circuits the generator's own corpus build: the + corpus-build step (``_build_tool_pool`` / ``_initialize_corpus``) is + replaced with a direct adoption of the supplied array. Byte parity is + preserved because the block-decode path depends only on the corpus + CONTENT (supplied identical by contract) and on ``_hash_id_corpus_rng``, + whose seed derives from the corpus-rng SEED (read without consuming + state) -- never on RNG state the skipped build would have advanced + (the build consumes only the independent ``_template_rng`` stream). + + Both the coding corpus POOL text (``_template_rng.shuffle`` block order in + ``CodingContentGenerator._build_tool_pool``) and the per-``(trace_id, + hash_id)`` reseed key (``_hash_id_corpus_rng.seed`` derived off + ``dataset.coding_content.corpus``) read the GLOBAL RNG ``root_seed`` via + ``rng.derive``. The caller (:meth:`__init__`) wraps this in + :func:`_seeded_global_rng` so an explicit ``root_seed`` is honored even + when the global manager is unset (e.g. a ``spawn``-started parallel parse + worker that did NOT inherit the parent's seeded ``_manager``). + """ + # In a forkserver worker the configured tokenizer is preloaded into + # the helper heap and CoW-shared (see aiperf.dataset._tokenizer_preload); + # prefer it to avoid each worker re-reading the tokenizer from disk. The + # preload and the on-demand fallback both key on the SAME (name, trust, + # revision) triple, so a hit and a fallback build produce identical + # content. + from aiperf.dataset._tokenizer_preload import ( + _env_revision, + _env_trust_remote_code, + get_preloaded, + ) + + trust_remote_code = _env_trust_remote_code() + revision = _env_revision() + + tokenizer = get_preloaded( + tokenizer_name, trust_remote_code=trust_remote_code, revision=revision + ) + if tokenizer is None: + # A transient load failure (network blip, HF down, auth) must NOT + # silently degrade to the builtin tokenizer: the run would then + # produce builtin-sized content while claiming the REQUESTED + # name/revision. Fail loudly instead so a transient failure never + # yields mislabeled content. resolve_alias=True mirrors the + # forkserver preload (aiperf.dataset._tokenizer_preload._preload) + # so a hit and a miss resolve the SAME tokenizer. + tokenizer = Tokenizer.from_pretrained( + tokenizer_name, + trust_remote_code=trust_remote_code, + revision=revision, + resolve_alias=True, + ) + + if prompt_corpus == PROMPT_CORPUS_CODING: + # Local import: the coding generator pulls in the procedural code + # builders, kept off the import path for the SONNET/topology paths. + from aiperf.dataset.generator.coding_content import ( + CodingContentGenerator, + ) + + if shared_corpus is None: + ccg = CodingContentGenerator(config=PromptConfig(), tokenizer=tokenizer) + else: + + class _SharedCorpusCodingContentGenerator(CodingContentGenerator): + """Adopts the supplied corpus instead of building the pool.""" + + def _build_tool_pool(self) -> None: + self._tool_pool = shared_corpus + + ccg = _SharedCorpusCodingContentGenerator( + config=PromptConfig(), tokenizer=tokenizer + ) + return ccg, ccg._hash_id_corpus_rng + + # prefix-prompt pool disabled (``prefix_prompts=None``): weka synthesis + # never draws from it, and a pool would only consume the independent + # prefix RNG -- the block/tail RNGs are separate, so output is unchanged. + if shared_corpus is None: + pg = PromptGenerator( + prompts=PromptConfig(), + prefix_prompts=None, + tokenizer=tokenizer, + ) + else: + + class _SharedCorpusPromptGenerator(PromptGenerator): + """Adopts the supplied corpus instead of tokenizing the file.""" + + def _initialize_corpus(self) -> None: + self._tokenized_corpus = shared_corpus + self._corpus_size = len(shared_corpus) + + pg = _SharedCorpusPromptGenerator( + prompts=PromptConfig(), + prefix_prompts=None, + tokenizer=tokenizer, + ) + return pg, HashIdRandomGenerator.from_base_rng(pg._corpus_rng) + + # --- agentx-faithful content callbacks -------------------------------- + + def _decode_block_tokens( + self, + hash_ids: list[int], + *, + block_size: int | None = None, + cache: dict[int, list[int]] | None = None, + trace_id: str | None = None, + ) -> list[int]: + """Concatenate per-hash-id token blocks (mirrors agentx exactly). + + Plain ``corpus[start:start+bs]`` window seeded per ``(trace_id, + hash_id)``; NO block-separation token, full-size blocks. + + ``trace_id`` scopes the hash-id namespace. Weka traces declare + ``hash_id_scope: "local"`` (ids are a per-trace namespace), so the weka + trie path passes the trace's own id -- hash_id 0 in trace A and trace B + then decode to DIFFERENT bytes, preventing manufactured cross-trace + KV-cache sharing at the server. A ``trace_id`` caller MUST supply its + own ``cache``: the shared ``pg._cache`` is keyed by bare hash id and + would leak one trace's bytes into another. ``None`` keeps the GLOBAL + namespace (the linear mooncake path, and dynamo's recorded chained + sequence hashes, which are content-global identity by construction). + + ``block_size`` / ``cache`` also serve NON-weka trie block sizes (dynamo + records 16-token blocks): the defaults keep the global single-namespace + path byte-exact, while a caller decoding at a different size MUST also + supply its own cache -- mixing block sizes in the shared cache would + return wrong-sized blocks across adapters. + """ + pg = self._pg + r = self._hash_id_corpus_rng + bs = self._block_size if block_size is None else block_size + corpus = pg._tokenized_corpus + corpus_size = pg._corpus_size + if cache is None and trace_id is not None: + raise ValueError( + "trace-scoped block decode requires a per-trace cache; the " + "shared cache is keyed by bare hash id and would leak bytes " + "across traces" + ) + cache = pg._cache if cache is None else cache + tokens: list[int] = [] + for h in hash_ids: + cached = cache.get(h) + if cached is None: + # Preserve the bare call shape on the default path: injected + # test doubles (and any external RNG duck-type) predate the + # trace_id kwarg. + if trace_id is None: + r.reseed_for_hash_id(h) + else: + r.reseed_for_hash_id(h, trace_id=trace_id) + start = r.randrange(corpus_size) + end = start + bs + cached = list(corpus[start:end]) + if end > corpus_size: + cached.extend(corpus[: end - corpus_size]) + cache[h] = cached + tokens.extend(cached) + return tokens + + def _decode_block_tokens_offset_cached( + self, + hash_ids: list[int], + *, + block_size: int, + offset_cache: dict[int, int], + ) -> list[int]: + """Memory-lean twin of :meth:`_decode_block_tokens` (GLOBAL namespace only). + + Byte-for-byte identical output to :meth:`_decode_block_tokens` for the + same ``(hash_ids, block_size)`` and RNG object, but caches one ``int`` + corpus offset per hash id instead of the decoded ``list[int]`` block -- + ``~4-8 B`` vs ``~block_size`` ints per entry, the decode-cache tier's + single largest on-peak shave at corpus scale. The + ``test_weka_memory_regressions`` differential test pins the equality + across miss/hit paths and list/``np.int32`` corpus backings. + + The equality rests on :meth:`~aiperf.common.hash_id_random_generator.HashIdRandomGenerator.reseed_for_hash_id` + being a FULL reseed: ``start`` is a pure function of ``(seed, scope, + hash_id, corpus_size)``, so the miss path here issues the IDENTICAL + ``reseed_for_hash_id(h)`` + ``randrange(corpus_size)`` call pair (same + RNG object, same order -> same RNG trajectory) that :meth:`_decode_block_tokens` + issues on its first occurrence of ``h``. Both paths then read the same + ``corpus[start:start+bs]`` window (wraparound included). Repeat + (hit-path) decodes re-slice from the cached offset rather than reissuing + RNG calls, so neither method touches the RNG on a repeat -- the + trajectories stay aligned. + + **Corpus-immutability assumption.** A cached offset reproduces the + original bytes ONLY while ``_tokenized_corpus`` / ``_corpus_size`` are + unchanged, or rebound to a BYTE-IDENTICAL array (the + :meth:`attach_shared_corpus` contract -- note :func:`get_or_build_synthesizer` + rebinds on a cache HIT). A byte-identical rebind keeps ``_corpus_size`` + equal, so as a cheap tripwire the first-seen ``corpus_size`` is snapshotted + and asserted per call: a contract-violating rebind to a DIFFERENT-sized + corpus fails loud rather than silently emitting wrong bytes. (A same-size + byte-different rebind is out of scope -- the size check is a guard, not a + content hash; the contract above forbids it.) + """ + pg = self._pg + r = self._hash_id_corpus_rng + bs = block_size + corpus = pg._tokenized_corpus + corpus_size = pg._corpus_size + snapshot = self._offset_cache_corpus_size + if snapshot is None: + self._offset_cache_corpus_size = corpus_size + elif snapshot != corpus_size: + raise RuntimeError( + f"corpus size changed under the offset cache ({snapshot} -> " + f"{corpus_size}): cached offsets no longer reproduce the original " + "bytes (see the corpus-immutability assumption on " + "_decode_block_tokens_offset_cached)." + ) + tokens: list[int] = [] + for h in hash_ids: + # ``get`` (not truthiness): offset 0 is a legal cached start. + start = offset_cache.get(h) + if start is None: + r.reseed_for_hash_id(h) + start = r.randrange(corpus_size) + offset_cache[h] = start + end = start + bs + if end > corpus_size: + # Wraparound: verbatim two-step shape from _decode_block_tokens. + block = list(corpus[start:end]) + block.extend(corpus[: end - corpus_size]) + tokens.extend(block) + else: + # DELIBERATE shape deviation from _decode_block_tokens (waived + # per the Task-3 amendment): extend straight from the slice, + # skipping its redundant ``list()`` copy -- the hit path runs + # once per covered block, and that copy alone cost ~6.5% of a + # full corpus-scale parse. Byte-identity holds: extend-from- + # slice appends exactly the values ``list(corpus[start:end])`` + # would (for an np backing, iterating the view yields the same + # np scalars ``list(view)`` yields), and the window is the same. + tokens.extend(corpus[start:end]) + return tokens + + def _sample_partial_tail_tokens(self, n_tokens: int, seed: str) -> list[int]: + """Deterministic per-seed partial-block tokens sized to ``n_tokens``. + + sha256-keyed corpus offset (PYTHONHASHSEED-independent, cross-process + stable). Mirrors agentx ``HashIdsPromptSynthesisMixin``. + """ + if n_tokens <= 0: + return [] + pg = self._pg + corpus_size = pg._corpus_size + digest = hashlib.sha256(seed.encode()).digest() + offset = int.from_bytes(digest[:8], "big") % max(corpus_size - n_tokens, 1) + return list(pg._tokenized_corpus[offset : offset + n_tokens]) + + def _decode_tokens_to_text(self, tokens: list[int]) -> str: + return self._pg.tokenizer.decode(tokens) + + +# Process-level synthesizer cache keyed by (tokenizer, corpus, root_seed). In a +# parallel-parse worker this is populated once (by trace_parallel's +# _init_worker, which also attaches the shared-memory corpus) and reused across +# every trace the worker handles -- so the ~600K-token coding pool is built ONCE +# per process and the corpus array is shared across workers, instead of rebuilt +# per trace (the rebuild-per-trace path ballooned RSS into tens of GB). +_WORKER_SYNTH_CACHE: dict[tuple[str, str, int | None], CorpusContentSynthesizer] = {} + + +def get_or_build_synthesizer( + tokenizer_name: str, + *, + prompt_corpus: str, + root_seed: int | None, + shared_corpus: SharedCorpus | None = None, +) -> CorpusContentSynthesizer: + """Return a cached synthesizer for the key, building one on first miss. + + The block cache is content-addressed by ``hash_id`` (see + :meth:`CorpusContentSynthesizer._decode_block_tokens`), so reusing one + synthesizer across traces is byte-identical to building a fresh one per + trace -- the serial path (no cache populated) and the parallel path (cache + populated in ``_init_worker``) produce the same wire bytes. + + ``shared_corpus`` (the parallel parser's parent-built shm array) is adopted + instead of built: a cache MISS constructs the synthesizer directly on it + (skipping the private corpus build entirely); a cache HIT rebinds the + existing synthesizer's corpus to it (releasing any privately-built copy). + Byte parity is unaffected -- the array is byte-identical to what the key + would build. + """ + key = (tokenizer_name, prompt_corpus, root_seed) + synth = _WORKER_SYNTH_CACHE.get(key) + if synth is None: + synth = CorpusContentSynthesizer( + tokenizer_name=tokenizer_name, + prompt_corpus=prompt_corpus, # type: ignore[arg-type] + root_seed=root_seed, + shared_corpus=shared_corpus, + ) + _WORKER_SYNTH_CACHE[key] = synth + elif shared_corpus is not None: + synth.attach_shared_corpus(shared_corpus) + return synth + + +__all__ = [ + "PROMPT_CORPUS_CODING", + "PromptCorpus", + "CorpusContentSynthesizer", +] diff --git a/src/aiperf/dataset/graph/adapters/shared/idle_gap.py b/src/aiperf/dataset/graph/adapters/shared/idle_gap.py new file mode 100644 index 0000000000..5c8500b2ce --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/shared/idle_gap.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared idle-gap warp entry defaults for recorded-trace adapters (weka, dynamo). + +Recorded traces replay their inter-request delays through the shared +:class:`~aiperf.dataset.graph.segment_ir.trie_content.ActiveIdleWarp`. The run +path resolves the cap from the dataset's ``synthesis.idle_gap_cap_seconds`` +(``--synthesis-idle-gap-cap``) and forwards it via +``GraphParseContext.idle_gap_cap_seconds``; the default below is only the +adapter-entry fallback for direct callers with no run config (CLI tooling, +tests). Both adapters share ONE sentinel + default so a dynamo replay warps +exactly like a weka replay under the same knob. +""" + +from __future__ import annotations + +from typing import Any + +# Default per-trace idle-gap cap (seconds) when the entry kwarg is left at the +# sentinel -- the value the recorded weka workloads were captured/replayed +# against. ``workload_detect._GRAPH_IDLE_GAP_CAP_DEFAULT`` and the +# ``synthesis.idle_gap_cap_seconds`` config default mirror it for the run path. +DEFAULT_IDLE_GAP_CAP_SECONDS = 60.0 + +# Sentinel for adapter-entry ``idle_gap_cap_seconds`` kwargs: left in place the +# cap resolves to :data:`DEFAULT_IDLE_GAP_CAP_SECONDS`; pass an explicit float +# (or ``None`` to disable warping) to override. A sentinel -- not the default as +# the literal -- lets callers pass ``None`` to mean "disable" without it +# colliding with "use the default". +IDLE_GAP_CAP_USE_DEFAULT: Any = object() + + +def resolve_idle_gap_cap(value: float | None | Any) -> float | None: + """Resolve an entry kwarg to a concrete cap (sentinel -> the 60s default).""" + if value is IDLE_GAP_CAP_USE_DEFAULT: + return DEFAULT_IDLE_GAP_CAP_SECONDS + return value diff --git a/src/aiperf/dataset/graph/adapters/shared/output_cap.py b/src/aiperf/dataset/graph/adapters/shared/output_cap.py new file mode 100644 index 0000000000..564f273b0b --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/shared/output_cap.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared wire output-cap resolution for recorded-trace adapters (weka, dynamo). + +Recorded replay always pins each call's generation to the recording via the +native ``LlmNode.max_tokens`` field (``Turn.max_tokens`` naming) -- folded into +the store envelope by ``store_builder._trie_envelope`` and mapped by the worker +to the endpoint's wire token field (``max_completion_tokens``, or +``max_tokens`` under ``--use-legacy-max-tokens``; see +``graph/worker_materialize.py::_apply_dispatch_overrides``). Both adapters +resolve the cap through :func:`wire_output_cap` so zero-output turns are +handled identically. +""" + +from __future__ import annotations + +from aiperf.common.aiperf_logger import AIPerfLogger + +_logger = AIPerfLogger(__name__) + + +def wire_output_cap(recorded_out: int, *, node_id: str) -> int: + """Resolve a node's recorded output length to its wire generation cap. + + A recorded length of 0 (a zero-output or aborted turn, or a capture with no + ``output_tokens``) is not a sendable cap -- ``max_tokens: 0`` is rejected or + degenerate on OpenAI-compatible servers -- so it upgrades to 1 with a + warning rather than replaying that turn unbounded or uncapped. + """ + if recorded_out > 0: + return recorded_out + _logger.warning( + f"node {node_id!r}: recorded output length is 0; upgrading the wire " + "generation cap to max_output_tokens=1" + ) + return 1 diff --git a/src/aiperf/dataset/graph/adapters/shared/peak_context.py b/src/aiperf/dataset/graph/adapters/shared/peak_context.py new file mode 100644 index 0000000000..74c82f2b78 --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/shared/peak_context.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-trace PEAK context-length helpers for the recorded-trace adapters. + +A later filter-then-cap trace selection pass drops traces whose peak context +(input + output tokens on the largest single request) exceeds +``--max-context-length``. These pure functions compute that peak straight off +each adapter's recorded schema -- no graph build, no tokenization -- so the +selector can screen a corpus cheaply. + +Both helpers mirror the token accounting their trie builders already apply, so +the screened peak matches the value the built graph would carry: + +- weka: :func:`weka_trace_peak_context` mirrors + :func:`~aiperf.dataset.graph.adapters.weka.trie_build._build_llm_node`'s + ``max_osl`` cap -- only TOP-LEVEL leaves are capped to ``min(out, max_osl)``; + subagent-body leaves (any nesting depth) stay uncapped, matching + ``_top_level_leaf_ids``. +- dynamo: :func:`dynamo_tree_peak_context` mirrors + :func:`~aiperf.dataset.graph.adapters.dynamo.trie_lowering.dynamo_trie_nodes`'s + per-record token read (``replay.input_length`` else ``input_tokens`` else 1, + plus ``output_tokens`` or 0). ``max_osl`` does NOT cap dynamo output. +""" + +from __future__ import annotations + +from collections.abc import Iterable + +from aiperf.dataset.graph.adapters.dynamo.trace_reader import AgentTraceRecord +from aiperf.dataset.graph.adapters.weka.trace_models import ( + WekaRequest, + WekaSubagentEntry, + WekaTrace, +) + + +def weka_trace_peak_context(trace: WekaTrace, *, max_osl: int | None) -> int: + """Peak ``input_length + output`` over every leaf request in a weka trace. + + Top-level leaves (those directly in ``trace.requests``) use + ``min(output_length, max_osl)`` when ``max_osl`` is set, matching the + dispatch ``max_tokens`` cap the trie builder applies to top-level chain + requests. Subagent-body leaves (inside any :class:`WekaSubagentEntry`, at + any nesting depth) always use the raw recorded ``output_length``. Returns + ``0`` for a trace with no leaf requests. + """ + peak = 0 + for req in trace.requests: + if isinstance(req, WekaSubagentEntry): + peak = max(peak, _uncapped_leaf_peak(req.requests)) + else: + output = ( + req.output_length + if max_osl is None + else min(req.output_length, max_osl) + ) + peak = max(peak, req.input_length + output) + return peak + + +def _uncapped_leaf_peak(requests: list[WekaRequest]) -> int: + """Peak ``input_length + output_length`` over subagent-body leaves (uncapped). + + Recurses into nested :class:`WekaSubagentEntry` markers so a + subagent-within-subagent body is screened at its raw recorded ``out``. + """ + peak = 0 + for req in requests: + if isinstance(req, WekaSubagentEntry): + peak = max(peak, _uncapped_leaf_peak(req.requests)) + else: + peak = max(peak, req.input_length + req.output_length) + return peak + + +def dynamo_tree_peak_context(records_or_tree: Iterable[AgentTraceRecord]) -> int: + """Peak ``input_length + output_tokens`` over a dynamo session-tree's records. + + ``records_or_tree`` is any iterable of :class:`AgentTraceRecord` making up a + tree (a root session plus its ``parent_session_id`` descendants). Each + record's input length is read exactly as the trie lowering reads it: + ``request.replay.input_length`` when replay metadata is present, else + ``request.input_tokens``, else ``1``; output uses ``request.output_tokens`` + or ``0``. ``max_osl`` does NOT cap dynamo output. Returns ``0`` for an empty + iterable. + """ + peak = 0 + for record in records_or_tree: + req = record.request + if req is not None and req.replay is not None: + input_length = int(req.replay.input_length) + elif req is not None and req.input_tokens: + input_length = int(req.input_tokens) + else: + input_length = 1 + output_tokens = ( + int(req.output_tokens) if req is not None and req.output_tokens else 0 + ) + peak = max(peak, input_length + output_tokens) + return peak + + +__all__ = [ + "dynamo_tree_peak_context", + "weka_trace_peak_context", +] diff --git a/src/aiperf/dataset/graph/adapters/shared/selection.py b/src/aiperf/dataset/graph/adapters/shared/selection.py new file mode 100644 index 0000000000..ff25ec387e --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/shared/selection.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Schema-only filter-then-cap trace selection shared by the recorded adapters. + +The graph plane historically ignored ``--num-dataset-entries`` and +``--max-context-length`` (ai-dynamo/aiperf#1106): every recorded trace was +built and dispatch silently cloned traces to fill lanes. This module is the ONE +selection primitive both the weka and dynamo loaders call to honor those knobs +BEFORE the expensive build: + +* FILTER: drop a candidate whose peak context (input + output tokens on its + largest single request, computed schema-only by the + :mod:`~aiperf.dataset.graph.adapters.shared.peak_context` helpers) exceeds + ``max_context_length``. +* CAP: keep the FIRST ``num_dataset_entries`` eligible candidates (in the + adapter's deterministic scan order). + +"Filter THEN cap" -- the cap is applied to the ELIGIBLE set, never to the raw +prefix (which would drop below N whenever a rejected trace fell in the first N). +Scanning stops as soon as the cap is reached, so on a capped load only the first +N-eligible-plus-rejected-so-far candidates are ever examined; the returned +:class:`SelectionStats` reflect only what was scanned. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass +from typing import TypeVar + +from aiperf.common.aiperf_logger import AIPerfLogger + +_T = TypeVar("_T") + +_logger = AIPerfLogger(__name__) + + +@dataclass(slots=True) +class SelectionStats: + """Tally of one filter-then-cap selection pass, for the load-summary report. + + ``scanned`` counts candidates examined (fewer than the corpus size when the + cap short-circuits the scan); ``rejected_by_maxctx`` counts those dropped + for exceeding ``max_context_length``; ``largest_observed`` is the peak + context seen across every scanned candidate (eligible or not); ``eligible`` + counts candidates that passed the filter and were kept; ``loaded`` is the + number ultimately handed to the build (``eligible`` under early-stop, since + the scan halts exactly when the cap is filled). + """ + + scanned: int = 0 + rejected_by_maxctx: int = 0 + largest_observed: int = 0 + eligible: int = 0 + loaded: int = 0 + + +def filter_then_cap( + candidates: Iterable[tuple[_T, int]], + *, + num_dataset_entries: int | None, + max_context_length: int | None, +) -> tuple[list[_T], SelectionStats]: + """Filter ``candidates`` by peak context, then cap to the first N eligible. + + ``candidates`` yields ``(item, peak_context)`` pairs in the adapter's + DETERMINISTIC scan order (dir files sorted by name, HF rows in stream order, + dynamo trees sorted by root session id). It should be LAZY: this function + stops pulling from it the moment ``num_dataset_entries`` eligible items are + kept, so a capped load never computes peaks for the whole corpus. + + ``max_context_length is None`` disables the filter (every candidate is + eligible); ``num_dataset_entries is None`` disables the cap (every eligible + candidate is kept). Returns the kept items in scan order plus the + :class:`SelectionStats` for the scanned prefix. + """ + stats = SelectionStats() + kept: list[_T] = [] + for item, peak in candidates: + stats.scanned += 1 + if peak > stats.largest_observed: + stats.largest_observed = peak + if max_context_length is not None and peak > max_context_length: + stats.rejected_by_maxctx += 1 + continue + stats.eligible += 1 + kept.append(item) + if num_dataset_entries is not None and len(kept) >= num_dataset_entries: + break + stats.loaded = len(kept) + return kept, stats + + +def log_selection_summary( + stats: SelectionStats, + *, + source: str, + num_dataset_entries: int | None, + max_context_length: int | None, +) -> None: + """Emit the once-per-build filter-then-cap summary (mirrors the trace loader). + + The graph counterpart of + :meth:`~aiperf.dataset.loader.base_trace_loader.BaseTraceLoader._log_filtering_summary`: + called at the PARENT-side finalize point where the single + :class:`SelectionStats` is produced (once per build, before any worker + fan-out), so it logs exactly once no matter how the build parallelizes. + Every adapter finalize point routes through this ONE helper so the summary + text is uniform across the weka and dynamo loaders. When no selection knob + is set the selection scan never runs, so this is never called (byte-identical + quiet path). + """ + _logger.info( + lambda: ( + f"graph trace selection [{source}]: scanned {stats.scanned:,}, " + f"rejected_by_maxctx {stats.rejected_by_maxctx:,} " + f"(--max-context-length={max_context_length}), " + f"eligible {stats.eligible:,}, loaded {stats.loaded:,} " + f"(--num-dataset-entries={num_dataset_entries}), " + f"largest_observed_peak_context {stats.largest_observed:,}" + ) + ) + + +__all__ = [ + "SelectionStats", + "filter_then_cap", + "log_selection_summary", +] diff --git a/src/aiperf/dataset/graph/adapters/weka/__init__.py b/src/aiperf/dataset/graph/adapters/weka/__init__.py new file mode 100644 index 0000000000..94d902c8cb --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/weka/__init__.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weka trace-format loader. + +The public entry points live in :mod:`.trace`; the remaining modules are +internal trie/parallel-reconstruction stages (content synthesis lives in +:mod:`aiperf.dataset.graph.adapters.shared.content`). +""" + +from aiperf.dataset.graph.adapters.weka.trace import ( + WekaTraceAdapter, + from_weka_trace, +) + +__all__ = ["WekaTraceAdapter", "from_weka_trace"] diff --git a/src/aiperf/dataset/graph/adapters/weka/trace.py b/src/aiperf/dataset/graph/adapters/weka/trace.py new file mode 100644 index 0000000000..b7c73e11a9 --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/weka/trace.py @@ -0,0 +1,913 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Adapter: convert Weka KV-cache-tester agentic-coding traces to ParsedGraph. + +A Weka trace is a single JSON object whose ``requests`` list interleaves +normal API calls (``type: "n"``), streaming API calls (``type: "s"``), and +subagent markers (``type: "subagent"``) with their own nested request lists. + +This adapter emits a **flat segment-trie IR** via +:func:`~aiperf.dataset.graph.adapters.weka.trie_build.build_trie_graph`: +one ``LlmNode`` per recorded normal/streaming request (including recursive +subagent-inner requests) connected by dependency-only ``StaticEdge``s. +Dependency edges are derived from a global interval-order pass over the +recorded request intervals (finished-before + rank). + +Prompt content is real-content and content-addressed: ``build_trie_graph`` walks +the recorded ``hash_ids`` prefix tree, assigns frozen per-block ``(role, +starts_new_message)`` tags, and emits one deduplicated ``SegmentPool`` entry per +message so shared block-aligned prefixes produce identical segment ids (a real +prefix-cache hit). Each node carries its ``prompt_segment_ids`` chain in +``metadata["trie"]``; the recorded ``out`` is plumbed to ``max_tokens`` (dispatch +overrides) and the stream flag tracks the request type. There is no +dispatch-time rewrite. + +This module performs the Weka JSON validation, streaming/parallel ingest, and +trace-record conversion; the flat-trie construction itself lives in +``weka/trie_build.py``, the shared trie core in ``segment_ir/``, and native +JSONL reading in ``parser.py``. +""" + +from __future__ import annotations + +import re +from collections.abc import Iterator +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import msgspec +import orjson +from pydantic import ValidationError + +from aiperf.common.constants import IS_WINDOWS +from aiperf.dataset.graph.adapters.shared.idle_gap import ( + DEFAULT_IDLE_GAP_CAP_SECONDS as _DEFAULT_IDLE_GAP_CAP_SECONDS, +) +from aiperf.dataset.graph.adapters.shared.idle_gap import ( + IDLE_GAP_CAP_USE_DEFAULT as _USE_DEFAULT, +) +from aiperf.dataset.graph.adapters.shared.selection import SelectionStats +from aiperf.dataset.graph.adapters.weka import trace_parallel +from aiperf.dataset.graph.adapters.weka.trace_models import ( + EmptyWekaTraceError, + WekaHashScopeError, + WekaSchemaError, + WekaTraceAdapterError, +) + +if TYPE_CHECKING: + from aiperf.dataset.graph.adapters.weka.trace_models import WekaTrace +from aiperf.dataset.graph.models import ParsedGraph +from aiperf.dataset.graph.parse_context import ( + UNSET, + GraphParseContext, + publish_ctx_tokenizer_env, +) + +_DEFAULT_TAG = "from-weka-trace" + +# Bounded sniff budget for ``WekaTraceAdapter.can_load`` — Weka traces are +# single-document JSON with the discriminator keys near the top-level object. +_DETECTION_SNIFF_BYTES = 4096 + +# HuggingFace repo-id pattern: ``org/name`` with no path separators beyond the +# single slash and no file extension. The weka HF corpora live at e.g. +# ``semianalysisai/cc-traces-weka-061526``. A ``--graph`` arg that is NOT an +# existing path AND that matches this shape AND carries the weka corpus marker +# is treated as a weka HF dataset id. The marker requirement is what keeps an +# arbitrary ``org/name`` (``meta-llama/Llama-3``, ``openai/gpt``) from being +# silently routed into ``datasets.load_dataset``. +_HF_REPO_ID_RE = re.compile(r"^[A-Za-z0-9][\w.-]*/[\w.-]+$") + +# A repo id is only treated as a weka HF corpus when this case-insensitive +# marker appears in it. The published weka corpora carry ``weka`` in the repo +# name (e.g. ``semianalysisai/cc-traces-weka-061526``); any non-existent +# ``org/name`` string WITHOUT it stays on the unchanged linear pipeline. +_HF_WEKA_MARKER = "weka" + +# File extensions that disqualify a string from being an HF repo id even when it +# happens to match the org/name shape (e.g. ``foo/bar.json``). +_HF_EXCLUDED_SUFFIXES = (".json", ".jsonl", ".yaml", ".yml") + +# The complete set of top-level keys a genuine on-disk weka trace object may +# carry: the required + optional fields of :class:`WekaTrace`, plus the ``kind`` +# marker the native JSONL writer stamps on serialized rows. Any top-level key +# outside this set marks the object as a FOREIGN format (mooncake / sharegpt / +# synthetic) that merely happens to share the weka discriminator keys, so we +# reject it rather than ignoring the extras (mirrors WekaTrace's extra="forbid"). +_WEKA_ALLOWED_KEYS: frozenset[str] = frozenset( + { + "id", + "models", + "block_size", + "hash_id_scope", + "tool_tokens", + "system_tokens", + "requests", + "totals", + "kind", + } +) + + +# The hash-id namespace scopes the adapter supports. Kept in sync with +# ``WekaTrace.hash_id_scope``'s Literal; the explicit pre-validation gate in +# ``_parse_trace_dict`` uses this so an unrecognized scope raises +# WekaHashScopeError instead of a generic ValidationError. +_SUPPORTED_HASH_ID_SCOPES: frozenset[str] = frozenset({"local", "global"}) + + +# The idle-gap entry default + ``_USE_DEFAULT`` sentinel are shared with the +# dynamo adapter (imported above from ``adapters.shared.idle_gap``): the run +# path passes the resolved cap explicitly; the default only applies when +# ``idle_gap_cap_seconds`` is left at the sentinel (direct adapter / CLI +# tooling / tests). + + +def _resolve_parse_kwargs( + *, + tag: str, + idle_gap_cap_seconds: float | None | object, + content_root_seed: int | None, + content_tokenizer: str | None, + prompt_corpus: str | None, + max_osl: int | None, +) -> dict[str, Any]: + """Resolve the per-trace parse kwargs shared by EVERY weka route. + + One spelling for the ``_parse_trace_dict`` keyword set: the single-file + path, the directory pool, the HF pool, and the streaming build plane all + pass this exact dict, so a knob added here reaches every route instead of + needing four separately-spelled kwarg sets that could silently drift. + Resolves the ``_USE_DEFAULT`` idle-gap sentinel and pins the effective content + seed to a concrete int (:func:`resolve_effective_root_seed`) so serial and + pool-worker parses synthesize identical bytes at any parallel threshold. + """ + from aiperf.dataset.graph.adapters.shared.content import ( + resolve_effective_root_seed, + ) + + if idle_gap_cap_seconds is _USE_DEFAULT: + idle_gap_cap_seconds = _DEFAULT_IDLE_GAP_CAP_SECONDS + return { + "tag": tag, + "idle_gap_cap_seconds": idle_gap_cap_seconds, + "content_root_seed": resolve_effective_root_seed(content_root_seed), + "content_tokenizer": content_tokenizer, + "prompt_corpus": prompt_corpus, + "max_osl": max_osl, + } + + +def from_weka_trace( + path: str | Path, + *, + tag: str = _DEFAULT_TAG, + idle_gap_cap_seconds: float | None | object = _USE_DEFAULT, + workers: int | None = None, + content_root_seed: int | None = None, + content_tokenizer: str | None = None, + prompt_corpus: str | None = None, + max_osl: int | None = None, + num_dataset_entries: int | None = None, + max_context_length: int | None = None, + selection_out: list[SelectionStats] | None = None, +) -> ParsedGraph: + """Parse a Weka KV-cache-tester trace into a :class:`ParsedGraph`. + + Args: + path: Path to a ``.json`` file or a directory of ``.json`` files. + tag: Base tag added to every emitted ``TraceRecord``. + idle_gap_cap_seconds: Per-trace idle-gap cap (seconds) on the wall clock + that stamps each node's ``arrival_offset_us`` — gaps longer than this + are compressed to the cap (agentx ``_IdleGapTimeWarp`` parity), so the + warped trace duration and the snapshot-at-t* resume instant match + agentx. The run path resolves this from the dataset's + ``synthesis.idle_gap_cap_seconds`` (``--synthesis-idle-gap-cap``); + left unset here it defaults to 60.0. Pass ``None`` to disable warping + (raw recorded ``t``). + workers: Override worker count for parallel parsing of directory + inputs. ``None`` defers to the + ``AIPERF_DATASET_WEKA_GRAPH_PARALLEL_WORKERS`` env var (``0``/unset + auto-scales like AgentX). + content_root_seed: Run ``--random-seed`` pinned through to the + real-content synthesizer so the corpus pool + per-hash reseed key + derive from the run seed even in a ``spawn``-started parse worker + that did not inherit the parent's seeded global RNG ``_manager``. + ``None`` keeps the ambient global seed (bootstrap-seeded in-process + path). See :class:`CorpusContentSynthesizer`. + content_tokenizer: Tokenizer name the content synthesizer decodes block / + partial-tail token IDs with. Derived from the RUN config by + :func:`resolve_graph_content_tokenizer` (no env override) and threaded + here so the synthesized content matches what the run dispatches and + counts against. ``None`` (direct adapter / tooling with no run config) + falls back to ``builtin``. Threaded like ``content_root_seed`` so + spawn-started parse workers decode identically to the in-process path. + prompt_corpus: Named corpus the real-content synthesizer draws block / + partial-tail text from. Threaded through to + :class:`CorpusContentSynthesizer` like ``content_root_seed``. ``None`` + defaults to ``"coding"``. + max_osl: ``--synthesis-max-osl`` cap. Threaded through so the content + synthesizer caps each top-level chain's native ``LlmNode.max_tokens`` + to ``min(recorded out, max_osl)`` (agentx + ``_cap_output`` parity); subagent-body chains are left uncapped. + ``None`` (default) leaves the recorded ``out`` uncapped. Threaded + like ``content_root_seed`` so spawn-started parse workers cap + identically to the in-process path. + num_dataset_entries: ``--num-dataset-entries`` cap on the number of + distinct traces built (the graph-plane fix for ai-dynamo/aiperf#1106, + where the loader built every trace and dispatch cloned to fill lanes). + ``None`` (default) builds every eligible trace. Applies to the + multi-trace directory and HF corpus routes only; a single ``.json`` + file is one trace and is never capped. + max_context_length: ``--max-context-length`` per-trace ceiling on peak + context (input+output tokens on the largest single request, computed + schema-only via :func:`weka_trace_peak_context`). Traces whose peak + exceeds it are rejected BEFORE the build; the first + ``num_dataset_entries`` eligible traces are then kept (filter THEN + cap). ``None`` (default) applies no context filter. + selection_out: Optional sink; when provided AND a selection knob is set, + the single :class:`SelectionStats` for the filter-then-cap scan is + appended for the caller's load-summary report. Untouched when both + knobs are ``None`` (no selection performed). + + Returns: + A :class:`ParsedGraph` with one :class:`TraceRecord` per Weka trace + (one trace per file for directory input). + + Raises: + WekaTraceAdapterError: any conversion failure (see the + ``WekaTraceAdapterError`` subclasses defined in ``trace_models`` and + re-exported here). + """ + parse_kwargs = _resolve_parse_kwargs( + tag=tag, + idle_gap_cap_seconds=idle_gap_cap_seconds, + content_root_seed=content_root_seed, + content_tokenizer=content_tokenizer, + prompt_corpus=prompt_corpus, + max_osl=max_osl, + ) + + # HuggingFace-direct path: ``--graph `` that is not an existing + # filesystem path loads the corpus straight from the `datasets` library and + # parses each row dict in-memory through the SAME shared core as files. No + # temporary JSON is materialized. + if _looks_like_hf_dataset_id(path): + return _from_hf_dataset( + _hf_dataset_id_str(path), + workers=workers, + parse_kwargs=parse_kwargs, + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + max_osl=max_osl, + selection_out=selection_out, + ) + + p = Path(path) + if p.is_dir(): + files = _list_directory_json_files(p) + files = _apply_weka_selection( + files, + source=str(p), + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + max_osl=max_osl, + selection_out=selection_out, + ) + return trace_parallel.parse_items( + trace_parallel.file_work_items(files), + source_label=str(p), + item_count=len(files), + workers=workers, + parse_kwargs=parse_kwargs, + ) + return _parse_single_file(p, **parse_kwargs) + + +def _list_directory_json_files(directory: Path) -> list[Path]: + """Return the directory's ``*.json`` children sorted by name (stable).""" + files = sorted( + c for c in directory.iterdir() if c.is_file() and c.suffix.lower() == ".json" + ) + if not files: + raise EmptyWekaTraceError( + f"weka_trace directory {str(directory)!r} contains no .json files" + ) + return files + + +def _parse_single_file( + path: Path, + *, + tag: str, + idle_gap_cap_seconds: float | None = None, + content_root_seed: int | None = None, + content_tokenizer: str | None = None, + prompt_corpus: str | None = None, + max_osl: int | None = None, +) -> ParsedGraph: + """Single-file path: validate schema, build topology, build records. + + Split out so :func:`from_weka_trace` can pass the same per-file logic into + the parallel directory parser without a self-recursive lambda that would + loop through the directory branch. Defers dict -> trace conversion to + :func:`_parse_trace_dict`, the shared core used by the HF-source path too. + """ + try: + raw = orjson.loads(path.read_bytes()) + except orjson.JSONDecodeError as e: + raise WekaSchemaError(f"{path}: invalid JSON: {e}") from e + + return _parse_trace_dict( + raw, + source=str(path), + tag=tag, + idle_gap_cap_seconds=idle_gap_cap_seconds, + content_root_seed=content_root_seed, + content_tokenizer=content_tokenizer, + prompt_corpus=prompt_corpus, + max_osl=max_osl, + ) + + +def _from_hf_dataset( + repo_id: str, + *, + workers: int | None = None, + parse_kwargs: dict[str, Any], + num_dataset_entries: int | None = None, + max_context_length: int | None = None, + max_osl: int | None = None, + selection_out: list[SelectionStats] | None = None, +) -> ParsedGraph: + """Load a weka corpus directly from a HuggingFace dataset, no temp files. + + Resolves row dicts via :func:`_load_hf_rows` as a streaming iterator and + parses each through :func:`parse_items` — the SAME process pool the + directory path uses, over the SAME shared core + (:func:`_parse_trace_dict`). The per-row :class:`ParsedGraph` objects are + merged with the SAME multi-graph merge the directory path uses, so + HF-sourced and directory-sourced corpora produce identical structure. + Above ``WEKA_GRAPH_PARALLEL_THRESHOLD`` rows the parse fans out across + worker processes (so per-row builds are not a serial bottleneck); at or + below it the parse stays serial in-process. + + ``split`` / ``revision`` are resolved from the + ``Environment.DATASET.WEKA_HF_*`` knobs. HuggingFace rows are always read in + streaming mode; bound the ingested rows with a split slice + (``WEKA_HF_SPLIT=train[:N]``) or the loader cap (``--num-dataset-entries``). + + ``datasets`` is imported lazily here so non-HF runs never pay its import. + """ + from aiperf.common.environment import Environment + + split = Environment.DATASET.WEKA_HF_SPLIT + revision = Environment.DATASET.WEKA_HF_REVISION + + rows = _load_hf_rows(repo_id, split=split, revision=revision) + if num_dataset_entries is not None or max_context_length is not None: + kept_rows, stats = _select_weka_rows( + rows, + repo_id=repo_id, + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + max_osl=max_osl, + ) + if selection_out is not None: + selection_out.append(stats) + return trace_parallel.parse_items( + trace_parallel.row_work_items(kept_rows, repo_id), + source_label=repo_id, + item_count=len(kept_rows), + workers=workers, + parse_kwargs=parse_kwargs, + ) + return trace_parallel.parse_items( + trace_parallel.row_work_items(rows, repo_id), + source_label=repo_id, + workers=workers, + parse_kwargs=parse_kwargs, + ) + + +def _load_hf_rows( + repo_id: str, + *, + split: str, + revision: str | None, +) -> Iterator[dict[str, Any]]: + """Yield Weka trace row dicts from a HuggingFace dataset lazily. + + The published Weka corpora are large JSONL splits. Always using HF streaming + mode avoids the Arrow/full-list materialization path; bound the streamed + rows with a split slice (``split="train[:N]"``) instead. Each yielded row is + shallow-copied to a plain ``dict`` so it is picklable for the forkserver + parse pool and detached from the dataset row view. + + A load failure is wrapped with BOTH interpretations of the arg: ``repo_id`` + only reaches here through the weka-marker HF heuristic + (:func:`_looks_like_hf_dataset_id`), which fires on any non-existent + weka-marked ``org/name`` string -- including a typo'd local path -- so the + error must not present the arg as an HF id only. + """ + from datasets import load_dataset + + try: + streamed = load_dataset(repo_id, split=split, streaming=True, revision=revision) + except Exception as e: + raise WekaTraceAdapterError( + f"no such local file or directory {repo_id!r}, and loading it as a " + f"HuggingFace weka dataset id failed: {e}. If you meant a local " + "path, check the spelling; if you meant a HuggingFace dataset id, " + "verify the repo exists and is accessible." + ) from e + for row in streamed: + yield dict(row) + + +def stream_weka_trace_segment_payloads( + source: str, + *, + tag: str = _DEFAULT_TAG, + idle_gap_cap_seconds: float | None | object = _USE_DEFAULT, + workers: int | None = None, + content_root_seed: int | None = None, + content_tokenizer: str | None = None, + prompt_corpus: str | None = None, + max_osl: int | None = None, + num_dataset_entries: int | None = None, + max_context_length: int | None = None, + selection_out: list[SelectionStats] | None = None, +) -> Iterator[Any]: + """Yield worker-built trace segment payloads for ANY weka source. + + The build-plane streaming counterpart of :func:`from_weka_trace`, now for + HF corpora AND local files/directories: the SAME run-derived content knobs + must thread into the per-item parse via ``_resolve_parse_kwargs`` so the + streamed store's content and node ordinals stay consistent with the run's + parse elsewhere in THIS process. For HF sources, ``split`` / ``revision`` + are resolved from the ``Environment.DATASET.WEKA_HF_*`` knobs. + + ``num_dataset_entries`` / ``max_context_length`` drive the SAME schema-only + filter-then-cap selection as :func:`from_weka_trace` (ai-dynamo/aiperf#1106), + applied BEFORE the payload fan-out so only kept traces are built here too -- + this is the entry the DatasetManager build plane actually drains, so the run + honors the knobs. ``selection_out`` receives the :class:`SelectionStats` when + a knob is set. + """ + parse_kwargs = _resolve_parse_kwargs( + tag=tag, + idle_gap_cap_seconds=idle_gap_cap_seconds, + content_root_seed=content_root_seed, + content_tokenizer=content_tokenizer, + prompt_corpus=prompt_corpus, + max_osl=max_osl, + ) + + item_count: int | None = None + if _looks_like_hf_dataset_id(source): + from aiperf.common.environment import Environment + + split = Environment.DATASET.WEKA_HF_SPLIT + revision = Environment.DATASET.WEKA_HF_REVISION + rows = _load_hf_rows(_hf_dataset_id_str(source), split=split, revision=revision) + if num_dataset_entries is not None or max_context_length is not None: + kept_rows, stats = _select_weka_rows( + rows, + repo_id=source, + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + max_osl=max_osl, + ) + if selection_out is not None: + selection_out.append(stats) + item_count = len(kept_rows) + items = trace_parallel.row_work_items(kept_rows, source) + else: + items = trace_parallel.row_work_items(rows, source) + else: + p = Path(source) + files = _list_directory_json_files(p) if p.is_dir() else [p] + files = _apply_weka_selection( + files, + source=source, + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + max_osl=max_osl, + selection_out=selection_out, + ) + item_count = len(files) + items = trace_parallel.file_work_items(files) + + yield from trace_parallel.iter_item_segment_payloads( + items, + source_label=source, + item_count=item_count, + workers=workers, + parse_kwargs=parse_kwargs, + ) + + +def _validate_weka_trace(raw: dict[str, Any], *, source: str) -> WekaTrace: + """Schema-validate one weka trace row (NO content synthesis / graph build). + + Shared by the build core (:func:`_parse_trace_dict`) and the pre-build + selection scan (:func:`_select_weka_rows`) so both surface identical errors + (hash-scope, schema, empty-requests) for the same row and the selector can + screen a corpus at schema cost only. + """ + from aiperf.dataset.graph.adapters.weka.trace_models import WekaTrace + + if not isinstance(raw, dict): + raise WekaSchemaError( + f"{source}: top-level value must be an object, got {type(raw).__name__}" + ) + + # Surface the explicit hash-scope error before Pydantic's Literal turns it + # into a generic ValidationError, so users see the precise cause. + scope = raw.get("hash_id_scope") + if scope is not None and scope not in _SUPPORTED_HASH_ID_SCOPES: + raise WekaHashScopeError( + f"{source}: hash_id_scope={scope!r} is not supported (supported " + "scopes: 'local' = per-trace hash namespace, 'global' = one hash " + "namespace shared across all traces in the corpus)" + ) + + try: + weka_trace = WekaTrace.model_validate(raw) + except ValidationError as e: + raise WekaSchemaError(f"{source}: {e}") from e + + if not weka_trace.requests: + raise EmptyWekaTraceError(f"{source}: trace {weka_trace.id!r} has no requests") + return weka_trace + + +def _apply_weka_selection( + files: list[Path], + *, + source: str, + num_dataset_entries: int | None, + max_context_length: int | None, + max_osl: int | None, + selection_out: list[SelectionStats] | None, +) -> list[Path]: + """Filter-then-cap a directory's ``.json`` files by schema-only peak context. + + No-op (returns ``files`` unchanged, nothing appended) when both knobs are + ``None`` -- the ctx-less / knob-less build path stays byte-identical. Files + arrive already sorted by name (:func:`_list_directory_json_files`), the + deterministic scan order the cap is applied over. + """ + if num_dataset_entries is None and max_context_length is None: + return files + + from aiperf.dataset.graph.adapters.shared.peak_context import ( + weka_trace_peak_context, + ) + from aiperf.dataset.graph.adapters.shared.selection import ( + filter_then_cap, + log_selection_summary, + ) + + def _candidates() -> Iterator[tuple[Path, int]]: + for f in files: + try: + raw = orjson.loads(f.read_bytes()) + except orjson.JSONDecodeError as e: + raise WekaSchemaError(f"{f}: invalid JSON: {e}") from e + trace = _validate_weka_trace(raw, source=str(f)) + yield f, weka_trace_peak_context(trace, max_osl=max_osl) + + kept, stats = filter_then_cap( + _candidates(), + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + ) + # Parent-side finalize point for the directory/file path (one caller runs per + # build); mutually exclusive with the HF-row path's own summary. + log_selection_summary( + stats, + source=source, + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + ) + if selection_out is not None: + selection_out.append(stats) + return kept + + +def _select_weka_rows( + rows: Iterator[dict[str, Any]], + *, + repo_id: str, + num_dataset_entries: int | None, + max_context_length: int | None, + max_osl: int | None, +) -> tuple[list[dict[str, Any]], SelectionStats]: + """Filter-then-cap streamed HF row dicts by schema-only peak context. + + Rows are consumed LAZILY in stream order (the deterministic scan order); the + cap short-circuits the stream so a capped load pulls only the scanned prefix. + The kept row dicts (already shallow-copied by :func:`_load_hf_rows`) are + returned for the SAME parse fan-out the full corpus would take. + """ + from aiperf.dataset.graph.adapters.shared.peak_context import ( + weka_trace_peak_context, + ) + from aiperf.dataset.graph.adapters.shared.selection import ( + filter_then_cap, + log_selection_summary, + ) + + def _candidates() -> Iterator[tuple[dict[str, Any], int]]: + for index, row in enumerate(rows): + trace = _validate_weka_trace(row, source=f"{repo_id}#{index}") + yield row, weka_trace_peak_context(trace, max_osl=max_osl) + + kept, stats = filter_then_cap( + _candidates(), + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + ) + # Parent-side finalize point for the HF-row path (one caller runs per build). + log_selection_summary( + stats, + source=repo_id, + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + ) + return kept, stats + + +def _parse_trace_dict( + raw: dict[str, Any], + *, + source: str, + tag: str, + idle_gap_cap_seconds: float | None = None, + content_root_seed: int | None = None, + content_tokenizer: str | None = None, + prompt_corpus: str | None = None, + max_osl: int | None = None, +) -> ParsedGraph: + """Shared dict -> :class:`ParsedGraph` core. + + Operates on an in-memory Weka trace row ``dict`` already matching the schema + (the same dict a ``.json`` file or a HuggingFace dataset row contains), so + the file path (read file -> dict) and the HF path (``load_dataset`` -> row + dict) feed the identical topology / aux / flat-chain / content-synthesis + pipeline. Only the SOURCE of the dict differs. + + ``source`` is a human-readable origin label (a file path or an + ``org/name#row`` HF locator) used only for error messages and + ``ParsedGraph.source_path``. + """ + weka_trace = _validate_weka_trace(raw, source=source) + + from aiperf.common.tokenizer import BUILTIN_TOKENIZER_NAME + from aiperf.dataset.graph.adapters.weka.trie_build import ( + build_trie_graph, + ) + from aiperf.dataset.graph.models import TraceRecord + + parsed, pool = build_trie_graph( + weka_trace, + tokenizer_name=content_tokenizer or BUILTIN_TOKENIZER_NAME, + prompt_corpus=prompt_corpus or "coding", + root_seed=content_root_seed, + idle_gap_cap_seconds=idle_gap_cap_seconds, + max_osl=max_osl, + ) + # The trivial trie IR is a single top graph with no subgraphs; attach the + # trace record (with the base tag) so the schedule plane + worker address + # its nodes by trace id + node ordinal, and surface the + # SegmentPool so the build plane drains it into a GraphSegmentBackingStore. + trie_trace = TraceRecord(id=weka_trace.id, tags=[tag]) + return msgspec.structs.replace( + parsed, + traces=[trie_trace], + segment_pool=pool, + ) + + +class WekaTraceAdapter: + """Weka KV-cache-tester trace v1 graph adapter. + + See module docstring for topology / prompt-synthesis details. Implements + :class:`GraphAdapterProtocol` (registered under ``graph_adapter.weka_trace`` + in ``plugins.yaml`` with ``detection_priority: 85``). + """ + + @classmethod + def can_load(cls, path: Path) -> bool: + """Return True if ``path`` looks like a Weka trace file or directory. + + HuggingFace id: a non-existent ``org/name`` path (no workload file + extension) that carries the weka corpus marker is recognized as a weka + HF dataset id so ``detect_format`` resolves it to ``weka_trace`` and the + load goes straight through ``datasets.load_dataset`` — no file is read. + An arbitrary ``org/name`` WITHOUT the marker is NOT treated as weka. + Single ``.json`` file: read the first ~4 KB, parse as JSON, require the + discriminator key set ``{id: str, models: list, block_size: int, + hash_id_scope: "local"|"global", requests: list}``, AND reject objects + carrying any top-level key outside the WekaTrace field set. Directory: pick the + lexicographically-first ``*.json`` child and apply the file check. + Sniffs are bounded — no full-file parse on detection. + """ + if _looks_like_hf_dataset_id(path): + return True + if path.is_dir(): + try: + candidates = sorted( + c + for c in path.iterdir() + if c.is_file() and c.suffix.lower() == ".json" + ) + except OSError: + return False + if not candidates: + return False + return cls._file_matches(candidates[0]) + if path.suffix.lower() == ".json": + return cls._file_matches(path) + return False + + @classmethod + def _file_matches(cls, path: Path) -> bool: + """Bounded JSON sniff on a single ``.json`` file. + + We attempt a streaming decode of the first ``_DETECTION_SNIFF_BYTES``; + if that slice contains a complete JSON object we evaluate it directly. + Otherwise we fall back to a full-file parse only when the head slice + already shows the expected discriminator-key prefix, so detection + remains O(1) on non-Weka inputs. + """ + try: + with path.open("rb") as f: + head = f.read(_DETECTION_SNIFF_BYTES) + except OSError: + return False + if not head: + return False + try: + doc = orjson.loads(head) + except orjson.JSONDecodeError: + # Head slice was truncated mid-document. Only do a full parse if + # the head contains all five discriminator-key prefixes; this keeps + # detection cheap on unrelated large JSON files. + if not _head_has_signature_keys(head): + return False + try: + with path.open("rb") as f: + doc = orjson.loads(f.read()) + except (OSError, orjson.JSONDecodeError): + return False + return _is_weka_trace_object(doc) + + @classmethod + def parse(cls, path: Path, ctx: GraphParseContext | None = None) -> ParsedGraph: + """Convert ``path`` into a :class:`ParsedGraph` via :func:`from_weka_trace`. + + ``ctx`` carries the run-derived knobs (seed / tokenizer / corpus / + max_osl / idle-gap cap), each forwarded ONLY when set so a ctx-less + parse matches the :func:`from_weka_trace` defaults byte-for-byte. + ``idle_gap_cap_seconds`` is TRI-STATE: ``UNSET`` keeps the entry's + ``_USE_DEFAULT`` default, while an explicit ``None`` forwards verbatim and + DISABLES warping (the user's ``synthesis.idle_gap_cap_seconds: null``). + The run tokenizer trust/revision publish to the loader-preload env + before any callbacks are built (:func:`publish_ctx_tokenizer_env`). + ``ctx.num_dataset_entries`` / ``ctx.max_context_length`` forward the + schema-only filter-then-cap trace selection (ai-dynamo/aiperf#1106) when + set; the run's build plane drains :func:`stream_weka_trace_segment_payloads` + (also selection-aware), so both this oracle parse and the run honor them. + """ + publish_ctx_tokenizer_env(ctx) + kwargs: dict[str, Any] = {} + if ctx is not None: + if ctx.content_root_seed is not None: + kwargs["content_root_seed"] = ctx.content_root_seed + if ctx.content_tokenizer is not None: + kwargs["content_tokenizer"] = ctx.content_tokenizer + if ctx.prompt_corpus is not None: + kwargs["prompt_corpus"] = ctx.prompt_corpus + if ctx.max_osl is not None: + kwargs["max_osl"] = ctx.max_osl + if ctx.idle_gap_cap_seconds is not UNSET: + kwargs["idle_gap_cap_seconds"] = ctx.idle_gap_cap_seconds + if ctx.num_dataset_entries is not None: + kwargs["num_dataset_entries"] = ctx.num_dataset_entries + if ctx.max_context_length is not None: + kwargs["max_context_length"] = ctx.max_context_length + return from_weka_trace(path, **kwargs) + + +# Discriminator-key prefixes we require to appear in the head sniff before +# we are willing to fall back to a full-file parse on a truncated head. +_SIGNATURE_KEY_PREFIXES: tuple[bytes, ...] = ( + b'"id"', + b'"models"', + b'"block_size"', + b'"hash_id_scope"', + b'"requests"', +) + + +def _head_has_signature_keys(head: bytes) -> bool: + return all(prefix in head for prefix in _SIGNATURE_KEY_PREFIXES) + + +def _hf_dataset_id_str(arg: str | Path) -> str: + """Forward-slash string form of a candidate HF ``org/name`` repo id. + + An HF id that round-trips through ``Path`` flips its ``/`` to ``\\`` on + Windows, which breaks the single-slash repo-id shape check and the + canonical-repo prefix compare, and would 404 against the hub. HF repo ids + are always forward-slash, so normalize before matching or loading. + """ + s = str(arg) + if IS_WINDOWS: + s = s.replace("\\", "/") + return s + + +def _looks_like_hf_dataset_id(arg: str | Path) -> bool: + """True if ``arg`` is a weka HuggingFace ``org/name`` repo id, not a local path. + + The published weka HF corpora are referenced by ``--graph `` and + carry the weka corpus marker in the repo name (e.g. + ``semianalysisai/cc-traces-weka-061526``). We treat the arg as a weka HF id + only when ALL of the following hold: + + * it is NOT an existing filesystem path (an existing path -- even one shaped + like ``org/name`` -- keeps the file/dir behavior), + * its leading path component does NOT exist as a local directory (a typo'd + relative path like ``traces/weka-061526`` under an existing ``traces/`` + dir is a local-path mistake, not an HF repo id), + * it has no recognized workload file extension, + * it matches the repo-id shape (single slash, ``org/name`` -- which also + rejects path-like markers such as a ``./`` prefix or a trailing ``/``), + and + * it carries the case-insensitive weka marker (:data:`_HF_WEKA_MARKER`). + + The marker is the tightening: an arbitrary ``org/name`` string with no weka + marker (``meta-llama/Llama-3``, ``openai/gpt``, ``a/b``) is NOT a weka HF id + and stays on the unchanged linear pipeline rather than being sent to + ``datasets.load_dataset``. When the heuristic DOES fire and the HF load then + fails, :func:`_load_hf_rows` raises with both interpretations (local path vs + HF id) so a misrouted arg is still diagnosable. + """ + s = _hf_dataset_id_str(arg) + p = Path(s) + if p.exists(): + return False + lowered = s.lower() + if any(lowered.endswith(suffix) for suffix in _HF_EXCLUDED_SUFFIXES): + return False + if not _HF_REPO_ID_RE.match(s): + return False + if _HF_WEKA_MARKER not in lowered: + return False + # The regex guarantees exactly one slash, so ``p.parent`` is the leading + # component; if it exists locally the arg is a typo'd local path. + return not p.parent.exists() + + +def _is_weka_trace_object(doc: Any) -> bool: + """True only for a strict on-disk weka :class:`WekaTrace` discriminator. + + Requires the five discriminator keys with their expected types + (``id: str``, ``models: list``, ``block_size: int`` (not bool), + ``hash_id_scope`` in :data:`_SUPPORTED_HASH_ID_SCOPES`, ``requests: list``) + AND rejects any object carrying a top-level key outside + :data:`_WEKA_ALLOWED_KEYS`. The foreign-key rejection is what prevents a + mooncake / sharegpt / synthetic object that merely happens to share the + five weka keys from being misclassified as a weka graph workload -- + mirroring ``WekaTrace``'s ``extra="forbid"`` config. + """ + if not isinstance(doc, dict): + return False + if not isinstance(doc.get("id"), str): + return False + if not isinstance(doc.get("models"), list): + return False + if not isinstance(doc.get("block_size"), int) or isinstance( + doc.get("block_size"), bool + ): + return False + if doc.get("hash_id_scope") not in _SUPPORTED_HASH_ID_SCOPES: + return False + if not isinstance(doc.get("requests"), list): + return False + return _WEKA_ALLOWED_KEYS.issuperset(doc.keys()) + + +__all__ = [ + "EmptyWekaTraceError", + "WekaHashScopeError", + "WekaSchemaError", + "WekaTraceAdapter", + "WekaTraceAdapterError", + "from_weka_trace", +] diff --git a/src/aiperf/dataset/graph/adapters/weka/trace_models.py b/src/aiperf/dataset/graph/adapters/weka/trace_models.py new file mode 100644 index 0000000000..326abab8c7 --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/weka/trace_models.py @@ -0,0 +1,248 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Adapter-private Pydantic models for the Weka KV-cache-tester trace format. + +``WekaSubagentEntry.requests`` is ``list[WekaRequest]``, allowing inner +subagent markers: subagent-within-subagent nesting, which the trie builder +flattens recursively into ordinary flat ``LlmNode``s. + +These models are private to the adapter; nothing outside the weka adapter +imports them, so this file alone defines the adapter's input contract. +""" + +from __future__ import annotations + +from typing import Annotated, Any, Literal, TypeAlias + +from pydantic import ConfigDict, Field + +from aiperf.common.models import AIPerfBaseModel + + +class WekaTraceAdapterError(ValueError): + """Base error raised when a Weka trace cannot be converted to a ParsedGraph.""" + + +class EmptyWekaTraceError(WekaTraceAdapterError): + """The file/directory/HF source yields zero usable traces. + + Raised when a trace's ``requests`` list is empty, or when a directory or + HuggingFace source produces no parsable traces (no ``.json`` files, zero + rows, or an empty merge result). + """ + + +class WekaSchemaError(WekaTraceAdapterError): + """A Weka trace failed Pydantic schema validation. + + Wraps ``pydantic.ValidationError`` with the offending file path and trace + ``id`` for clear CLI surfacing. + """ + + +class WekaHashScopeError(WekaTraceAdapterError): + """The trace declares an unrecognized ``hash_id_scope``. + + Supported scopes are ``"local"`` (per-trace hash namespace) and + ``"global"`` (one hash namespace shared across every trace in the corpus, + reproducing recorded cross-trace KV-cache sharing). Any other value is + rejected here with the precise cause rather than surfacing as a generic + Pydantic Literal error. + """ + + +class WekaNormalRequest(AIPerfBaseModel): + """One normal (``type: "n"``) API call in a Weka trace.""" + + model_config = ConfigDict(extra="forbid", populate_by_name=True) + + t: float = Field( + ge=0, + allow_inf_nan=False, + description="Request timestamp in seconds from conversation start.", + ) + type: Literal["n"] = Field(description="Discriminator: normal API call.") + model: str = Field(description="Model identifier for this request.") + input_length: int = Field(ge=0, alias="in", description="Input token count.") + output_length: int = Field(ge=0, alias="out", description="Output token count.") + hash_ids: list[Annotated[int, Field(ge=0)]] = Field( + default_factory=list, description="KV-cache block hash IDs." + ) + input_types: list[str] = Field( + default_factory=list, description="Content-type annotations for input." + ) + output_types: list[str] = Field( + default_factory=list, description="Content-type annotations for output." + ) + stop: str = Field( + default="", description="Stop reason: '', 'tool_use', 'end_turn'." + ) + api_time: float | None = Field( + default=None, + ge=0, + allow_inf_nan=False, + description="Server processing time in seconds.", + ) + think_time: float | None = Field( + default=None, + ge=0, + allow_inf_nan=False, + description="Client delay in seconds before this request.", + ) + + +class WekaStreamingRequest(AIPerfBaseModel): + """One streaming (``type: "s"``) API call in a Weka trace. + + Structurally identical to :class:`WekaNormalRequest` except for the + discriminator value and an optional ``ttft`` field (recorded + time-to-first-token in seconds). + """ + + model_config = ConfigDict(extra="forbid", populate_by_name=True) + + t: float = Field( + ge=0, + allow_inf_nan=False, + description="Request timestamp in seconds from conversation start.", + ) + type: Literal["s"] = Field(description="Discriminator: streaming API call.") + model: str = Field(description="Model identifier for this request.") + input_length: int = Field(ge=0, alias="in", description="Input token count.") + output_length: int = Field(ge=0, alias="out", description="Output token count.") + hash_ids: list[Annotated[int, Field(ge=0)]] = Field( + default_factory=list, description="KV-cache block hash IDs." + ) + input_types: list[str] = Field( + default_factory=list, description="Content-type annotations for input." + ) + output_types: list[str] = Field( + default_factory=list, description="Content-type annotations for output." + ) + stop: str = Field( + default="", description="Stop reason: '', 'tool_use', 'end_turn'." + ) + api_time: float | None = Field( + default=None, + ge=0, + allow_inf_nan=False, + description="Server processing time in seconds.", + ) + think_time: float | None = Field( + default=None, + ge=0, + allow_inf_nan=False, + description="Client delay in seconds before this request.", + ) + ttft: float | None = Field( + default=None, + ge=0, + allow_inf_nan=False, + description="Recorded time-to-first-token in seconds.", + ) + + +class WekaSubagentEntry(AIPerfBaseModel): + """A ``type: "subagent"`` marker with its nested inner requests. + + The parent's next ``WekaNormalRequest`` in the outer list is understood + to occur after this subagent completes (when ``status != "async_launched"``). + Inner ``requests`` may themselves contain :class:`WekaSubagentEntry` markers; + the trie builder flattens each recursively into ordinary flat ``LlmNode``s. + """ + + model_config = ConfigDict(extra="forbid") + + t: float = Field( + ge=0, + allow_inf_nan=False, + description="Spawn timestamp in seconds from conversation start.", + ) + type: Literal["subagent"] = Field(description="Discriminator: subagent marker.") + agent_id: str = Field(description="Opaque subagent identifier, e.g. 'agent_001'.") + subagent_type: str = Field(description="Subagent type, e.g. 'Explore'.") + duration_ms: int | None = Field( + default=None, + ge=0, + description="Wall-clock duration of the subagent. None for subagents " + "with status='async_launched' (telemetry not captured).", + ) + total_tokens: int | None = Field( + default=None, + ge=0, + description="Total tokens across all subagent inner requests. None " + "for status='async_launched'.", + ) + tool_use_count: int | None = Field( + default=None, + ge=0, + description="Tool calls made by the subagent. None for " + "status='async_launched'.", + ) + status: str = Field(description="'completed', 'async_launched', or other status.") + requests: list[WekaRequest] = Field( + description="Inner requests of the subagent. May contain nested " + "WekaSubagentEntry markers (subagent-within-subagent nesting); the " + "trie builder flattens them recursively.", + ) + models: list[str] = Field(description="Models used by the subagent.") + tool_tokens: int = Field( + default=0, ge=0, description="Subagent's tools prefix token count." + ) + system_tokens: int = Field( + default=0, ge=0, description="Subagent's system prefix token count." + ) + + +WekaRequest: TypeAlias = Annotated[ + WekaNormalRequest | WekaStreamingRequest | WekaSubagentEntry, + Field(discriminator="type"), +] + + +class WekaTrace(AIPerfBaseModel): + """A single Weka trace file.""" + + model_config = ConfigDict(extra="forbid") + + id: str = Field(description="Trace identifier (session ID).") + models: list[str] = Field(description="Models used in the trace.") + block_size: int = Field(gt=0, description="Cache block size in tokens.") + hash_id_scope: Literal["local", "global"] = Field( + description=( + "Hash ID namespace scope. 'local' scopes hashes per-trace: equal " + "hash ids in different traces are different logical blocks and " + "synthesize different bytes. 'global' shares ONE hash namespace " + "across every trace in the corpus: equal hash ids synthesize " + "byte-identical content, reproducing recorded cross-trace " + "KV-cache sharing. Any other value is rejected with " + "WekaHashScopeError." + ) + ) + tool_tokens: int = Field(default=0, ge=0, description="Tools prefix token count.") + system_tokens: int = Field( + default=0, ge=0, description="System prefix token count." + ) + requests: list[WekaRequest] = Field( + description="Interleaved normal/streaming requests and subagent markers." + ) + totals: dict[str, Any] | None = Field( + default=None, description="Optional trace-level summary; opaque." + ) + + +# Resolve forward reference inside WekaSubagentEntry.requests. +WekaSubagentEntry.model_rebuild() + + +__all__ = [ + "EmptyWekaTraceError", + "WekaHashScopeError", + "WekaNormalRequest", + "WekaRequest", + "WekaSchemaError", + "WekaStreamingRequest", + "WekaSubagentEntry", + "WekaTrace", + "WekaTraceAdapterError", +] diff --git a/src/aiperf/dataset/graph/adapters/weka/trace_parallel.py b/src/aiperf/dataset/graph/adapters/weka/trace_parallel.py new file mode 100644 index 0000000000..a789f735d9 --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/weka/trace_parallel.py @@ -0,0 +1,629 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Parallel multi-process parse dispatch for the Weka graph adapter. + +Every weka route (local directory, HuggingFace corpus, streaming build plane) +funnels work items through :func:`_map_items`: at or below +``WEKA_GRAPH_PARALLEL_THRESHOLD`` (default 8) items parse serially in-process; +above it they stream through a bounded, ordered forkserver pool window. Each +worker parses one item independently and returns a msgpack-encoded +``ParsedGraph`` (or its trie segment payloads) — Pydantic-pickle has +historically broken cross-process transfer of ``ParsedGraph`` instances, so +encoding at the worker boundary side-steps that. The main process decodes and +either merges into a single multi-graph :class:`ParsedGraph` whose traces are +sorted by trace id (byte-deterministic) or streams per-trace payloads into the +unified segment store. +""" + +from __future__ import annotations + +import contextlib +import itertools +import multiprocessing +import os +import signal +import threading +from collections import deque +from collections.abc import Callable, Iterable, Iterator +from dataclasses import dataclass +from multiprocessing import shared_memory +from pathlib import Path +from typing import Any + +import numpy as np + +from aiperf.common.utils import allow_daemon_children +from aiperf.dataset.graph.adapters.weka.trace_models import ( + EmptyWekaTraceError, +) +from aiperf.dataset.graph.codecs import ( + decode_parsed_graph_msgpack, + encode_parsed_graph_msgpack, +) +from aiperf.dataset.graph.merge import merge_parsed_graphs +from aiperf.dataset.graph.models import ( + ParsedGraph, +) + +_POOL_JOIN_TIMEOUT_S = 10.0 + + +def _loader_pool_context(tokenizer_name: str | None): + """Return the forkserver (Linux) / spawn (macOS) mp context for the pool. + + The default ``fork`` start method DEADLOCKS here: the parent process has + already loaded the HF real-content tokenizer and exercised its Rust rayon + thread pool, so forked workers inherit broken rayon state and hang at ~2% + CPU while holding the parent's threads and many GB RSS. A forkserver helper + is a fresh interpreter free of that state; it also preloads the configured + tokenizer into its heap so workers CoW-share one copy (see + :mod:`aiperf.dataset._mp_context` / :mod:`aiperf.dataset._tokenizer_preload`). + """ + from aiperf.dataset._mp_context import get_loader_mp_context + + return get_loader_mp_context(preload_tokenizer=tokenizer_name) + + +@dataclass(slots=True) +class _WorkerInitArgs: + """Static args shipped to every pool worker via ``Pool(initializer=...)``. + + ``shm_name``/``corpus_len`` describe the parent-built shared-memory corpus + block (int32 token ids); the worker attaches it read-only and points its + cached :class:`CorpusContentSynthesizer` at that array, so the ~600K-token + coding pool is built ONCE in the parent and CoW-shared across all workers + instead of rebuilt per worker (the rebuild-per-trace path ballooned RSS into + tens of GB and abruptly killed pool workers -> ``BrokenProcessPool``). + """ + + tokenizer_name: str | None + prompt_corpus: str + root_seed: int | None + shm_name: str | None + corpus_len: int + + +# Per-worker handle to the attached shared-memory corpus; kept alive for the +# worker's lifetime so the mmap is not released between tasks. +_worker_corpus_shm: shared_memory.SharedMemory | None = None + + +def _install_hard_exit_on_sigterm() -> None: + """Replace the worker SIGTERM handler with ``os._exit(0)``. + + Pool teardown (``terminate()``) SIGTERMs each worker, but the worker's + CoW-shared HF tokenizer carries Rust ``rayon`` threads that do not unwind on + SIGTERM, so the default Python finalizer path wedges. Workers are stateless + (the parent owns shared memory and collects results), so a hard exit is the + right behavior. ``signal.signal`` only works on the main thread, so a + ``ValueError`` (unit test invoking ``_init_worker`` off-thread) is ignored. + """ + + def _hard_exit(_signum, _frame): # noqa: ANN001 + os._exit(0) + + with contextlib.suppress(ValueError): + signal.signal(signal.SIGTERM, _hard_exit) + + +def _init_worker(args: _WorkerInitArgs) -> None: + """Pool worker init: attach the shared corpus + warm the synthesizer cache. + + Builds (and caches) one :class:`CorpusContentSynthesizer` for the run's + ``(tokenizer, corpus, seed)`` CONSTRUCTED ON the parent-built shared-memory + array: passing ``shared_corpus`` into ``get_or_build_synthesizer`` makes the + generator skip its own corpus build entirely, so the worker never pays the + ~600K-token pool build only to discard it for the shm attachment. Every + trace this worker parses then reuses that single CoW-shared pool. Fail-soft: + any error here leaves the cache unpopulated and the per-trace path builds a + private synthesizer on demand (correct, just not memory-optimal). + """ + global _worker_corpus_shm + + _install_hard_exit_on_sigterm() + os.environ.setdefault("HF_HUB_OFFLINE", "1") + os.environ.setdefault("TRANSFORMERS_OFFLINE", "1") + + if args.tokenizer_name is None or args.shm_name is None: + return + + try: + from aiperf.dataset.graph.adapters.shared.content import ( + get_or_build_synthesizer, + ) + + shm = shared_memory.SharedMemory(name=args.shm_name) + _worker_corpus_shm = shm + corpus = np.ndarray((args.corpus_len,), dtype=np.int32, buffer=shm.buf) + + get_or_build_synthesizer( + args.tokenizer_name, + prompt_corpus=args.prompt_corpus, + root_seed=args.root_seed, + shared_corpus=corpus, + ) + except Exception: # init must never crash the worker + _worker_corpus_shm = None + + +def _shutdown_pool(pool, *, timeout_s: float = _POOL_JOIN_TIMEOUT_S) -> None: + """Drain a ``multiprocessing.Pool`` without the SIGTERM teardown wedge. + + ``with Pool(...)`` exit calls ``terminate()`` which SIGTERMs every worker; + the CoW-shared HF tokenizer's ``rayon`` threads do not unwind on SIGTERM, so + ``terminate()``+``join()`` hangs. Graceful path: ``close()`` lets each worker + drain and exit via the hard-exit handler / normal sentinel; ``join()`` then + returns promptly. Falls back to ``terminate()`` (bounded by ``timeout_s``) + if a worker wedges anyway, so a stuck worker never blocks the CLI. + """ + pool.close() + done = threading.Event() + + def _wait() -> None: + try: + pool.join() + finally: + done.set() + + threading.Thread(target=_wait, daemon=True).start() + if done.wait(timeout=timeout_s): + return + pool.terminate() + done.wait(timeout=timeout_s) + + +def _build_shared_corpus( + tokenizer_name: str | None, + *, + prompt_corpus: str, + root_seed: int | None, +) -> tuple[shared_memory.SharedMemory | None, int]: + """Build the deterministic corpus ONCE in the parent into shared memory. + + Returns ``(shm, corpus_len)``; ``(None, 0)`` when real-content synthesis is + off (``tokenizer_name is None``) so workers skip corpus attachment. The + corpus is byte-identical to what each worker would build for the same + ``(tokenizer, corpus, seed)``, so determinism / parity is preserved. + """ + if tokenizer_name is None: + return None, 0 + + from aiperf.dataset.graph.adapters.shared.content import ( + get_or_build_synthesizer, + ) + + synth = get_or_build_synthesizer( + tokenizer_name, prompt_corpus=prompt_corpus, root_seed=root_seed + ) + corpus_tokens = synth.corpus_tokens() + corpus_arr = np.ascontiguousarray(corpus_tokens, dtype=np.int32) + corpus_len = int(corpus_arr.shape[0]) + if corpus_len == 0: + return None, 0 + shm = shared_memory.SharedMemory( + create=True, size=corpus_len * np.dtype(np.int32).itemsize + ) + np.ndarray((corpus_len,), dtype=np.int32, buffer=shm.buf)[:] = corpus_arr + return shm, corpus_len + + +def _run_pool_streaming( + worker_fn: Callable[[Any], Any], + work_items: Iterable[Any], + *, + workers: int, + root_seed: int | None, + content_tokenizer: str | None = None, + prompt_corpus: str | None = None, + prefetch_multiplier: int | None = None, + item_timeout_s: float | None = None, + timeout_hint: str | None = None, +) -> Iterator[Any]: + """Dispatch ``worker_fn`` across a forkserver ``Pool``, YIELDING each result + as ``imap``-ordered completion produces it. + + Each yielded value is whatever ``worker_fn`` returns -- a msgpack + ``ParsedGraph`` blob for the merged consumer, or a pickled + ``TraceSegmentPayload`` list for the streaming consumer. An eager + ``list(pool.imap(...))`` would hold EVERY worker's result in memory at once, + and the caller would then decode ALL of them into a list before merging. For + a large real corpus (393 weka traces, each with full real-content synthesis + -- multi-turn message lists + decoded token text) that peak is what blows + host RAM. This generator instead streams through a bounded ordered + apply-async window, so the pool feeder cannot consume the full input iterator + before first result. The consumer can decode -> write-to-mmap -> DROP each + result before the next arrives, keeping resident memory flat at ~one trace. + + Builds the shared-memory corpus in the parent and opens a raw ``ctx.Pool`` + (NOT ``ProcessPoolExecutor`` -- the latter raises ``BrokenProcessPool`` on + any worker death, which the rayon-threaded weka workers hit during teardown; + a raw ``Pool`` + graceful :func:`_shutdown_pool` tolerates it). AIPerf + services are daemonic and Python refuses to spawn children from a daemon, + so the daemon flag is cleared around the pool (restored in ``finally``). + + Adapter-neutral tuning: ``prefetch_multiplier`` / ``item_timeout_s`` / + ``timeout_hint`` default to ``None`` meaning "use the weka + ``WEKA_GRAPH_PARALLEL_*`` settings + weka error message" (the historical + behavior every weka caller relies on). The dynamo session-tree build passes + its own ``DYNAMO_GRAPH_PARALLEL_*`` values and error message so the SAME + pool lifecycle serves both adapters without a weka-coupled window/timeout. + """ + from aiperf.common.environment import Environment + from aiperf.common.tokenizer import BUILTIN_TOKENIZER_NAME + + tokenizer_name = content_tokenizer or BUILTIN_TOKENIZER_NAME + prompt_corpus = prompt_corpus or "coding" + if prefetch_multiplier is None: + prefetch_multiplier = ( + Environment.DATASET.WEKA_GRAPH_PARALLEL_PREFETCH_MULTIPLIER + ) + + shm, corpus_len = _build_shared_corpus( + tokenizer_name, prompt_corpus=prompt_corpus, root_seed=root_seed + ) + init_args = _WorkerInitArgs( + tokenizer_name=tokenizer_name, + prompt_corpus=prompt_corpus, + root_seed=root_seed, + shm_name=shm.name if shm is not None else None, + corpus_len=corpus_len, + ) + + try: + with allow_daemon_children(): + ctx = _loader_pool_context(tokenizer_name) + pool = ctx.Pool(workers, _init_worker, (init_args,)) + try: + window_size = workers * prefetch_multiplier + yield from _bounded_ordered_pool_map( + pool, + worker_fn, + work_items, + window_size, + timeout_s=item_timeout_s, + timeout_hint=timeout_hint, + ) + finally: + _shutdown_pool(pool) + finally: + if shm is not None: + shm.close() + shm.unlink() + + +def _bounded_ordered_pool_map( + pool: Any, + worker_fn: Callable[[Any], Any], + work_items: Iterable[Any], + window_size: int, + *, + timeout_s: float | None = None, + timeout_hint: str | None = None, +) -> Iterator[Any]: + """Yield pool results in input order while consuming at most one window ahead. + + Each head-of-window ``AsyncResult`` is waited on with a bounded + ``get(timeout=...)``: a raw ``multiprocessing.Pool`` repopulates a killed + worker but never completes its in-flight task, so an unbounded ``get()`` + turns a worker OOM-kill into a silent indefinite CLI hang. On expiry a + ``RuntimeError`` naming that likely cause is raised instead. Order + preservation is unchanged. + + ``timeout_s`` defaults to ``None`` meaning "read + ``WEKA_GRAPH_PARALLEL_ITEM_TIMEOUT_SECONDS``" (the weka behavior); + ``timeout_hint`` overrides the raised message so the dynamo session-tree + build can name its own env var / unit ("batch" not "trace") on expiry. + """ + from aiperf.common.environment import Environment + + if timeout_s is None: + timeout_s = Environment.DATASET.WEKA_GRAPH_PARALLEL_ITEM_TIMEOUT_SECONDS + iterator = iter(work_items) + pending: deque[Any] = deque() + exhausted = False + window = max(1, window_size) + + def _fill_window() -> None: + nonlocal exhausted + while not exhausted and len(pending) < window: + try: + item = next(iterator) + except StopIteration: + exhausted = True + return + pending.append(pool.apply_async(worker_fn, (item,))) + + _fill_window() + while pending: + result = pending.popleft() + try: + item_result = result.get(timeout=timeout_s) + except multiprocessing.TimeoutError: + raise RuntimeError( + timeout_hint + or ( + f"weka graph parse worker produced no result within " + f"{timeout_s:g}s for one trace: the worker process was most " + "likely killed mid-parse (OOM kill / external SIGKILL) -- a " + "raw multiprocessing Pool cannot complete a killed worker's " + "in-flight task. Reduce worker count " + "(AIPERF_DATASET_WEKA_GRAPH_PARALLEL_WORKERS) to lower peak " + "memory, or raise " + "AIPERF_DATASET_WEKA_GRAPH_PARALLEL_ITEM_TIMEOUT_SECONDS if a " + "single trace legitimately parses slower than the timeout." + ) + ) from None + yield item_result + _fill_window() + + +def _get_threshold() -> int: + from aiperf.common.environment import Environment + + return max(0, Environment.DATASET.WEKA_GRAPH_PARALLEL_THRESHOLD) + + +def _configured_workers(override: int | None) -> int: + if override is not None: + return override + + from aiperf.common.environment import Environment + + return Environment.DATASET.WEKA_GRAPH_PARALLEL_WORKERS + + +def _get_workers(*, item_count: int | None, override: int | None) -> int: + from aiperf.common.environment import Environment + + configured = _configured_workers(override) + if configured > 0: + resolved = configured + else: + cpu = os.cpu_count() or 1 + resolved = min( + max(cpu - 1, 1), + Environment.DATASET.WEKA_GRAPH_PARALLEL_AUTO_MAX_WORKERS, + ) + if item_count is not None: + resolved = min(resolved, item_count) + return max(1, resolved) + + +@dataclass(slots=True) +class _WorkItem: + """One parse unit: a local trace file OR an in-memory HF row. + + Exactly one of ``path`` / ``row`` is set. File items ship only the path + string (the worker reads the bytes itself — far cheaper to pickle than the + parsed dict); row items ship the already-loaded dict (HF rows have no + file to re-read). ``source`` is the human-readable origin label used in + errors and ``ParsedGraph.source_path``. + """ + + source: str + """Origin label: a file path or an ``org/name#row`` HF locator.""" + + path: str | None + """Local trace file path; ``None`` for row items.""" + + row: dict[str, Any] | None + """In-memory weka trace dict; ``None`` for file items.""" + + +def file_work_items(files: Iterable[Path]) -> Iterator[_WorkItem]: + """Wrap local trace files as work items.""" + for p in files: + yield _WorkItem(source=str(p), path=str(p), row=None) + + +def row_work_items(rows: Iterable[Any], source_prefix: str) -> Iterator[_WorkItem]: + """Wrap HF row dicts as work items with stable ``org/name#index`` labels. + + Rows arrive as plain dicts already shallow-copied by :func:`_load_hf_rows` + (``yield dict(row)``), so they are picklable for the forkserver pool and + detached from the dataset row view -- this wrapper does not re-copy. + """ + for index, row in enumerate(rows): + yield _WorkItem(source=f"{source_prefix}#{index}", path=None, row=row) + + +def _parse_item(item: _WorkItem, parse_kwargs: dict[str, Any]) -> ParsedGraph: + """Parse ONE work item through the shared per-trace core. + + This is the only per-item entry for both the serial in-process path and + the pool workers: file items go through ``_parse_single_file`` (bytes read + where the parse runs) and row items through ``_parse_trace_dict``. Workers + deliberately do NOT re-enter the public ``from_weka_trace`` dispatcher — + it would re-run HF-id and directory detection per item. + """ + from aiperf.dataset.graph.adapters.weka.trace import ( + _parse_single_file, + _parse_trace_dict, + ) + + if item.path is not None: + return _parse_single_file(Path(item.path), **parse_kwargs) + assert item.row is not None + return _parse_trace_dict(item.row, source=item.source, **parse_kwargs) + + +def _parse_item_to_msgpack(task: tuple[_WorkItem, dict[str, Any]]) -> bytes: + """Pool worker entry: parse one item, return a msgpack ``ParsedGraph`` blob. + + A single picklable ``(item, parse_kwargs)`` argument drops cleanly into + the pool; msgpack at the worker boundary side-steps the historically + broken cross-process pickling of ``ParsedGraph`` instances. + """ + item, parse_kwargs = task + return encode_parsed_graph_msgpack(_parse_item(item, parse_kwargs)) + + +def _parse_item_to_segment_payloads( + task: tuple[_WorkItem, dict[str, Any]], +) -> list[Any]: + """Pool worker entry: parse one item, return its trie segment payloads.""" + from aiperf.dataset.graph.segment_ir.store_builder import ( + iter_trace_segment_payloads, + ) + + item, parse_kwargs = task + return list(iter_trace_segment_payloads(_parse_item(item, parse_kwargs))) + + +def _prefetch_items( + items: Iterable[_WorkItem], limit: int +) -> tuple[list[_WorkItem], Iterator[_WorkItem], bool]: + """Pull at most ``limit`` items; return ``(prefetched, tail, exhausted)``. + + ``exhausted`` is True only when the source yielded fewer than ``limit`` + items, i.e. ``prefetched`` is the complete input and its length is the + exact item count (used to cap the worker fan-out). + """ + iterator = iter(items) + prefetched: list[_WorkItem] = [] + for _ in range(limit): + try: + prefetched.append(next(iterator)) + except StopIteration: + return prefetched, iterator, True + return prefetched, iterator, False + + +def _map_items( + items: Iterable[_WorkItem], + *, + worker_fn: Callable[[tuple[_WorkItem, dict[str, Any]]], Any], + local_fn: Callable[[_WorkItem], Any], + decode_fn: Callable[[Any], Any], + source_label: str, + item_count: int | None, + threshold: int | None, + workers: int | None, + parse_kwargs: dict[str, Any], +) -> Iterator[Any]: + """Map work items through the ONE serial-or-pool dispatch, in input order. + + At or below the parallel threshold every item runs through ``local_fn`` + in-process (no pool, no codec round-trip). Above it, items stream through + :func:`_run_pool_streaming` (bounded ordered window, per-item timeout, + graceful shutdown, shared-memory corpus) and each pool result passes + through ``decode_fn``. Only ``threshold + 1`` items are prefetched to pick + the path, so lazy sources are never fully consumed up front. When the item + count is known (``item_count`` or an exhausted prefetch) it caps the + worker fan-out. + """ + # EmptyWekaTraceError is already a module-top import from trace_models; + # do not add a lazy import here. + effective_threshold = max( + 0, threshold if threshold is not None else _get_threshold() + ) + prefetched, remaining, exhausted = _prefetch_items(items, effective_threshold + 1) + if not prefetched: + raise EmptyWekaTraceError( + f"weka_trace source {source_label!r} yielded zero traces" + ) + + if len(prefetched) <= effective_threshold: + for item in prefetched: + yield local_fn(item) + return + + known_count = ( + item_count + if item_count is not None + else (len(prefetched) if exhausted else None) + ) + effective_workers = _get_workers(item_count=known_count, override=workers) + tasks = ((item, parse_kwargs) for item in itertools.chain(prefetched, remaining)) + for result in _run_pool_streaming( + worker_fn, + tasks, + workers=effective_workers, + root_seed=parse_kwargs.get("content_root_seed"), + content_tokenizer=parse_kwargs.get("content_tokenizer"), + prompt_corpus=parse_kwargs.get("prompt_corpus"), + ): + yield decode_fn(result) + + +def parse_items( + items: Iterable[_WorkItem], + *, + source_label: str, + item_count: int | None = None, + threshold: int | None = None, + workers: int | None = None, + parse_kwargs: dict[str, Any] | None = None, +) -> ParsedGraph: + """Parse work items and merge into ONE multi-graph ``ParsedGraph``. + + The eager-merge consumer of :func:`_map_items`, backing the whole in-memory + ``ParsedGraph`` that :func:`from_weka_trace` / ``WekaTraceAdapter.parse`` + return (reached via ``parse_graph_workload`` and pinned by the registry/run + parse-parity tests). The DatasetManager build plane does NOT take this route + — it drains worker-built segment payloads through + :func:`iter_item_segment_payloads` instead. Pool results stream straight + into the merge in input order, so the parent never holds every blob plus + every decoded graph at once; the post-merge sort by trace id keeps the + result byte-deterministic across worker counts. + """ + kwargs = parse_kwargs or {} + per_item = _map_items( + items, + worker_fn=_parse_item_to_msgpack, + local_fn=lambda item: _parse_item(item, kwargs), + decode_fn=decode_parsed_graph_msgpack, + source_label=source_label, + item_count=item_count, + threshold=threshold, + workers=workers, + parse_kwargs=kwargs, + ) + return merge_parsed_graphs(per_item) + + +def iter_item_segment_payloads( + items: Iterable[_WorkItem], + *, + source_label: str, + item_count: int | None = None, + threshold: int | None = None, + workers: int | None = None, + parse_kwargs: dict[str, Any] | None = None, +) -> Iterator[Any]: + """Yield per-trace trie segment payloads for work items, memory-bounded. + + The streaming consumer of :func:`_map_items` — the build plane serializes + each trace's envelopes into the unified store and DROPS the payloads + before the next arrive, so resident memory stays at ~one trace regardless + of corpus size. Serial and pool paths emit the same payload shape. + """ + from aiperf.dataset.graph.segment_ir.store_builder import ( + iter_trace_segment_payloads, + ) + + kwargs = parse_kwargs or {} + + def _local(item: _WorkItem) -> list[Any]: + return list(iter_trace_segment_payloads(_parse_item(item, kwargs))) + + for payloads in _map_items( + items, + worker_fn=_parse_item_to_segment_payloads, + local_fn=_local, + decode_fn=lambda blobs: blobs, + source_label=source_label, + item_count=item_count, + threshold=threshold, + workers=workers, + parse_kwargs=kwargs, + ): + yield from payloads + + +__all__ = [ + "file_work_items", + "iter_item_segment_payloads", + "parse_items", + "row_work_items", +] diff --git a/src/aiperf/dataset/graph/adapters/weka/trie_build.py b/src/aiperf/dataset/graph/adapters/weka/trie_build.py new file mode 100644 index 0000000000..cc2a93ace8 --- /dev/null +++ b/src/aiperf/dataset/graph/adapters/weka/trie_build.py @@ -0,0 +1,458 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Trivial dependency-only ``ParsedGraph`` builder for a single Weka trace. + +:func:`build_trie_graph` is the simplest possible IR realization of a +:class:`WekaTrace`: one :class:`LlmNode` per recorded ``n``/``s`` request and +plain :class:`StaticEdge` "waits-for" dependency edges between them. There are +NO reducers / channels / subgraph / spawn / await primitives and NO +chain-detection / ``::fa``-``::aux`` classification -- the conversation +structure and the dependency structure are derived purely from the recorded +hash-id prefix tree and the recorded spawn + timing facts. + +The trie machinery itself (content-parent resolution, frozen block tags, +message assembly, covered-count ISL gate, interval-order edges) is the shared +format-agnostic core in +:mod:`~aiperf.dataset.graph.segment_ir.trie_content` +(:func:`~aiperf.dataset.graph.segment_ir.trie_content.build_trie_ir`). +This module only flattens the recorded weka requests into normalized +:class:`~aiperf.dataset.graph.segment_ir.trie_content.TrieNode`s and +assembles the resulting per-node artifacts into weka ``LlmNode``s: + +1. **Conversation structure = hash-id prefix tree.** Every request (top-level + ``n``/``s`` AND subagent-inner ``n``/``s``, recursively) is collected in + recorded ``t`` order; a request's CONTENT-PARENT is its longest full-prefix + (else branch-point) predecessor + (:func:`~aiperf.dataset.graph.segment_ir.trie_content.resolve_content_parents`). + +2. **Segments (message units).** Each covered block gets a frozen + ``(role, starts_new_message)`` tag + (:func:`~aiperf.dataset.graph.segment_ir.trie_content.assign_block_tags`); + :func:`~aiperf.dataset.graph.segment_ir.trie_content.assemble_messages` + groups those tags into messages and emits one content-addressed + :class:`SegmentPool` entry per message, chained root->tip. Because the tags + are frozen per trie position, a shared block prefix yields identical message + ids -- a real KV-cache prefix. The node's own assistant output is appended + as one ``"assistant"`` segment sized to the recorded ``out`` so + the successor prompt bytes stay block-exact. + +3. **Dependency edges (interval order).** R's incoming edges come from + :func:`~aiperf.dataset.graph.segment_ir.interval_order.build_interval_edges`: a candidate ``A`` yields an edge into R iff + ``A`` FINISHED before R STARTED (raw ``A.t + A.api_time <= R.t``) AND + ``rank(A) < rank(R)``, after async exclusion and a frontier (transitive- + reduction) filter. The latest-ending frontier predecessor carries the warped + end-to-start delay; every other is an AND-fan-in wait (delay 0). With NO + finished-before cause R is wired ``StaticEdge(START, R)`` with + ``min_start_delay_us = R.start * 1e6`` -- UNLESS R's recorded causal parent + (spawner / chain-prev, stamped in :func:`_walk`) was still IN FLIGHT at R's + start, in which case + :func:`~aiperf.dataset.graph.segment_ir.interval_order.apply_start_anchors` + collapses R's incoming edges to one start-anchored ``StaticEdge(parent, R)`` + (``delay_after_predecessor_start_us`` = warped start-to-start gap) so recorded + mid-flight concurrency tracks the parent's dispatch instead of freezing to the + recorded wall clock. + +4. **Concurrency is emergent.** Two requests sharing a cause with no edge + between them stay edge-free; the builder adds an inter-sibling edge ONLY when + the start-anchor pass wires an in-flight causal parent (point 3). +""" + +from __future__ import annotations + +from aiperf.common.constants import MICROS_PER_SECOND +from aiperf.dataset.graph.adapters.shared.output_cap import wire_output_cap +from aiperf.dataset.graph.adapters.weka.trace_models import ( + EmptyWekaTraceError, + WekaSubagentEntry, + WekaTrace, + WekaTraceAdapterError, +) +from aiperf.dataset.graph.models import ( + ExpectedTokens, + LlmNode, + ParsedGraph, + ProvenanceSpec, +) +from aiperf.dataset.graph.segment_ir.envelope import stamp_prompt_segment_ids +from aiperf.dataset.graph.segment_ir.pool import SegmentPool +from aiperf.dataset.graph.segment_ir.trie_content import ( + ReconCallbacks, + TrieNode, + TrieNodeBuild, + TrieRequest, + assemble_trie_graph, + build_trie_ir, +) + + +def build_trie_graph( + trace: WekaTrace, + *, + callbacks: ReconCallbacks | None = None, + tokenizer_name: str = "builtin", + prompt_corpus: str = "coding", + root_seed: int | None = None, + idle_gap_cap_seconds: float | None = None, + max_osl: int | None = None, +) -> tuple[ParsedGraph, SegmentPool]: + """Build a dependency-only :class:`ParsedGraph` + :class:`SegmentPool`. + + ``callbacks`` supplies the reconstructor's block-decode / partial-tail / + decode-to-text functions. When ``None`` (production), a + :class:`~aiperf.dataset.graph.adapters.shared.content.CorpusContentSynthesizer` + is built from ``(tokenizer_name, prompt_corpus, root_seed)`` and its + byte-faithful callbacks are used, decoding each hash block to the trace's + own ``block_size`` tokens. Unit tests pass deterministic stubs. + + ``idle_gap_cap_seconds`` caps inter-request-start idle gaps: every recorded + node start ``request.t`` is mapped through an + :class:`~aiperf.dataset.graph.segment_ir.trie_content.ActiveIdleWarp` + so node arrival offsets, edge ``delay_after_predecessor_us``, and root + ``min_start_delay_us`` all sit on the SAME warped clock. When ``None`` + (warp disabled) the raw ``request.t`` timeline is used unchanged. Without + the cap a recorded multi-hour idle gap survives into the warmup phase and + parks it forever; the cap compresses each over-long gap to ``cap_seconds``. + + ``max_osl`` (``--synthesis-max-osl``) caps each TOP-LEVEL chain request's + native ``LlmNode.max_tokens`` to ``min(recorded out, max_osl)`` + (agentx ``_cap_output`` parity); subagent-body requests are left uncapped. + ``None`` leaves the recorded ``out`` uncapped everywhere. The cap touches + dispatch only -- synthesized history content stays sized to the recorded + ``out`` so successor prompt bytes (and ISL) are unchanged. + + The returned :class:`ParsedGraph` carries the trace's nodes + edges on its + single top-level ``graph``; ``graphs`` stays empty (a single top-level + graph, no per-trace graph variants in this trivial IR). + + Raises: + EmptyWekaTraceError: the recorded ``requests`` flatten to zero + normal/streaming leaves (e.g. only subagent markers with empty + bodies) -- such a trace could never fire a node. + """ + if callbacks is None: + callbacks = _default_callbacks( + tokenizer_name, + prompt_corpus, + root_seed, + trace_id=trace.id, + block_size=trace.block_size, + hash_scope=trace.hash_id_scope, + ) + + nodes = _flatten_requests(trace.requests, root_scope=trace.id) + if not nodes: + raise EmptyWekaTraceError( + f"trace {trace.id!r}: requests flatten to zero normal/streaming leaf " + "requests (only subagent markers with empty bodies) -- the trace " + "could never fire a node" + ) + top_level_ids = _top_level_leaf_ids(trace.requests, trace.id) + pool = SegmentPool() + result = build_trie_ir( + nodes, + block_size=trace.block_size, + callbacks=callbacks, + pool=pool, + idle_gap_cap_seconds=idle_gap_cap_seconds, + # A recorded request with zero covered blocks (in < block_size, + # hash_ids []) still carried a real prompt: synthesize the sampled + # sub-block user message (dynamo parity) instead of the legacy empty + # prompt, which produced an unreplayable empty messages array. + small_prompt_fallback=True, + ) + + def build_node(node: TrieNode) -> LlmNode: + return _build_llm_node( + node, + build=result.builds[node.node_id], + max_osl=max_osl if node.node_id in top_level_ids else None, + ) + + graph = assemble_trie_graph( + nodes, + result, + build_node=build_node, + provenance=ProvenanceSpec(source="weka_trace", tool="aiperf-weka-trie/1"), + ) + return ParsedGraph(graph=graph), pool + + +# --- flattening ----------------------------------------------------------- + + +def _guard_scope(scope: str, seen: set[str]) -> None: + """Fail loud on scope ids the ``{scope}:{turn}`` grammar cannot carry. + + ``:`` is the ONE reserved character of the identity grammar (node ids are + ``{scope}:{turn}``; runtime ids append ``::{nonce}``), and a duplicate + scope (a repeated ``agent_id``, or an ``agent_id`` equal to the trace id) + would merge two trajectories' node namespaces. + """ + if not scope or ":" in scope: + raise WekaTraceAdapterError( + f"scope id {scope!r} cannot form node ids: scope ids must be " + "non-empty and contain no ':' (the reserved identity separator)" + ) + if scope in seen: + raise WekaTraceAdapterError( + f"duplicate trajectory scope {scope!r}: agent_ids must be unique " + "within a trace and distinct from the trace id" + ) + seen.add(scope) + + +def _flatten_requests(requests: list, *, root_scope: str) -> list[TrieNode]: + """Collect every recorded leaf request in recorded ``t`` order. + + Recurses into subagent markers (and nested markers). Node ids are + ``{scope}:{turn}`` -- the trajectory scope (the trace id for the root + chain, the recorded ``agent_id`` for each subagent) plus the 0-based turn + within that scope. Pure recorded data: identical logical traces produce + identical node ids regardless of file layout, and a node id IS the legacy + ``(conversation, turn_index)`` coordinate of its trajectory. + """ + out: list[TrieNode] = [] + seen: set[str] = set() + _guard_scope(root_scope, seen) + _walk( + requests, + out, + scope=root_scope, + scopes_seen=seen, + async_ancestors=frozenset(), + ) + return out + + +def _top_level_leaf_ids(requests: list, root_scope: str) -> frozenset[str]: + """Node ids of the leaves sitting DIRECTLY in the trace's top-level list. + + Mirrors :func:`_walk`'s id assignment for the root scope: the k-th + top-level leaf is ``{root_scope}:{k}``; anything under a subagent marker + lives in that agent's own scope. Consumed by the ``max_osl`` cap, which + applies to top-level chain requests only (agentx ``_cap_output`` parity). + """ + return frozenset( + f"{root_scope}:{k}" + for k, req in enumerate( + req for req in requests if not isinstance(req, WekaSubagentEntry) + ) + ) + + +_ASYNC_LAUNCHED_STATUS = "async_launched" + + +def _walk( + requests: list, + out: list[TrieNode], + *, + scope: str, + scopes_seen: set[str], + async_ancestors: frozenset[str], + inherited_causal: str | None = None, +) -> None: + """Depth-first walk appending leaves to ``out`` in recorded order. + + Recurses into subagent markers, threading each ``async_launched`` subtree's + id into ``async_ancestors`` so the interval-order edge builder can exclude + fire-and-forget children (:func:`~aiperf.dataset.graph.segment_ir.interval_order.build_interval_edges`). Dependency structure is + derived globally afterward from the recorded finished-before intervals and + ranks, not from the walk order (see :func:`~aiperf.dataset.graph.segment_ir.interval_order.build_interval_edges`). + + Each leaf's ``causal_parent_id`` is stamped here: the previous n/s leaf in + this same list (chain-prev), else the nearest preceding n/s leaf in an + enclosing list (``inherited_causal`` -- the spawner), else ``None``. A + subagent marker inherits the current ``prev_leaf_id`` (the leaf that spawned + it) as its subtree's ``inherited_causal``, falling back to this list's own + ``inherited_causal`` when no preceding leaf exists. This feeds + :func:`~aiperf.dataset.graph.segment_ir.interval_order.apply_start_anchors` + (wired inside ``build_trie_ir``): when the causal parent is still in flight + at a node's recorded start, its incoming edges collapse to one start-anchored + edge -- recorded mid-flight concurrency tracks the parent causally instead of + freezing to the recorded wall clock. + """ + prev_leaf_id: str | None = None + turn = 0 + for req in requests: + if isinstance(req, WekaSubagentEntry): + child_scope = req.agent_id + _guard_scope(child_scope, scopes_seen) + child_async = ( + async_ancestors | {child_scope} + if req.status == _ASYNC_LAUNCHED_STATUS + else async_ancestors + ) + _walk( + req.requests, + out, + scope=child_scope, + scopes_seen=scopes_seen, + async_ancestors=child_async, + inherited_causal=( + prev_leaf_id if prev_leaf_id is not None else inherited_causal + ), + ) + else: + node_id = f"{scope}:{turn}" + turn += 1 + out.append( + TrieNode( + node_id=node_id, + request=TrieRequest( + hash_ids=list(req.hash_ids), + input_length=req.input_length, + output_length=req.output_length, + t=req.t, + api_time=req.api_time or 0.0, + model=req.model, + streaming=req.type == "s", + ttft=req.ttft if req.type == "s" else None, + ), + order=len(out), + async_ancestors=async_ancestors, + causal_parent_id=( + prev_leaf_id if prev_leaf_id is not None else inherited_causal + ), + ) + ) + prev_leaf_id = node_id + + +# --- node construction ------------------------------------------------------ + + +def _build_llm_node( + node: TrieNode, + *, + build: TrieNodeBuild, + max_osl: int | None = None, +) -> LlmNode: + """Assemble ``node``'s prompt-segment path + assistant output into the LlmNode. + + The node carries NO inline prompt (``prompt=[]``): on the trie route the + prompt content lives ONLY in the run's ``SegmentPool``, addressed by + ``metadata["trie"]["prompt_segment_ids"]`` (stamped below). The store build, + ``graph_meta`` sidecar (``strip_replay_text`` forces ``prompt=[]``), and the + worker all materialize from the segment pool / mmap store -- none reads + ``node.prompt`` -- matching the dynamo and dag_jsonl adapters' convention. + For in-process debugging, recover the content with + ``segment_pool.materialize(read_prompt_segment_ids(node))``. + + The prompt path comes from + :func:`~aiperf.dataset.graph.segment_ir.trie_content.assemble_messages`, + which groups the node's frozen per-block ``(role, starts_new_message)`` tags + into messages and emits one content-addressed pool entry per message. Because + the tags are frozen per trie position, a shared block prefix produces + byte-identical message ids -- a real KV-cache prefix -- with NO relabeling. + The recorded ``out`` is emitted as one trailing ``"assistant"`` segment + parented at the prompt tip. + + ``max_osl`` (already resolved to this node: the caller passes ``None`` for + subagent-body nodes) caps the wire generation cap (the native + ``LlmNode.max_tokens``, endpoint-mapped by the + worker) to ``min(recorded out, max_osl)``; a recorded ``out`` of 0 + upgrades to 1 with a warning (:func:`wire_output_cap`). The response + SEGMENT stays sized to the recorded ``out`` so successor prompt content is + unchanged. + """ + req = node.request + max_tokens = ( + req.output_length if max_osl is None else min(req.output_length, max_osl) + ) + llm = LlmNode( + prompt=[], + output=f"{node.node_id}_out", + streaming=req.streaming, + model=req.model, + max_tokens=wire_output_cap(max_tokens, node_id=node.node_id), + arrival_offset_us=int(round(node.start * MICROS_PER_SECOND)), + # Recorded token expectations (dynamo parity): the RECORDED lengths, + # not the capped wire value. Cache fields stay None -- weka records no + # engine cache counts, and the loader's theoretical prefix-cache + # prediction is a model, not a recording. + expected=ExpectedTokens( + input_tokens=req.input_length, + output_tokens=req.output_length, + ), + ) + # ONLY prompt_segment_ids is stamped (dynamo parity): the build-synthesis + # response_id / hash_ids companions reached no build, sidecar, store, or + # dispatch consumer, and the per-node hash-list copy was pure build-plane + # RAM at corpus scale. + return stamp_prompt_segment_ids(llm, build.prompt_path) + + +# --- default production callbacks ----------------------------------------- + + +def _default_callbacks( + tokenizer_name: str, + prompt_corpus: str, + root_seed: int | None, + trace_id: str | None = None, + *, + block_size: int = 64, + hash_scope: str = "local", +) -> ReconCallbacks: + """Build byte-faithful reconstructor callbacks from a real synthesizer. + + Imported lazily so the heavy corpus build stays off the import path for + callers (and tests) that inject their own stub callbacks. + + ``hash_scope`` is the trace's declared ``hash_id_scope`` and picks the + block-decode namespace. Under ``"local"``, ``trace_id`` scopes the hash-id + namespace per trace: block decode reseeds per ``(trace_id, hash_id)``, so + equal hash ids across traces synthesize DIFFERENT bytes (no manufactured + cross-trace KV-cache sharing). Under ``"global"``, block decode reseeds per + bare ``hash_id`` -- the SAME namespace dynamo's content-global recorded + hashes use -- so equal hash ids across traces synthesize byte-identical + blocks, reproducing recorded cross-trace KV-cache sharing. Sharing needs no + cross-file coordination: the reseed is a pure function of ``(root seed, + hash_id)``, so independently-parsed traces (including parallel pool + workers) decode identical bytes for equal ids. + + Partial-tail seeds -- which the shared trie driver keys only by node id + (``"{node_id}:response"`` / ``"{node_id}:tiny"``) -- get the trace id + prefixed under BOTH scopes: trajectory-local node ids recur across traces and + are a driver artifact, not part of the recorded hash namespace, so leaving + them unscoped would fabricate identical response bytes corpus-wide. + + ``block_size`` is the TRACE's recorded block size: each hash id decodes to + exactly that many corpus tokens (mirrors ``dynamo_recon_callbacks``). The + synthesizer's own default is weka-legacy 64; threading the trace value keeps + a ``block_size != 64`` trace's covered-count ISL intact. Decodes always use + a PRIVATE per-build cache: the synthesizer instance (and its shared cache) + is reused across traces and block sizes, and mixing block sizes in the + shared bare-hash-id cache would hand wrong-sized blocks across builds. + """ + from aiperf.dataset.graph.adapters.shared.content import ( + get_or_build_synthesizer, + ) + + synth = get_or_build_synthesizer( + tokenizer_name, prompt_corpus=prompt_corpus, root_seed=root_seed + ) + trace_cache: dict[int, list[int]] = {} + + decode_trace_id = None if hash_scope == "global" else trace_id + + def _decode_scoped(hash_ids: list[int]) -> list[int]: + return synth._decode_block_tokens( + hash_ids, block_size=block_size, cache=trace_cache, trace_id=decode_trace_id + ) + + def _tail_scoped(n_tokens: int, seed: str) -> list[int]: + return synth._sample_partial_tail_tokens(n_tokens, f"{trace_id}:{seed}") + + return ReconCallbacks( + decode_block_tokens=_decode_scoped, + sample_partial_tail_tokens=_tail_scoped, + decode_tokens_to_text=synth._decode_tokens_to_text, + ) + + +__all__ = [ + "ReconCallbacks", + "build_trie_graph", +] diff --git a/src/aiperf/dataset/graph/auto_derive.py b/src/aiperf/dataset/graph/auto_derive.py new file mode 100644 index 0000000000..63d6b91cbc --- /dev/null +++ b/src/aiperf/dataset/graph/auto_derive.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Auto-derivation: linear-chat synthesis, channel auto-derive. + +Pure transform on ParsedGraph. +""" + +from __future__ import annotations + +from typing import Any + +import msgspec + +from aiperf.dataset.graph.decode import decode_node +from aiperf.dataset.graph.models import ( + END_NODE_ID, + START_NODE_ID, + ChannelSpec, + ChannelType, + GraphRecord, + LlmNode, + ParsedGraph, + ReducerName, + StaticEdge, + TraceRecord, + _collect_content_block_channels, + _collect_prompt_array_channels, +) + +_LINEAR_CHAT_NODE_ID = "_llm" +_MESSAGES_CHANNEL = "messages" + + +def auto_derive(parsed: ParsedGraph) -> ParsedGraph: + """Apply all auto-derivation transforms in order. Returns a new ParsedGraph.""" + pb = parsed + pb = _lift_trace_messages(pb) + pb = _synthesize_linear_chat_if_needed(pb) + pb = _auto_derive_channels(pb) + return normalize_runtime_channels(pb) + + +def _lift_trace_messages(pb: ParsedGraph) -> ParsedGraph: + """Lift the ``traces[].messages`` shorthand into ``initial_state.messages``. + + The shorthand is documented as equivalent to ``initial_state.messages`` + (see ``TraceRecord.messages``), so it applies regardless of whether the + graph declares explicit nodes — not only on the linear-chat synthesis + path. An already-present ``initial_state.messages`` wins. + """ + if not any(t.messages is not None for t in pb.traces): + return pb + return msgspec.structs.replace( + pb, traces=[_lift_messages_into_init(t) for t in pb.traces] + ) + + +def _synthesize_linear_chat_if_needed(pb: ParsedGraph) -> ParsedGraph: + if pb.graph.nodes: + return pb + system_prompt = pb.graph.system + prompt: list[Any] = [] + if system_prompt is not None: + prompt.append({"role": "system", "content": system_prompt}) + prompt.append("@" + _MESSAGES_CHANNEL) + llm_node = decode_node( + { + "prompt": prompt, + "output": _MESSAGES_CHANNEL, + "streaming": True, + } + ) + new_graph = msgspec.structs.replace( + pb.graph, + nodes={_LINEAR_CHAT_NODE_ID: llm_node}, + state={ + **pb.graph.state, + _MESSAGES_CHANNEL: ChannelSpec( + type=ChannelType.MESSAGES, reducer=ReducerName.ADD_MESSAGES + ), + }, + edges=list(pb.graph.edges) + + [ + StaticEdge(source=START_NODE_ID, target=_LINEAR_CHAT_NODE_ID), + StaticEdge(source=_LINEAR_CHAT_NODE_ID, target=END_NODE_ID), + ], + ) + return msgspec.structs.replace(pb, graph=new_graph) + + +def _lift_messages_into_init(t: TraceRecord) -> TraceRecord: + if t.messages is None: + return t + new_init = dict(t.initial_state) + new_init.setdefault("messages", t.messages) + return msgspec.structs.replace(t, initial_state=new_init, messages=None) + + +def _auto_derive_channels(pb: ParsedGraph) -> ParsedGraph: + declared = dict(pb.graph.state) + for node in pb.graph.nodes.values(): + if not isinstance(node, LlmNode): + continue + for ch in _collect_prompt_array_channels(node): + if ch not in declared: + declared[ch] = ChannelSpec( + type=ChannelType.MESSAGES, reducer=ReducerName.ADD_MESSAGES + ) + for ch in _collect_content_block_channels(node): + if ch not in declared: + declared[ch] = ChannelSpec( + type=ChannelType.TEXT, reducer=ReducerName.OVERWRITE + ) + if declared == pb.graph.state: + return pb + return msgspec.structs.replace( + pb, graph=msgspec.structs.replace(pb.graph, state=declared) + ) + + +def normalize_runtime_channels(parsed: ParsedGraph) -> ParsedGraph: + """Return ``parsed`` with safe inferred channel declarations added.""" + # Collect the trace-driven channels that must be declared on the top-level + # graph; subgraphs only gain their own node-output channels. + trace_channels: list[str] = [] + for trace in parsed.traces: + trace_channels.extend(trace.initial_state) + for outputs in trace.replay_outputs.values(): + trace_channels.extend(outputs) + + new_graph = _normalize_graph(parsed.graph, extra_channels=trace_channels) + if new_graph is parsed.graph: + return parsed + return msgspec.structs.replace(parsed, graph=new_graph) + + +def _normalize_graph(graph: GraphRecord, *, extra_channels: Any) -> GraphRecord: + new_state = dict(graph.state) + for node in graph.nodes.values(): + for channel in _node_output_channels(node): + _declare_missing(new_state, channel) + for channel in extra_channels: + _declare_missing(new_state, channel) + if new_state == graph.state: + return graph + return msgspec.structs.replace(graph, state=new_state) + + +def _node_output_channels(node: Any) -> list[str]: + write_channels = getattr(node, "write_channels", None) + if isinstance(write_channels, list): + return [channel for channel in write_channels if isinstance(channel, str)] + return [] + + +def _declare_missing(state: dict[str, ChannelSpec], channel: str) -> None: + if channel in state: + return + state[channel] = _default_channel_spec(channel) + + +def _default_channel_spec(channel: str) -> ChannelSpec: + if channel == "messages": + return ChannelSpec(type=ChannelType.MESSAGES, reducer=ReducerName.ADD_MESSAGES) + return ChannelSpec(type=ChannelType.TEXT, reducer=ReducerName.OVERWRITE) diff --git a/src/aiperf/dataset/graph/codecs.py b/src/aiperf/dataset/graph/codecs.py new file mode 100644 index 0000000000..6091f211c3 --- /dev/null +++ b/src/aiperf/dataset/graph/codecs.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Cached msgspec codecs for the graph IR (canonical struct (de)serialization). + +Decoders snapshot their target type at first use, so they are module-level and +built once. msgpack is the cross-process store format. +""" + +from __future__ import annotations + +from typing import Any + +import msgspec + +from aiperf.dataset.graph.models import ParsedGraph + +_PG_MSGPACK_ENCODER = msgspec.msgpack.Encoder() +_PG_MSGPACK_DECODER = msgspec.msgpack.Decoder(ParsedGraph) + + +def encode_parsed_graph_msgpack(pg: ParsedGraph) -> bytes: + """Encode a :class:`ParsedGraph` to canonical msgpack bytes (cross-process).""" + return _PG_MSGPACK_ENCODER.encode(pg) + + +def decode_parsed_graph_msgpack(data: bytes) -> ParsedGraph: + """Decode canonical msgpack bytes back into a :class:`ParsedGraph`.""" + return _PG_MSGPACK_DECODER.decode(data) + + +GRAPH_META_SIDECAR_FILENAME = "graph_meta.msgpack" + +# Sidecar frame ``kind`` discriminator: every frame carries a node-typed +# ``ParsedGraph`` and MUST declare ``kind`` explicitly (see the strict decode). +_SIDECAR_KIND_PARSED_GRAPH = "parsed_graph" + +# Bumped 2->3 when the kind-less decode default was retired: the decoder now +# REQUIRES an explicit ``kind`` header, so old persisted sidecars/cache entries +# (written without it) version-gate out and self-heal on the next run. +# Bumped 3->4 with the verbatim raw-JSON segment variant (``Segment.wire_json``): +# pre-v4 stores normalized every segment to ``{"role", "content"}`` (dropping key +# order and extra keys). Unlike the 2->3 bump, v4 adds NO reader-side gate: the +# decoder tolerates a pre-v4 blob because ``Segment.wire_json`` defaults to ``None`` +# (a normalized role/content segment), so the version here is advisory provenance, +# not a compatibility gate. msgspec auto-encodes the new dataclass field, so the +# cross-process ``ParsedGraph`` codec round-trips it without an explicit hook. +GRAPH_META_SCHEMA_VERSION: int = 4 + +_SIDECAR_DECODER = msgspec.msgpack.Decoder() + + +def encode_graph_meta_sidecar( + pg: ParsedGraph, + *, + source_fingerprint: dict[str, Any], + schema_version: int = GRAPH_META_SCHEMA_VERSION, +) -> bytes: + """Encode a content-free structural ``ParsedGraph`` + identity header to bytes. + + The frame is ``[header, pg_bytes]``; ``pg_bytes`` is the canonical + ``encode_parsed_graph_msgpack`` output kept as a nested blob so the existing + typed decoder round-trips the graph unchanged. The header declares an explicit + ``kind`` so the strict decoder can reject kind-less pre-v3 artifacts. + """ + header = { + "kind": _SIDECAR_KIND_PARSED_GRAPH, + "schema_version": schema_version, + "source_fingerprint": source_fingerprint, + } + return _PG_MSGPACK_ENCODER.encode([header, encode_parsed_graph_msgpack(pg)]) + + +def decode_graph_meta_sidecar( + data: bytes, +) -> tuple[ParsedGraph, dict[str, Any], int]: + """Decode a sidecar frame into ``(parsed_graph, source_fingerprint, schema_version)``. + + The blob is a content-free node-typed :class:`ParsedGraph`. The header ``kind`` + is REQUIRED and must equal ``"parsed_graph"``: a kind-less frame (a pre-v3 + persisted artifact) or an unsupported kind raises ``ValueError``; the + TimingManager surfaces it as ``InvalidStateError`` (the sidecar is mandatory; + no re-parse fallback exists). + + Raises ``msgspec.DecodeError`` / ``ValueError`` on a malformed frame. + """ + frame = _SIDECAR_DECODER.decode(data) + if not isinstance(frame, list) or len(frame) != 2: + raise ValueError("graph_meta sidecar frame must be [header, blob]") + header, blob = frame + if not isinstance(header, dict) or not isinstance(blob, bytes): + raise ValueError("graph_meta sidecar frame has the wrong shape") + if "source_fingerprint" not in header or "schema_version" not in header: + raise ValueError("graph_meta sidecar header missing required keys") + if header.get("kind") != _SIDECAR_KIND_PARSED_GRAPH: + raise ValueError( + "unsupported or missing graph_meta sidecar kind; rebuild the graph store" + ) + return ( + decode_parsed_graph_msgpack(blob), + header["source_fingerprint"], + int(header["schema_version"]), + ) diff --git a/src/aiperf/dataset/graph/decode.py b/src/aiperf/dataset/graph/decode.py new file mode 100644 index 0000000000..5e2a6b1174 --- /dev/null +++ b/src/aiperf/dataset/graph/decode.py @@ -0,0 +1,184 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Single home for loose/legacy graph input -> typed IR coercion. + +Coerces hand-authored / foreign dicts into the typed IR: + - LlmNode dispatch + kind->node_type rename + messages->prompt alias (a node + setting BOTH ``prompt`` and ``messages`` is rejected; so is a prompt-less + node that only carries ``expected.input_tokens`` — synth-token prompt + fabrication is not supported on the native lowering path) + - static-edge coercion; non-finite delay values are rejected on edges and + nodes (NaN/Inf discipline: they cross the msgpack boundary and would gate + the successor forever at runtime) + - vendor-key folding into the ``extra`` catch-all for ProvenanceSpec (that + model does NOT forbid unknown fields, so msgspec would silently drop the + vendor keys otherwise) + +Trusted writers and the store bypass this with ``codecs.*`` directly. +""" + +from __future__ import annotations + +import math +from typing import Any + +import msgspec + +from aiperf.dataset.graph.models import ( + GraphRecord, + LlmNode, + NodeKind, + ProvenanceSpec, + StaticEdge, +) + + +class GraphDecodeError(ValueError): + """Raised when loose graph input cannot be coerced to typed IR.""" + + +_VALID_KINDS = {k.value for k in NodeKind} + +# Known (declared) field names of the vendor-key catch-all model. Any +# key NOT in the set (and not "extra") is folded into the model's `extra` dict. +_PROVENANCE_FIELDS = set(ProvenanceSpec.__struct_fields__) - {"extra"} + + +def _fold_extra(raw: Any, known: set[str], *, loc: str) -> Any: + """Fold unknown keys of a catch-all model's dict into its ``extra`` field.""" + if not isinstance(raw, dict): + return raw + extra_raw = raw.get("extra") + if extra_raw is not None and not isinstance(extra_raw, dict): + raise GraphDecodeError( + f"{loc}.extra: must be a mapping of vendor keys, got " + f"{type(extra_raw).__name__}" + ) + out: dict[str, Any] = {} + extra: dict[str, Any] = dict(extra_raw or {}) + for k, v in raw.items(): + if k == "extra": + continue + if k in known: + out[k] = v + else: + extra[k] = v + if extra: + out["extra"] = extra + return out + + +def _has_expected_input_tokens(raw: dict[str, Any]) -> bool: + exp = raw.get("expected") + return isinstance(exp, dict) and isinstance(exp.get("input_tokens"), int) + + +def _normalize_llm(raw: dict[str, Any], loc: str) -> dict[str, Any]: + out = dict(raw) + if "dispatch_overrides" in out: + raise GraphDecodeError( + f"{loc}: 'dispatch_overrides' was renamed 'extra_body' " + f"(Turn-naming standardization); model / streaming / max_tokens / " + f"raw_tools are first-class node fields" + ) + if "prompt" in out and "messages" in out: + raise GraphDecodeError( + f"{loc}: node sets both 'prompt' and 'messages'; 'messages' is an " + f"alias for 'prompt' — keep exactly one" + ) + if "prompt" not in out and "messages" in out: + out["prompt"] = out.pop("messages") + if "prompt" not in out and _has_expected_input_tokens(out): + raise GraphDecodeError( + f"{loc}: node has no prompt; synth-token fabrication is not " + f"supported on the native lowering path" + ) + return out + + +def _require_finite(loc: str, field: str, value: float | None) -> None: + if value is not None and not math.isfinite(value): + raise GraphDecodeError(f"{loc}: {field} must be finite, got {value!r}") + + +def decode_node(raw: Any, node_id: str | None = None) -> Any: + """Coerce one loose node spec (dict or already-typed struct) into typed IR. + + ``node_id`` (when the caller knows it, e.g. :func:`decode_graph`) is used + only for ``graph.nodes.``-style error locations. + """ + if isinstance(raw, LlmNode): + return raw + loc = f"graph.nodes.{node_id}" if node_id is not None else "node" + if not isinstance(raw, dict): + raise GraphDecodeError(f"Node spec must be a mapping, got {type(raw).__name__}") + raw_kind = raw.get("kind") + explicit = raw.get("node_type") or raw_kind + body = {k: v for k, v in raw.items() if k != "kind"} + + if explicit is not None and explicit not in _VALID_KINDS: + field = "kind" if raw_kind is not None else "node_type" + raise GraphDecodeError( + f"unknown node {field} {explicit!r}; expected one of {sorted(_VALID_KINDS)}" + ) + + body = _normalize_llm(body, loc) + body.pop("node_type", None) # tag is implied by the chosen class + try: + node = msgspec.convert(body, type=LlmNode) + except msgspec.ValidationError as exc: + raise GraphDecodeError( + f"{loc}: node decode failed for kind {explicit!r}: {exc}" + ) from exc + _require_finite(loc, "min_start_delay_us", node.min_start_delay_us) + return node + + +def decode_edge(raw: Any) -> StaticEdge: + """Coerce one loose edge spec (dict or already-typed struct) into typed IR.""" + if isinstance(raw, StaticEdge): + return raw + if not isinstance(raw, dict): + raise GraphDecodeError(f"Edge spec must be a mapping, got {type(raw).__name__}") + body = {k: v for k, v in raw.items() if k != "edge_type"} + try: + edge = msgspec.convert(body, type=StaticEdge) + except msgspec.ValidationError as exc: + raise GraphDecodeError(f"edge decode failed: {exc}") from exc + loc = f"graph.edges[{edge.source}->{edge.target}]" + _require_finite(loc, "delay_after_predecessor_us", edge.delay_after_predecessor_us) + _require_finite(loc, "min_start_delay_us", edge.min_start_delay_us) + _require_finite( + loc, + "delay_after_predecessor_start_us", + edge.delay_after_predecessor_start_us, + ) + _require_finite( + loc, + "delay_after_predecessor_first_token_us", + edge.delay_after_predecessor_first_token_us, + ) + return edge + + +def decode_graph(raw: Any) -> GraphRecord: + """Coerce a loose graph dict (or :class:`GraphRecord`) into typed IR.""" + if isinstance(raw, GraphRecord): + return raw + if not isinstance(raw, dict): + raise GraphDecodeError( + f"Graph spec must be a mapping, got {type(raw).__name__}" + ) + body = dict(raw) + if isinstance(body.get("nodes"), dict): + body["nodes"] = {nid: decode_node(n, nid) for nid, n in body["nodes"].items()} + if isinstance(body.get("edges"), list): + body["edges"] = [decode_edge(e) for e in body["edges"]] + if isinstance(body.get("provenance"), dict): + body["provenance"] = _fold_extra( + body["provenance"], _PROVENANCE_FIELDS, loc="graph.provenance" + ) + try: + return msgspec.convert(body, type=GraphRecord, from_attributes=False) + except msgspec.ValidationError as exc: + raise GraphDecodeError(f"graph decode failed: {exc}") from exc diff --git a/src/aiperf/dataset/graph/graph_meta_sidecar.py b/src/aiperf/dataset/graph/graph_meta_sidecar.py new file mode 100644 index 0000000000..07617303bf --- /dev/null +++ b/src/aiperf/dataset/graph/graph_meta_sidecar.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Write/locate the content-free structural ``ParsedGraph`` sidecar. + +The DatasetManager serializes the structural graph into the per-benchmark +sidecar directory (``aiperf_graph_meta_/``; the unified segment store +lives in its own ``aiperf_graph_segments_/``) on EVERY graph build route, +and advertises the written path on the graph-typed +``DatasetConfiguredNotification.client_metadata``. The sidecar is mandatory: +the TimingManager only ingests this artifact, from the broadcast path (it +never re-parses the workload and never re-derives the path from env +conventions), so a build that cannot land a catalog-consistent sidecar fails +the run. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import msgspec + +from aiperf.dataset.graph.codecs import ( + GRAPH_META_SIDECAR_FILENAME, + encode_graph_meta_sidecar, +) +from aiperf.dataset.graph.graph_path_catalog import build_graph_path_catalog +from aiperf.dataset.graph.models import ( + GraphRecord, + LlmNode, + ParsedGraph, + TraceRecord, +) + +GRAPH_META_DIR_PREFIX = "aiperf_graph_meta_" + +__all__ = [ + "catalogs_match", + "sidecar_matches_index", + "sidecar_path_for", + "strip_replay_text", + "write_graph_meta_sidecar", +] + + +def _strip_graph_node_content(graph: GraphRecord) -> GraphRecord: + """Return a copy of ``graph`` with every ``LlmNode``'s bulk content emptied. + + Two dominant, schedule-plane-irrelevant fields are cleared: + + * ``LlmNode.prompt`` -- the resolved prompt content, inline megabytes/trace; + * ``metadata["trie"]`` contents -> ``{}`` -- the ``prompt_segment_ids`` and + any dynamic-slot ``assembly``/``capture`` keys (the worker reads all of + those from the STORE, not the sidecar). The ``"trie"`` KEY is KEPT (empty) + so ``graph_ir_replay._is_trie_graph`` still routes the trie frontier-chop. + + All other metadata, the native dispatch fields (model / max_tokens / + raw_tools / extra_headers / theoretical prefix-cache counts), + node ids/types, edges, and ``arrival_offset_us`` are + preserved verbatim -- everything the schedule plane and the prefix-cache metric + actually consume. + """ + + def _strip(node): + if not isinstance(node, LlmNode): + return node + metadata = node.metadata or {} + new_metadata = {**metadata, "trie": {}} if "trie" in metadata else metadata + return msgspec.structs.replace(node, prompt=[], metadata=new_metadata) + + return msgspec.structs.replace( + graph, nodes={nid: _strip(n) for nid, n in graph.nodes.items()} + ) + + +def strip_replay_text(graph: ParsedGraph) -> ParsedGraph: + """Return a content-free copy for the graph_meta sidecar. + + Every trace's ``replay_outputs`` is cleared. For the trie IR + (``segment_pool is not None``) the pool content is emptied (the pool is kept + NON-None so the loaded graph still reads as a segment-store IR -- e.g. + ``TraceExecutor``'s dispatch-failure sentinel writes gate on + ``segment_pool is not None``) and each + ``LlmNode``'s bulk content is stripped via :func:`_strip_graph_node_content` + across the top graph and every per-trace graph: the + inline ``prompt`` and the ``metadata["trie"]`` contents (``prompt_segment_ids`` + and any ``assembly``/``capture`` slot keys) -> ``{}`` (the ``"trie"`` marker + key stays). Non-trie graphs carry + no pool (``None``) and keep their small channel-ref prompts + metadata + untouched. The native dispatch fields (model / max_tokens / raw_tools / + extra_headers / theoretical prefix-cache counts), node ids/types, edges, + and ``arrival_offset_us`` are preserved verbatim -- everything the TimingManager + and the prefix-cache metric actually consume; the worker reads prompt content + + ``prompt_segment_ids`` from the STORE, never the sidecar. + """ + from aiperf.dataset.graph.segment_ir.pool import SegmentPool + + stripped: list[TraceRecord] = [ + msgspec.structs.replace(t, replay_outputs={}) for t in graph.traces + ] + if graph.segment_pool is None: + return msgspec.structs.replace(graph, traces=stripped) + + return msgspec.structs.replace( + graph, + traces=stripped, + segment_pool=SegmentPool(), + graph=_strip_graph_node_content(graph.graph), + graphs={k: _strip_graph_node_content(g) for k, g in graph.graphs.items()}, + ) + + +def sidecar_path_for(base_path: Path, benchmark_id: str) -> Path: + """Canonical sidecar path inside the per-benchmark sidecar dir.""" + return ( + Path(base_path) + / f"{GRAPH_META_DIR_PREFIX}{benchmark_id}" + / GRAPH_META_SIDECAR_FILENAME + ) + + +def write_graph_meta_sidecar( + graph: ParsedGraph, + *, + base_path: Path, + benchmark_id: str, + source_fingerprint: dict[str, Any], + schema_version: int, +) -> Path: + """Encode and write the structural sidecar; return the written path.""" + out = sidecar_path_for(base_path, benchmark_id) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_bytes( + encode_graph_meta_sidecar( + strip_replay_text(graph), + source_fingerprint=source_fingerprint, + schema_version=schema_version, + ) + ) + return out + + +def catalogs_match( + graph: ParsedGraph, other_catalog: dict[str, dict[str, int]] +) -> bool: + """True iff ``graph``'s build catalog equals ``other_catalog`` exactly.""" + return build_graph_path_catalog(graph) == other_catalog + + +def sidecar_matches_index( + graph: ParsedGraph, index_offsets: dict[str, dict[tuple[int, str], Any]] +) -> bool: + """Best-effort load-time check: per-trace sidecar ordinals == store index ordinals. + + The store index inner keys are ``(node_ordinal, phase_variant)`` integer tuples; + compare the per-trace ordinal SETS. A mismatch means the sidecar's topology + diverged from the stored envelopes -> the caller + (``TimingManager._load_graph_sidecar``) raises ``InvalidStateError`` (the + sidecar is mandatory; no re-parse fallback exists). + """ + catalog = build_graph_path_catalog(graph) + for trace_id, node_ordinals in catalog.items(): + store_ordinals = {ord_ for (ord_, _variant) in index_offsets.get(trace_id, {})} + if set(node_ordinals.values()) - store_ordinals: + return False + return True diff --git a/src/aiperf/dataset/graph/graph_path_catalog.py b/src/aiperf/dataset/graph/graph_path_catalog.py new file mode 100644 index 0000000000..afa1943fab --- /dev/null +++ b/src/aiperf/dataset/graph/graph_path_catalog.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Stable node-id -> node-ordinal catalog for the graph store. + +Every live producer lowers to a flat graph of ``LlmNode``s, so the catalog key +is simply the bare node id. Ordinals are assigned by +``flat_trie_ordinals`` (the SAME helper the build plane uses), so +build-time manifest ordinals and dispatch-time catalog ordinals are +byte-identical by construction and stable across runs. +""" + +from __future__ import annotations + +import dataclasses + +from aiperf.dataset.graph.models import ( + ParsedGraph, + TraceRecord, +) + + +@dataclasses.dataclass(frozen=True, slots=True) +class CatalogContext: + """Immutable per-build catalog state threaded into :func:`node_ordinal_for`. + + Built once per parse via :func:`build_catalog_context`; passed explicitly + to the resolver so there is NO module-level mutable cache (which would + corrupt across concurrent traces). + """ + + # ``{trace_id: {node_id: node_ordinal}}`` -- the addressing map. + catalog: dict[str, dict[str, int]] + + +def build_catalog_context(parsed: ParsedGraph) -> CatalogContext: + """Build the :class:`CatalogContext` (the per-trace ordinal catalog).""" + return CatalogContext(catalog=build_graph_path_catalog(parsed)) + + +def build_graph_path_catalog(parsed: ParsedGraph) -> dict[str, dict[str, int]]: + """Build ``{trace_id: {node_id: node_ordinal}}`` for every trace.""" + return {trace.id: _catalog_for_trace(parsed, trace) for trace in parsed.traces} + + +def _catalog_for_trace(parsed: ParsedGraph, trace: TraceRecord) -> dict[str, int]: + """Assign a dense ordinal to each LlmNode id in the shared trie order. + + The build-plane unified store is keyed by + :func:`flat_trie_ordinals` -- the SHARED build/schedule ordinal + scheme. The dispatch adapter must resolve a fired node to the EXACT ordinal + the store was written at, so this returns the same helper's output + verbatim -- byte-identical to the build plane by construction, never a + re-derived ordering. + """ + # Imported locally to avoid a build<->schedule module import cycle. + from aiperf.dataset.graph.segment_ir.store_builder import ( + flat_trie_ordinals, + ) + + return flat_trie_ordinals(parsed, trace) + + +def node_ordinal_for( + context: CatalogContext, + trace_id: str, + node_key: str, +) -> int | None: + """Resolve a fired node's build-time ordinal, or ``None`` when unknown.""" + return context.catalog.get(trace_id, {}).get(node_key) + + +__all__ = [ + "CatalogContext", + "build_catalog_context", + "build_graph_path_catalog", + "node_ordinal_for", +] diff --git a/src/aiperf/dataset/graph/merge.py b/src/aiperf/dataset/graph/merge.py new file mode 100644 index 0000000000..b823158c25 --- /dev/null +++ b/src/aiperf/dataset/graph/merge.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Merge per-source ParsedGraphs into one multi-graph workload. + +Used by the weka per-item parse (many files/rows -> one workload), the dynamo +per-session-tree merge, and the weka streaming store build's structural-sidecar +merge. +""" + +from __future__ import annotations + +from collections.abc import Iterable + +import msgspec + +from aiperf.dataset.graph.models import GraphRecord, ParsedGraph, resolve_trace_graph +from aiperf.dataset.graph.segment_ir.pool import Segment, SegmentPool + + +class GraphMergeError(ValueError): + """Raised when per-source ParsedGraphs cannot be merged into one workload.""" + + +def merge_parsed_graphs(per_source: Iterable[ParsedGraph]) -> ParsedGraph: + """Merge per-source ``ParsedGraph`` outputs into one MULTI-GRAPH workload. + + Trace sources are HETEROGENEOUS — every trace is an independent + conversation whose topology (turn count, subagent fan-out) differs, so + each trace keeps its OWN ``GraphRecord`` under ``ParsedGraph.graphs`` + keyed by its trace id, and that trace's ``graph_ref`` selects it. A + single source may carry one trace (weka rows, dynamo captures) or many + (a multi-session dag_jsonl parse); either way each trace's graph is + resolved from ITS source via ``resolve_trace_graph``, which also threads + the per-trace graph through firing-plan / snapshot / conversation-source + computation downstream. + + This function intentionally accepts any iterable and folds each + ``ParsedGraph`` into the merged structures immediately. The HF path feeds it + a pool-result generator, so the parent never holds the complete list of + decoded per-trace graphs on top of the merged graph. + """ + seen_trace_ids: set[str] = set() + merged_traces = [] + graphs: dict[str, GraphRecord] = {} + merged_graph: GraphRecord | None = None + # Segment-trie IR: union every per-source pool's content-addressed + # entries. Ids are content-addressed (blake2b over parent_id/role/tokens), + # so the SAME id MUST carry identical content -- dedup is only correct under + # that invariant. Stays None when no per-source graph carries a pool so the + # merged graph carries none either. + merged_segments: dict[str, Segment] | None = None + + for pg in per_source: + if pg.segment_pool is not None: + if merged_segments is None: + merged_segments = {} + # Fail loud on an id that maps to divergent content: two Segments + # sharing an id but differing in value (Segment is a frozen + # dataclass, so == is a value comparison) can only mean a + # content-addressing / hash break, which would otherwise silently + # keep whichever entry wins the union. This is a correctness + # invariant, not a debug aid, so it is a real raise (not an assert). + for sid, seg in pg.segment_pool.by_id.items(): + prior = merged_segments.get(sid) + if prior is not None and prior != seg: + raise GraphMergeError( + f"segment id {sid!r} maps to divergent content across " + f"graph trace sources (content-addressing invariant " + f"broken): {prior!r} != {seg!r}" + ) + merged_segments[sid] = seg + if merged_graph is None: + # Frozen structs are immutable and never mutated below, so the + # default single-graph slot can share the first source's graph directly. + merged_graph = pg.graph + + # Key each trace's OWN topology by its trace id. A source is usually + # one TraceRecord over one top-level graph (weka rows, dynamo + # captures), but a source may itself be multi-trace (a dag_jsonl file + # with several independent root sessions streams ONE structural blob + # carrying every trace); resolving through the trace's graph_ref keeps + # each trace on its own graph instead of remapping all of them onto + # the first tree's ``pg.graph``. + for trace in pg.traces: + if trace.id in seen_trace_ids: + raise GraphMergeError( + f"duplicate trace id across graph trace sources: {trace.id!r}" + ) + seen_trace_ids.add(trace.id) + graphs[trace.id] = resolve_trace_graph(pg, trace) + merged_traces.append(msgspec.structs.replace(trace, graph_ref=trace.id)) + + if merged_graph is None or not merged_traces: + raise GraphMergeError("graph merge produced zero ParsedGraph results") + + merged_traces.sort(key=lambda t: t.id) + + merged_pool = ( + SegmentPool(_by_id=merged_segments) if merged_segments is not None else None + ) + + return ParsedGraph( + graph=merged_graph, + graphs=graphs, + traces=merged_traces, + segment_pool=merged_pool, + ) + + +__all__ = [ + "GraphMergeError", + "merge_parsed_graphs", +] diff --git a/src/aiperf/dataset/graph/models.py b/src/aiperf/dataset/graph/models.py new file mode 100644 index 0000000000..4b108e3afc --- /dev/null +++ b/src/aiperf/dataset/graph/models.py @@ -0,0 +1,383 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Graph schema models - canonical frozen msgspec structs for graph/trace records. + +See `docs/benchmark-modes/agentic.md` for the canonical spec. + +These are the canonical IR types. Trusted (typed) construction goes through the +struct constructors / `codecs.py`; loose / foreign dict input is +normalized to these types by `decode.py` (`decode_graph` / `decode_node` / +`decode_edge`). Every producer (the weka and dynamo trace adapters, and +the native lowering) emits flat `LlmNode` + `StaticEdge` graphs. +""" + +from __future__ import annotations + +from enum import Enum +from typing import Annotated, Any, Literal + +import msgspec + +from aiperf.dataset.graph.segment_ir.pool import SegmentPool + +# Reserved graph sentinels: ``START`` is the virtual entry (a valid edge source +# only), ``END`` the virtual exit (a valid edge target only); neither is ever a +# real declared node id. +START_NODE_ID = "START" +END_NODE_ID = "END" + + +class ChannelType(str, Enum): + TEXT = "text" + MESSAGES = "messages" + + +class ReducerName(str, Enum): + OVERWRITE = "overwrite" + ADD_MESSAGES = "add_messages" + + +class ChannelSpec(msgspec.Struct, frozen=True, kw_only=True, omit_defaults=True): + """A state-channel declaration.""" + + type: ChannelType = ChannelType.TEXT # Channel value type. + # How concurrent writes to the channel are merged. + reducer: ReducerName = ReducerName.OVERWRITE + + +class ChannelRequirement( + msgspec.Struct, + frozen=True, + kw_only=True, + forbid_unknown_fields=True, +): + """AND-fan-in input requirement on a node. + + A node fires only when every requirement in its `inputs` list is satisfied. + `count: int` requires exactly that many writes to the channel; `count: "all"` + means "all statically declared producers" (computed at Scheduler init). + + Not ``omit_defaults`` so ``count`` always round-trips in the serialized + requirement shape (``{"channel": str, "count": int|"all"}``). + """ + + channel: str # Input channel name. + # Required arrival count. 'all' resolves to the static topology count of + # declared producers. + count: Annotated[int, msgspec.Meta(ge=1)] | Literal["all"] = 1 + + +class StaticEdge( + msgspec.Struct, + frozen=True, + kw_only=True, + omit_defaults=True, + tag_field="edge_type", + tag="static", + forbid_unknown_fields=True, +): + """Unconditional edge: `{source: , target: }`.""" + + source: str # Source node id (or 'START'). + target: str # Destination node id (or 'END'). + # Idle / scheduling delay after the predecessor finishes, in microseconds. + # Edge delay represents wall-clock idle time between the predecessor's + # completion and the successor becoming ready. Not used for node execution + # time. None means no extra edge delay. + delay_after_predecessor_us: Annotated[float, msgspec.Meta(ge=0)] | None = None + # Minimum wait time on the successor after all predecessors are satisfied, + # in microseconds. Stacks with `delay_after_predecessor_us`; the runtime + # takes the later of the two when both are set. None means no minimum. + min_start_delay_us: Annotated[float, msgspec.Meta(ge=0)] | None = None + # Idle / scheduling delay after the predecessor DISPATCHES (its firing + # gate clears and it proceeds to its endpoint call), in microseconds. The + # successor does NOT wait for the predecessor's completion: the runtime + # schedules it at predecessor dispatch and gates it at dispatch + delay. + # Mutually exclusive with `delay_after_predecessor_us` (validator rule 54). + # None means this edge is not start-anchored. + delay_after_predecessor_start_us: Annotated[float, msgspec.Meta(ge=0)] | None = None + # Idle delay after the predecessor's OBSERVED FIRST TOKEN, in microseconds. + # Refines a start-anchored edge for children recorded post-TTFT: when the + # runtime observes the predecessor's first token it gates the successor at + # first_token_wall + this delay; when the predecessor terminates without one + # the gate falls back to dispatch + delay_after_predecessor_start_us. Only + # valid alongside delay_after_predecessor_start_us (validator rule 55). + delay_after_predecessor_first_token_us: ( + Annotated[float, msgspec.Meta(ge=0)] | None + ) = None + + +class NodeKind(str, Enum): + LLM = "llm" + + +class ExpectedTokens(msgspec.Struct, frozen=True, kw_only=True, omit_defaults=True): + """Engine-validation predictions.""" + + # Predicted input tokens served from cache (maps to OTel + # `gen_ai.usage.cache_read.input_tokens`). + cache_read_tokens: int | None = None + # Predicted input tokens written to cache on this request (maps to OTel + # `gen_ai.usage.cache_creation.input_tokens`). + cache_creation_tokens: int | None = None + input_tokens: int | None = None # Predicted total input tokens. + output_tokens: int | None = None # Predicted output tokens. + + +class _BaseNode( + msgspec.Struct, + frozen=True, + kw_only=True, + omit_defaults=True, + forbid_unknown_fields=True, +): + """Common fields for graph nodes.""" + + metadata: dict[str, Any] = {} # Span attributes under aiperf.meta.* + # Minimum wait time (in microseconds) after all predecessors are satisfied, + # before this node may start. The runtime takes the max of this and any + # incoming edge-level + # `min_start_delay_us` / `delay_after_predecessor_us` when computing the + # firing gate (see `TraceExecutor._apply_firing_delay`). None means no + # node-level minimum (edge-level gates still apply). + min_start_delay_us: Annotated[float, msgspec.Meta(ge=0)] | None = None + # Recorded wall-clock offset of this node from the trace's arrival, in + # microseconds. Populated by adapters that carry per-node timestamps (e.g. + # weka); used by snapshot-at-t* replay to partition nodes into warmup vs + # profiled. + arrival_offset_us: Annotated[int, msgspec.Meta(ge=0)] | None = None + # AND-fan-in input requirements. Empty list = OR-fan-in successor-walk + # fireability (default). Non-empty list = node fires only when every + # requirement is satisfied. + inputs: list[ChannelRequirement] = [] + + @property + def node_type(self) -> NodeKind: + # msgspec stores the tag in __struct_config__, not as an attribute. + return NodeKind(type(self).__struct_config__.tag) + + @property + def write_channels(self) -> list[str]: + """Channels this node writes to its parent graph. + + Each concrete node kind overrides this. Authoritative for any consumer + that needs to know what a node produces (visualizers, validators, + explain/export, etc.) -- instead of pattern-matching on `output` at + every callsite. + """ + return [] + + +class LlmNode( + _BaseNode, + frozen=True, + kw_only=True, + omit_defaults=True, + forbid_unknown_fields=True, + tag_field="node_type", + tag=NodeKind.LLM.value, +): + """LLM node.""" + + # Prompt grammar items; see `docs/benchmark-modes/agentic.md` "Prompt grammar". + prompt: list[Any] + output: str # Channel name to capture the model response into. + streaming: bool = True # Whether to dispatch streaming. + # Model name dispatched for this call, named like `Turn.model` (the + # recorded model for replay adapters). Folded into the wire body by the + # envelope builder (`store_builder._trie_envelope`); None falls through to + # the run's --model at materialize. + model: str | None = None + # Generation cap for this call, named/typed like `Turn.max_tokens`. The + # worker maps it to the endpoint's wire token field (max_completion_tokens, + # or max_tokens under --use-legacy-max-tokens); the envelope builder folds + # it into the wire body (`store_builder._trie_envelope`). A hand-authored + # positional `extra_body` entry wins over the fold. None leaves generation + # uncapped. + max_tokens: Annotated[int, msgspec.Meta(ge=1)] | None = None + # OpenAI-compatible tool definitions for this call, named like + # `Turn.raw_tools`. Folded into the wire body as `tools` by the envelope + # builder. dag_jsonl resolves lineage inheritance (the one per-turn field + # that walks history) BEFORE stamping, so the field always carries the + # effective definitions for this node. + raw_tools: list[dict[str, Any]] | None = None + # Per-call HTTP headers, named like `Turn.extra_headers`. Attached to the + # request HEADERS by the worker, never the body (dynamo stamps its + # x-dynamo-* session-identity headers here; the worker uniquifies them per + # replay instance). + extra_headers: dict[str, str] | None = None + expected: ExpectedTokens | None = None # Engine-prediction comparators. + # Non-native per-call request-body fields (temperature, top_p, seed, + # vendor tunables), named like `Turn.extra_body` and merged into the top + # level of the dispatched body. Model / stream / token cap / tools ride + # the native fields above, not here. None falls through to global CLI + # defaults. + extra_body: dict[str, Any] | None = None + # Leading hash-id blocks that would hit an infinite per-trace prefix + # cache, named like `Turn.theoretical_prefix_cache_hit_blocks`. Stamped by + # the shared trie build (`segment_ir.prefix_cache`); None when the node's + # request carries no hash blocks. + theoretical_prefix_cache_hit_blocks: Annotated[int, msgspec.Meta(ge=0)] | None = ( + None + ) + # Total hash-id blocks considered for the theoretical prefix-cache count, + # named like `Turn.theoretical_prefix_cache_total_blocks`. Pairs with + # `theoretical_prefix_cache_hit_blocks`. + theoretical_prefix_cache_total_blocks: Annotated[int, msgspec.Meta(ge=0)] | None = ( + None + ) + + @property + def write_channels(self) -> list[str]: + return [self.output] + + +class ProvenanceSpec(msgspec.Struct, frozen=True, kw_only=True, omit_defaults=True): + """Where this benchmark file came from.""" + + source: str = "hand-authored" # Origin tag. + tool: str = "manual" # Generating tool plus version, or 'manual'. + captured_at: str | None = None # ISO-8601 capture timestamp. + notes: str | None = None # Free-form notes. + # Catch-all for adapter-specific provenance keys; folded in at the codec + # boundary like the other catch-alls. + extra: dict[str, Any] = {} + + +NodeUnion = LlmNode + + +class GraphRecord( + msgspec.Struct, + frozen=True, + kw_only=True, + omit_defaults=True, + forbid_unknown_fields=True, +): + """Topology record.""" + + version: str = "2.0" # Schema major version. + provenance: ProvenanceSpec = msgspec.field( + default_factory=ProvenanceSpec + ) # Origin metadata. + system: str | None = None # Optional system prompt for linear chat. + state: dict[str, ChannelSpec] = {} # Explicit channel declarations. + nodes: dict[str, NodeUnion] = {} # Node id to node spec. + edges: list[StaticEdge] = [] # Edge list. + + +class TraceRecord( + msgspec.Struct, + frozen=True, + kw_only=True, + omit_defaults=True, + forbid_unknown_fields=True, +): + """Per-trace data.""" + + id: str # Stable trace identifier. + # Opaque labels. Round-tripped through the codec/sidecar; no runtime + # consumer reads them yet (authorable surface, kept for provenance). + tags: list[str] = [] + # Selects this trace's top-level graph from `ParsedGraph.graphs`. None (the + # default) means the trace runs against the single `ParsedGraph.graph`. A + # non-None value names a key in `ParsedGraph.graphs`. Resolution always goes + # through `resolve_trace_graph`. + graph_ref: str | None = None + # Linear-chat shorthand: equivalent to init.messages. + messages: list[dict[str, Any]] | None = None + # Initial channel values at t=0. + initial_state: dict[str, Any] = {} + # Node id to {channel: value}. Authorable surface for hand-written native + # files (read by `auto_derive` for channel inference); adapters never + # populate it, and the structural sidecar strip clears it. + replay_outputs: dict[str, dict[str, Any]] = {} + + +class ParsedGraph(msgspec.Struct, frozen=True, kw_only=True, omit_defaults=True): + """In-memory representation of a parsed graph workload file. Not yet auto-derived or validated.""" + + # Default / single-workload topology (defaults if absent in file). For + # single-graph workloads this is THE graph every trace runs against and + # `graphs` is empty. For multi-graph workloads it is the first source's graph + # as the default slot; per-trace resolution goes through + # `graphs[trace.graph_ref]` via `resolve_trace_graph`. + graph: GraphRecord = msgspec.field(default_factory=GraphRecord) + # Per-trace top-level graphs for multi-graph workloads, keyed by the key each + # `TraceRecord.graph_ref` names (the trace id, for the Weka + # heterogeneous-directory adapter). Empty for single-graph workloads. + # Resolve a trace's graph with `resolve_trace_graph`. + graphs: dict[str, GraphRecord] = {} + traces: list[TraceRecord] = [] # Trace records in file order. + # The content-addressed ``segment_ir.pool.SegmentPool`` whose entries every + # ``LlmNode``'s ``metadata["trie"]["prompt_segment_ids"]`` path indexes. + # Set by every live producer (weka/dynamo trie builds and the native + # lowering); the build plane drains it into a + # ``GraphSegmentUnifiedBackingStore`` so the worker can materialize each + # node's prompt. Typed ``SegmentPool | None`` (msgspec handles the plain + # dataclass) so the msgpack codec round-trips the real pool across worker + # processes instead of decoding it to a bare ``dict``. + segment_pool: SegmentPool | None = None + + +def resolve_trace_graph(parsed: ParsedGraph, trace: TraceRecord) -> GraphRecord: + """Return the top-level :class:`GraphRecord` a trace runs against. + + For single-graph workloads (``trace.graph_ref is None``) this is the + shared ``parsed.graph`` -- the byte-unchanged single-file / hand-authored + path. For multi-graph workloads (a heterogeneous directory of Weka traces) + ``trace.graph_ref`` names a key in ``parsed.graphs`` so each trace resolves + to its own distinct topology. Falls back to ``parsed.graph`` when the ref + is set but missing from ``parsed.graphs`` (defensive: tolerate a single bad + ref at runtime rather than crash the whole workload). + """ + ref = trace.graph_ref + if ref is None: + return parsed.graph + return parsed.graphs.get(ref, parsed.graph) + + +# --------------------------------------------------------------------------- +# Prompt-walk helpers +# +# Shared helpers for walking LLM-node prompt arrays and extracting channel refs. +# Used by ``auto_derive`` (to declare implicit channels) and ``native_lowering`` +# (channel-ref extraction while lowering prompt arrays). Keeping them here with +# the node models avoids drifting copies. +# --------------------------------------------------------------------------- + + +def _channel_ref(s: str) -> str | None: + """Return the channel name for a `@chan` ref, or None for plain text / `@@` escape.""" + if not s.startswith("@") or s.startswith("@@"): + return None + return s[1:] + + +def _collect_prompt_array_channels(node: LlmNode) -> list[str]: + """Return channel names referenced at the top level of a prompt array.""" + out: list[str] = [] + for item in node.prompt: + if isinstance(item, str): + ch = _channel_ref(item) + if ch is not None: + out.append(ch) + return out + + +def _collect_content_block_channels(node: LlmNode) -> list[str]: + """Return channel names referenced inside chat-message content blocks.""" + out: list[str] = [] + for item in node.prompt: + if not isinstance(item, dict): + continue + content = item.get("content") + if not isinstance(content, list): + continue + for block in content: + if not isinstance(block, str): + continue + ch = _channel_ref(block) + if ch is not None: + out.append(ch) + return out diff --git a/src/aiperf/dataset/graph/native_lowering.py b/src/aiperf/dataset/graph/native_lowering.py new file mode 100644 index 0000000000..6f05a26876 --- /dev/null +++ b/src/aiperf/dataset/graph/native_lowering.py @@ -0,0 +1,745 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Native-graph lowering onto the unified interned segment store. + +Lowers a hand-authored (native) :class:`ParsedGraph` onto the same content +plane the weka and dynamo adapters emit: every ``LlmNode`` prompt is +canonicalized to ``[{role, content: str}, ...]`` messages, interned into a +content-addressed :class:`SegmentPool`, and stamped with +``metadata["trie"]["prompt_segment_ids"]``. The build plane then drains the +pool into the ONE unified interned store exactly as it does for the trace +adapters (``dataset_manager`` routes on ``parsed.segment_pool``), and the +worker materializes prompts by handle. + +Dynamic content (dynamic-content-pool spec, Phase 3): a ``@channel`` prompt +reference whose channel is written by ancestor ``LlmNode``\\ s lowers to +DYNAMIC SLOTS — an ordered assembly program stamped as +``metadata["trie"]["assembly"]`` whose tokens interleave interned static +segments with references to producer nodes. At run time the sticky worker +fills slots from its dynamic pool (the producers' actual responses). + +- Array-level splices (messages channels): composition at the read point + reconstructs the FULL user/assistant alternation = init-seeded messages, then + for each completion-ordered ancestor writer its authored user turn (its + ``delta`` = its prompt minus its own ``@C`` splice, static, interned) followed + by a slot for its live reply. Every writer except the first must itself read + the channel (the accumulate shape) so injected input gates make build-time + order equal runtime commit order. The reading node is excluded (read observes + pre-write state), so the linear-chat shorthand stays fully static. This is why + the naive single-channel chain (``output: C`` on each turn, ``["@C", ]``) + yields a well-formed conversation with no re-stating of prior user turns. +- Block-level refs (text channels): at most one ancestor writer; the slot + value is the writer's response, composed INTO the containing message + (static text parts + slot-value parts, concatenated). +- Producers referenced by any slot are stamped ``capture: true``; readers get + an injected ``ChannelRequirement`` per spliced channel so their credits are + minted only after every producer resolved (completion gating regardless of + edge anchors). + +Trace-dependent content (init splices) lowers to PER-TRACE graphs +(``parsed.graphs[trace.id]`` + ``trace.graph_ref``) against one shared pool. +The lowering is a pure function of the parsed file (content-addressed ids, no +RNG, no clock), so the dataset plane and timing plane re-parses stamp +identical graphs. + +Canonical wire shape: message ``content`` is always a plain string; a content +list of string blocks is concatenated in order (no separator). Typed / +multimodal / directive blocks are gated. +""" + +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass, field +from typing import Any + +import msgspec + +from aiperf.dataset.graph.models import ( + ChannelRequirement, + GraphRecord, + LlmNode, + ParsedGraph, + StaticEdge, + TraceRecord, + _channel_ref, + _collect_content_block_channels, + _collect_prompt_array_channels, +) +from aiperf.dataset.graph.segment_ir.envelope import stamp_prompt_segment_ids +from aiperf.dataset.graph.segment_ir.pool import SegmentPool + +__all__ = ["lower_native_to_unified"] + + +def lower_native_to_unified(parsed: ParsedGraph) -> ParsedGraph: + """Stamp a native parse onto the unified segment-store content plane.""" + _gate_unsupported_topology(parsed) + plan = _analyze_dynamic_refs(parsed) + + pool = SegmentPool() + has_refs = any( + _collect_prompt_array_channels(node) or _collect_content_block_channels(node) + for node in parsed.graph.nodes.values() + ) + if not has_refs: + stamped = _stamp_graph(pool, parsed.graph, plan, trace=None) + return msgspec.structs.replace(parsed, graph=stamped, segment_pool=pool) + + if not parsed.traces: + raise NotImplementedError( + "traces: '@channel' prompt references splice per-trace initial_state; " + "a workload with no trace records has no content to lower" + ) + seen_ids: set[str] = set() + graphs: dict[str, GraphRecord] = {} + new_traces: list[TraceRecord] = [] + for trace in parsed.traces: + if trace.graph_ref is not None: + raise NotImplementedError( + f"traces[{trace.id}].graph_ref: pre-set graph refs are not " + f"supported by the unified-store lowering" + ) + if trace.id in seen_ids: + raise NotImplementedError( + f"traces[{trace.id}]: duplicate trace ids cannot key per-trace " + f"lowered graphs" + ) + seen_ids.add(trace.id) + graphs[trace.id] = _stamp_graph(pool, parsed.graph, plan, trace=trace) + new_traces.append(msgspec.structs.replace(trace, graph_ref=trace.id)) + first_graph = graphs[new_traces[0].id] + return msgspec.structs.replace( + parsed, + graph=first_graph, + graphs=graphs, + traces=new_traces, + segment_pool=pool, + ) + + +def _gate_unsupported_topology(parsed: ParsedGraph) -> None: + for nid, node in parsed.graph.nodes.items(): + if not isinstance(node, LlmNode): + raise NotImplementedError( + f"graph.nodes.{nid}: node_type '{node.node_type}' is not supported " + f"by the unified-store lowering (LlmNode only)" + ) + # replay_reducers flips the node's write to replace_channel, wiping the + # init seed the static composition relies on. + if "replay_reducers" in (node.metadata or {}): + raise NotImplementedError( + f"graph.nodes.{nid}.metadata.replay_reducers: reducer overrides " + f"are not supported by the unified-store lowering" + ) + for edge in parsed.graph.edges: + if not isinstance(edge, StaticEdge): + continue + # Rule-55 shape, enforced here because the profile pipeline lowers at + # parse time without running validate(): a first-token anchor with no + # dispatch fallback would otherwise be silently lowered as a + # completion edge (see _completion_adjacency). + if ( + edge.delay_after_predecessor_first_token_us is not None + and edge.delay_after_predecessor_start_us is None + ): + raise NotImplementedError( + f"graph.edges[{edge.source}->{edge.target}]: edge is " + f"first-token-anchored (delay_after_predecessor_first_token_us) " + f"but sets no delay_after_predecessor_start_us; the runtime " + f"needs the start delay as the dispatch fallback when the " + f"predecessor terminates without a first token (validator " + f"rule 55)" + ) + + +# --------------------------------------------------------------------------- +# Dynamic-ref analysis: which prompt refs become slots, and their gates +# --------------------------------------------------------------------------- + + +@dataclass(slots=True) +class _LoweringPlan: + """Per-parse dynamic-slot plan (topology-derived; trace-independent).""" + + array_slots: dict[tuple[str, str], list[str]] = field(default_factory=dict) + """(reader, channel) -> completion-ordered writer node ids (array splice).""" + block_slots: dict[tuple[str, str], str] = field(default_factory=dict) + """(reader, channel) -> single writer node id (block-level text ref).""" + capture_nodes: set[str] = field(default_factory=set) + """Producers referenced by any slot; stamped ``capture: true``.""" + injected_inputs: dict[str, dict[str, int]] = field(default_factory=dict) + """reader -> {channel: required write count} (completion gating).""" + + +def _completion_adjacency(graph: GraphRecord) -> dict[str, set[str]]: + """src -> targets over COMPLETION edges only. + + Start-anchored and first-token-anchored edges schedule the successor while + the predecessor is still in flight, so they contribute no completion + ordering; excluding them makes "ancestor" mean commit-before by + construction. + """ + adj: dict[str, set[str]] = {} + for edge in graph.edges: + if not isinstance(edge, StaticEdge): + continue + if edge.delay_after_predecessor_start_us is not None: + continue + adj.setdefault(edge.source, set()).add(edge.target) + return adj + + +def _ancestors(adj: dict[str, set[str]], target: str) -> set[str]: + rev: dict[str, set[str]] = {} + for src, dsts in adj.items(): + for dst in dsts: + rev.setdefault(dst, set()).add(src) + seen: set[str] = set() + queue: deque[str] = deque(rev.get(target, ())) + while queue: + node = queue.popleft() + if node in seen: + continue + seen.add(node) + queue.extend(rev.get(node, ())) + return seen + + +def _anchored_pairs(graph: GraphRecord) -> set[tuple[str, str]]: + return { + (e.source, e.target) + for e in graph.edges + if isinstance(e, StaticEdge) and e.delay_after_predecessor_start_us is not None + } + + +def _analyze_dynamic_refs(parsed: ParsedGraph) -> _LoweringPlan: + """Classify every ``@channel`` ref as static or slot; gate the unlowerable. + + Slot legality: writers must be completion-ancestors of + the reader; array-splice writers must form a completion-ordered chain in + which every writer except the first also reads the channel; block refs + take at most one writer and must not mix with init seeding; anchored + edges between a producer and its reader (or successive writers) are + contradictory timing intent. + """ + graph = parsed.graph + plan = _LoweringPlan() + writers: dict[str, list[str]] = {} + for nid, node in graph.nodes.items(): + writers.setdefault(node.output, []).append(nid) + adj = _completion_adjacency(graph) + anchored = _anchored_pairs(graph) + ancestors: dict[str, set[str]] = {nid: _ancestors(adj, nid) for nid in graph.nodes} + array_refs: dict[str, set[str]] = { + nid: set(_collect_prompt_array_channels(node)) + for nid, node in graph.nodes.items() + } + + for nid, node in graph.nodes.items(): + loc = f"graph.nodes.{nid}.prompt" + block_refs = set(_collect_content_block_channels(node)) + for ch in sorted(array_refs[nid] | block_refs): + all_writers = sorted(w for w in writers.get(ch, []) if w != nid) + if not all_writers: + continue # static (init-seeded or self-written) + # A reader sees only the channel's writers that COMMITTED before it + # (its completion-ancestors). DESCENDANT writers are future writes + # this read legitimately precedes (read-before-write) — excluded, + # not an error (this is what lets the accumulate chain's root read + # @C past its own downstream re-writers). CONCURRENT writers + # (neither ancestor nor descendant) are a genuine race → gate. + others = [w for w in all_writers if w in ancestors[nid]] + concurrent = [ + w + for w in all_writers + if w not in ancestors[nid] and nid not in ancestors[w] + ] + if concurrent: + raise NotImplementedError( + f"{loc}: '@{ch}' is written by LlmNode(s) {concurrent} " + f"concurrent with '{nid}' (neither ancestor nor descendant); " + f"a dynamic slot needs a deterministic commit order" + ) + if not others: + continue # only future (descendant) writers → static for this reader + for w in others: + if (w, nid) in anchored: + raise NotImplementedError( + f"graph.edges[{w}->{nid}]: a start/first-token-anchored " + f"edge from slot producer '{w}' to its reader '{nid}' " + f"is contradictory timing intent (the anchor dispatches " + f"the reader while the producer is still in flight)" + ) + if ch in block_refs and ch not in array_refs[nid]: + if len(others) > 1: + raise NotImplementedError( + f"{loc}: '@{ch}' block ref has multiple writers " + f"{others}; overwrite channels latch a single writer " + f"for the trace lifetime" + ) + if any(ch in t.initial_state for t in parsed.traces): + raise NotImplementedError( + f"{loc}: '@{ch}' block ref channel is both init-seeded " + f"and written by '{others[0]}'; the init value is never " + f"observable past the completion gate — remove one " + f"source" + ) + plan.block_slots[(nid, ch)] = others[0] + else: + _gate_array_splice_shape(loc, nid, node, ch, writers) + ordered = _ordered_writer_chain(loc, ch, others, ancestors, anchored) + for i, w in enumerate(ordered[1:], start=1): + if ch not in array_refs[w]: + raise NotImplementedError( + f"{loc}: '@{ch}' writers must chain through reads — " + f"writer '{w}' does not itself splice '@{ch}', so " + f"its commit order after " + f"'{ordered[i - 1]}' is not gated" + ) + # The ROOT writer has no committed ancestor writers, so the + # reader-side Gate A above never runs for it; a non-leading + # (or repeated) '@{ch}' there would survive to _delta_messages, + # which drops only a LEADING splice and re-expands the init + # seed for anything else — silently duplicating it in every + # reader's reconstruction. + for w in ordered: + _gate_array_splice_shape( + f"graph.nodes.{w}.prompt", w, graph.nodes[w], ch, writers + ) + plan.array_slots[(nid, ch)] = ordered + plan.capture_nodes.update(others) + plan.injected_inputs.setdefault(nid, {})[ch] = len(others) + + _gate_init_bearing_root_reads(parsed, plan, writers, adj, array_refs) + return plan + + +def _gate_array_splice_shape( + loc: str, nid: str, node: LlmNode, ch: str, writers: dict[str, list[str]] +) -> None: + """Gate A: a slot channel's ``@C`` array-splice must be single, and leading + for any node that also writes ``C``. + + Runs for every reconstructing reader AND every writer in its ordered + chain (the root writer has no committed ancestor writers, so only the + chain-side call reaches it). A repeated ``@C`` would double-expand the + whole reconstructed conversation. A writer whose ``@C`` is not + the first prompt item has an ill-defined ``delta`` (items before ``@C`` + would be misplaced after the reconstructed history), so its ``@C`` must + lead. + """ + positions = [ + i + for i, item in enumerate(node.prompt) + if isinstance(item, str) and _channel_ref(item) == ch + ] + if len(positions) > 1: + raise NotImplementedError( + f"{loc}: '@{ch}' appears {len(positions)} times as an array splice; " + f"a slot channel may be spliced at most once per node (a repeat " + f"would duplicate the reconstructed conversation)" + ) + is_writer = nid in writers.get(ch, []) + if is_writer and positions and positions[0] != 0: + raise NotImplementedError( + f"{loc}: node writes '{ch}' and splices '@{ch}', so '@{ch}' must be " + f"the first prompt item (its authored delta is everything after the " + f"history it reads); move '@{ch}' to the front" + ) + + +def _gate_init_bearing_root_reads( + parsed: ParsedGraph, + plan: _LoweringPlan, + writers: dict[str, list[str]], + adj: dict[str, set[str]], + array_refs: dict[str, set[str]], +) -> None: + """Gate B: an init-bearing messages channel's ROOT writer must read ``@C`` + — but only when some reader actually reconstructs the channel. + + A reconstructing reader splices ``@C`` with committed writers (an array + slot), expanding ``[init, …writer deltas/replies…]``. If such a reader + exists and the root writer does not read ``@C``, the root dispatches + without the seeded prefix while readers reconstruct ``[init, …]`` — + inserting context the root never saw (unfaithful) and diverging its + interned handles from the reconstruction. Requiring the root to read + ``@C`` makes every writer see the same prefix it is later reconstructed + with. + + When NO reader reconstructs the channel (e.g. a write-only channel whose + init seed is consumed only as a static block ref, the "rewrite your own + draft" workload), the divergence cannot occur and the gate must not fire. + """ + reconstructed = {ch for (_, ch) in plan.array_slots} + init_channels = {ch for t in parsed.traces for ch in t.initial_state} + for ch in sorted(init_channels): + ch_writers = writers.get(ch, []) + if not ch_writers: + continue # pure static init splice; no reconstruction + if ch not in reconstructed: + continue # nothing reconstructs [init, ...]; no divergence possible + ch_ancestors = {w: _ancestors(adj, w) for w in ch_writers} + roots = [ + w + for w in ch_writers + if not any(o in ch_ancestors[w] for o in ch_writers if o != w) + ] + for root in roots: + if ch not in array_refs.get(root, set()): + raise NotImplementedError( + f"graph.nodes.{root}.prompt: channel '{ch}' is init-seeded " + f"and written here, but this root writer does not splice " + f"'@{ch}'; it would dispatch without the seed while readers " + f"reconstruct it — add '@{ch}' as the first prompt item" + ) + + +def _ordered_writer_chain( + loc: str, + channel: str, + writer_ids: list[str], + ancestors: dict[str, set[str]], + anchored: set[tuple[str, str]], +) -> list[str]: + """Totally order writers by completion ancestry, or gate.""" + ordered = sorted( + writer_ids, key=lambda w: sum(1 for o in writer_ids if o in ancestors[w]) + ) + for earlier, later in zip(ordered, ordered[1:], strict=False): + if earlier not in ancestors[later]: + raise NotImplementedError( + f"{loc}: '@{channel}' writers {sorted(writer_ids)} are not " + f"totally completion-ordered ('{earlier}' and '{later}' are " + f"mutually unordered); parallel producers of one spliced " + f"channel have no deterministic assembly order" + ) + if (earlier, later) in anchored: + raise NotImplementedError( + f"graph.edges[{earlier}->{later}]: successive writers of " + f"spliced channel '@{channel}' must be completion-ordered; a " + f"start/first-token-anchored edge between them is " + f"contradictory timing intent" + ) + return ordered + + +# --------------------------------------------------------------------------- +# Stamping: canonical messages -> pool segments + assembly tokens +# --------------------------------------------------------------------------- + + +def _stamp_graph( + pool: SegmentPool, + graph: GraphRecord, + plan: _LoweringPlan, + trace: TraceRecord | None, +) -> GraphRecord: + stamped = { + nid: _stamp_node(pool, nid, node, graph.nodes, plan, trace) + for nid, node in graph.nodes.items() + } + return msgspec.structs.replace(graph, nodes=stamped) + + +def _stamp_node( + pool: SegmentPool, + node_id: str, + node: LlmNode, + nodes: dict[str, LlmNode], + plan: _LoweringPlan, + trace: TraceRecord | None, +) -> LlmNode: + assembly = _assemble_prompt(pool, node_id, node, nodes, plan, trace) + segment_ids = [token["seg"] for token in assembly if "seg" in token] + if not assembly: + raise NotImplementedError( + f"graph.nodes.{node_id}.prompt: an empty prompt cannot be lowered " + f"to the unified store" + ) + has_slots = any("seg" not in token for token in assembly) + extra: dict[str, Any] = {} + if has_slots: + extra["assembly"] = assembly + if node_id in plan.capture_nodes: + extra["capture"] = True + injected = plan.injected_inputs.get(node_id) + if injected: + node = msgspec.structs.replace( + node, inputs=_upsert_inputs(node.inputs, injected) + ) + return stamp_prompt_segment_ids(node, segment_ids, extra=extra or None) + + +def _upsert_inputs( + existing: list[ChannelRequirement], injected: dict[str, int] +) -> list[ChannelRequirement]: + """Merge injected completion gates with authored input requirements.""" + out: list[ChannelRequirement] = [] + remaining = dict(injected) + for req in existing: + needed = remaining.pop(req.channel, None) + if needed is not None and isinstance(req.count, int) and req.count < needed: + out.append(ChannelRequirement(channel=req.channel, count=needed)) + else: + out.append(req) + for channel, needed in sorted(remaining.items()): + out.append(ChannelRequirement(channel=channel, count=needed)) + return out + + +def _assemble_prompt( + pool: SegmentPool, + node_id: str, + node: LlmNode, + nodes: dict[str, LlmNode], + plan: _LoweringPlan, + trace: TraceRecord | None, +) -> list[dict[str, Any]]: + """Walk the prompt into assembly tokens, interning static messages. + + Token shapes (node-id references; the store builder resolves them to + handles/ordinals): ``{"seg": }`` static message, + ``{"s": {"src": }}`` array-splice slot, and + ``{"m": {"role": r, "parts": [{"t": text} | {"sv": }]}}`` + composed message. Static segments parent-chain in prompt order. + + An array ``@C`` slot expansion reconstructs full alternation: init messages, + then for each ordered writer its ``delta`` (authored user turn, static, + interned inline for this trace) followed by its reply slot. + """ + loc = f"graph.nodes.{node_id}.prompt" + assembly: list[dict[str, Any]] = [] + parent: str | None = None + + def _intern(message: dict[str, str]) -> None: + nonlocal parent + parent = pool.add_text( + role=message["role"], content=message["content"], parent_id=parent + ) + assembly.append({"seg": parent}) + + for item in node.prompt: + if isinstance(item, str): + channel = _channel_ref(item) + if channel is None: + raise NotImplementedError( + f"{loc}: top-level string items must be '@channel' messages splices" + ) + writer_chain = plan.array_slots.get((node_id, channel)) + if writer_chain is None: + for message in _init_messages(channel, trace, required=True): + _intern(message) + else: + for message in _init_messages(channel, trace, required=False): + _intern(message) + for writer in writer_chain: + for message in _delta_messages( + writer, nodes[writer], channel, plan, trace + ): + _intern(message) + assembly.append({"s": {"src": writer}}) + elif isinstance(item, dict): + role = item.get("role") + if not isinstance(role, str) or not role: + raise NotImplementedError( + f"{loc}: message items require a non-empty 'role' string" + ) + extra_keys = sorted(set(item) - {"role", "content"}) + if extra_keys: + raise NotImplementedError( + f"{loc}: message keys {extra_keys} are not representable in " + f"the unified store (role/content only); dropping them " + f"silently would corrupt the authored wire bytes" + ) + parts = _content_parts(loc, node_id, item.get("content"), plan, trace) + if any("sv" in part for part in parts): + assembly.append({"m": {"role": role, "parts": parts}}) + else: + _intern( + { + "role": role, + "content": "".join(part["t"] for part in parts), + } + ) + else: + raise NotImplementedError( + f"{loc}: unsupported prompt item type " + f"'{type(item).__name__}' (expected str or message dict)" + ) + return assembly + + +def _content_parts( + loc: str, + node_id: str, + content: Any, + plan: _LoweringPlan, + trace: TraceRecord | None, +) -> list[dict[str, Any]]: + """Resolve message content into ``{"t": text}`` / ``{"sv": node_id}`` parts.""" + if isinstance(content, str): + return [{"t": content}] + if isinstance(content, list): + parts: list[dict[str, Any]] = [] + for block in content: + if not isinstance(block, str): + raise NotImplementedError( + f"{loc}: non-string content blocks (directives / typed " + f"blocks) are not supported by the unified-store lowering" + ) + channel = _channel_ref(block) + if channel is None: + parts.append({"t": block[1:] if block.startswith("@@") else block}) + elif (node_id, channel) in plan.block_slots: + parts.append({"sv": plan.block_slots[(node_id, channel)]}) + else: + parts.append({"t": _init_text(loc, channel, trace)}) + return parts + raise NotImplementedError( + f"{loc}: message content must be a string or a list of string blocks" + ) + + +def _delta_messages( + writer_id: str, + writer: LlmNode, + channel: str, + plan: _LoweringPlan, + trace: TraceRecord | None, +) -> list[dict[str, str]]: + """A writer's authored contribution to conversation ``channel``: its prompt + messages minus the ``@channel`` splice it read (the recursive history). + + Computed inline per trace, in the WRITER's node_id context so block-slot + classification resolves against the writer. v1 requires the delta + to be **static** — it may not itself produce a live slot; gate loudly + otherwise (a writer that both contributes to a conversation and splices + another live channel is v2). Gate A (run for every writer in the slot + chain, root included) guarantees a writer's ``@channel`` splice, if + present, is the leading item, so dropping ``prompt[0]`` when it is that + splice yields the whole delta. + """ + loc = f"graph.nodes.{writer_id}.prompt (delta for @{channel})" + items = list(writer.prompt) + if items and isinstance(items[0], str) and _channel_ref(items[0]) == channel: + items = items[1:] # drop the leading @channel history read + out: list[dict[str, str]] = [] + for item in items: + if isinstance(item, str): + ch2 = _channel_ref(item) + if ch2 is None: + raise NotImplementedError( + f"{loc}: top-level string items must be '@channel' splices" + ) + if (writer_id, ch2) in plan.array_slots: + raise NotImplementedError( + f"{loc}: a conversation writer's delta may not splice live " + f"channel '@{ch2}' (v1 static-delta only)" + ) + out.extend(_init_messages(ch2, trace, required=True)) + elif isinstance(item, dict): + role = item.get("role") + if not isinstance(role, str) or not role: + raise NotImplementedError( + f"{loc}: message items require a non-empty 'role' string" + ) + extra_keys = sorted(set(item) - {"role", "content"}) + if extra_keys: + raise NotImplementedError( + f"{loc}: message keys {extra_keys} are not representable in " + f"the unified store (role/content only)" + ) + parts = _content_parts(loc, writer_id, item.get("content"), plan, trace) + if any("sv" in part for part in parts): + raise NotImplementedError( + f"{loc}: a conversation writer's delta may not splice a live " + f"block channel (v1 static-delta only)" + ) + out.append({"role": role, "content": "".join(p["t"] for p in parts)}) + else: + raise NotImplementedError( + f"{loc}: unsupported prompt item type '{type(item).__name__}'" + ) + return out + + +# --------------------------------------------------------------------------- +# Init-value resolution (per trace) +# --------------------------------------------------------------------------- + + +def _init_messages( + channel: str, trace: TraceRecord | None, *, required: bool +) -> list[dict[str, str]]: + if trace is None or channel not in trace.initial_state: + if not required: + return [] # slot-backed splice: init prefix is optional + trace_id = trace.id if trace is not None else "?" + raise NotImplementedError( + f"traces[{trace_id}].initial_state.{channel}: '@{channel}' splices " + f"init-seeded content; every trace firing the reader must supply it" + ) + value = trace.initial_state[channel] + trace_id = trace.id + loc = f"traces[{trace_id}].initial_state.{channel}" + if not isinstance(value, list): + raise NotImplementedError( + f"{loc}: a messages splice requires a list of message dicts" + ) + out: list[dict[str, str]] = [] + for i, message in enumerate(value): + if not isinstance(message, dict) or not isinstance(message.get("role"), str): + raise NotImplementedError( + f"{loc}[{i}]: messages must be dicts with a 'role' string" + ) + extra_keys = sorted(set(message) - {"role", "content"}) + if extra_keys: + raise NotImplementedError( + f"{loc}[{i}]: message keys {extra_keys} are not representable in " + f"the unified store (role/content only); dropping them silently " + f"would corrupt the authored wire bytes" + ) + content = _canonical_static_content(f"{loc}[{i}]", message.get("content")) + out.append({"role": message["role"], "content": content}) + return out + + +def _canonical_static_content(loc: str, content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for block in content: + if not isinstance(block, str): + raise NotImplementedError( + f"{loc}: non-string content blocks (directives / typed " + f"blocks) are not supported by the unified-store lowering" + ) + if _channel_ref(block) is not None: + raise NotImplementedError( + f"{loc}: '@' refs are not supported inside init-seeded " + f"message content" + ) + parts.append(block[1:] if block.startswith("@@") else block) + return "".join(parts) + raise NotImplementedError( + f"{loc}: message content must be a string or a list of string blocks" + ) + + +def _init_text(loc: str, channel: str, trace: TraceRecord | None) -> str: + if trace is None or channel not in trace.initial_state: + trace_id = trace.id if trace is not None else "?" + raise NotImplementedError( + f"traces[{trace_id}].initial_state.{channel}: '@{channel}' splices " + f"init-seeded content; every trace firing the reader must supply it" + ) + value = trace.initial_state[channel] + if not isinstance(value, str): + raise NotImplementedError( + f"{loc}: '@{channel}' text refs require a string initial_state value" + ) + return value diff --git a/src/aiperf/dataset/graph/parse_context.py b/src/aiperf/dataset/graph/parse_context.py new file mode 100644 index 0000000000..dd32c3e49c --- /dev/null +++ b/src/aiperf/dataset/graph/parse_context.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run-derived parse knobs threaded through the graph adapter protocol. + +:class:`GraphParseContext` is the single carrier for every run-config-derived +knob a graph adapter needs to parse byte-identically to the run. The parser +passes it opaquely (``adapter_cls.parse(path, ctx)``); each adapter maps the +fields it consumes onto its own entry function and ignores the rest. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +# Sentinel for ctx fields where ``None`` is a meaningful user value. Compare +# with ``is UNSET``, never ``==``. msgspec's singleton survives pickle and +# deepcopy with identity intact, so round-trips cannot silently break checks. +from msgspec import UNSET, UnsetType + + +@dataclass(frozen=True, slots=True) +class GraphParseContext: + """Run-derived knobs an adapter needs to parse byte-identically to the run. + + Field semantics: + * ``UNSET`` (where the type admits it) / ``None`` (elsewhere) means + "adapter default" — a ctx-less parse behaves exactly like today's + protocol-default entry. + * A SET value — including an explicit ``None`` on tri-state fields — is + forwarded to the adapter entry VERBATIM. ``idle_gap_cap_seconds=None`` + means DISABLE WARPING (the user's ``synthesis.idle_gap_cap_seconds: + null``), which is distinct from unset. + Adapters read the fields they consume and ignore the rest. Adapters must + only forward a field when it is set (not UNSET/None per the rules above), + so entry defaults that are not None (e.g. dynamo ``prompt_corpus="coding"``, + dag ``run_streaming=True``) are never clobbered by a partial ctx. + """ + + content_root_seed: int | None = None + """Run ``--random-seed`` for seed-dependent content synthesis (weka, dynamo).""" + + content_tokenizer: str | None = None + """Tokenizer name content synthesis decodes with (weka, dynamo).""" + + tokenizer_trust_remote_code: bool | None = None + """Run tokenizer trust flag; content-synthesizing adapters publish it to the + loader-preload env before building callbacks (weka, dynamo) so a registry + parse in a fresh process loads the same tokenizer the run does.""" + + tokenizer_revision: str | None = None + """Run tokenizer revision pin; published alongside trust (weka, dynamo).""" + + prompt_corpus: str | None = None + """Synthesis corpus selector (weka, dynamo).""" + + max_osl: int | None = None + """``--synthesis-max-osl`` cap on top-level chains (weka).""" + + num_dataset_entries: int | None = None + """Explicit ``entries`` cap on the run's default dataset + (``--num-dataset-entries``): the graph-plane ceiling on distinct traces + selected. ``None`` = unset (all eligible traces after filters).""" + + max_context_length: int | None = None + """``--max-context-length`` per-trace context ceiling (input+output tokens) + for graph-plane dataset selection. ``None`` = no context-length filter.""" + + idle_gap_cap_seconds: float | None | UnsetType = UNSET + """Idle-gap warp cap (weka, dynamo). TRI-STATE: UNSET = adapter default + (60s); a float = that cap; ``None`` = warping DISABLED (user's explicit + null).""" + + trajectory_start_max_ratio: float = 0.0 + """Resolved t* snapshot-window upper bound + (``--trajectory-start-max-ratio``, scenario-auto-applied when unset). + Consumed by the parse-time dynamic-slot gate; ``0.0`` = window OFF.""" + + default_model: str | None = None + """Worker dispatch fallback model stamped into node overrides (dag_jsonl).""" + + run_streaming: bool | None = None + """Resolved endpoint ``stream`` flag stamped onto nodes (dag_jsonl).""" + + delay_cap_seconds: float | None = None + """Legacy ``inter_turn_delay_cap_seconds`` clamp on authored delays (dag_jsonl).""" + + endpoint_extra: list[tuple[str, Any]] | None = None + """Run ``--extra-inputs`` pairs folded into node overrides (dag_jsonl).""" + + +def publish_ctx_tokenizer_env(ctx: GraphParseContext | None) -> None: + """Publish ``ctx``'s tokenizer trust/revision to the loader-preload env. + + No-op unless ``ctx.tokenizer_trust_remote_code`` is set (not None), so a + ctx-less or partial-ctx parse never clobbers values the run path already + published. ``tokenizer_revision`` passes verbatim — + :func:`~aiperf.dataset._mp_context.configure_loader_tokenizer_env` maps + ``None`` to ``"main"`` itself. Idempotent with the run path's publish + (same run-derived values). Lazy import: this module stays a pure-data + leaf at import time. + """ + if ctx is None or ctx.tokenizer_trust_remote_code is None: + return + from aiperf.dataset._mp_context import configure_loader_tokenizer_env + + configure_loader_tokenizer_env( + trust_remote_code=ctx.tokenizer_trust_remote_code, + revision=ctx.tokenizer_revision, + ) diff --git a/src/aiperf/dataset/graph/parser.py b/src/aiperf/dataset/graph/parser.py new file mode 100644 index 0000000000..b8d5f8a7ad --- /dev/null +++ b/src/aiperf/dataset/graph/parser.py @@ -0,0 +1,367 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Graph file parser backed by native readers and registered graph adapters.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, TypeAlias + +import msgspec +import orjson + +from aiperf.common.aiperf_logger import AIPerfLogger +from aiperf.dataset.graph.decode import GraphDecodeError, decode_graph +from aiperf.dataset.graph.models import ( + END_NODE_ID, + START_NODE_ID, + GraphRecord, + ParsedGraph, + StaticEdge, + TraceRecord, +) +from aiperf.dataset.graph.parse_context import GraphParseContext + + +class GraphParseError(ValueError): + """Raised when a graph workload file cannot be parsed.""" + + +def parse_graph( + path: str | Path, + *, + format: WorkloadFormat | None = None, + ctx: GraphParseContext | None = None, + **adapter_kwargs: Any, +) -> ParsedGraph: + """Parse a workload file into a ParsedGraph. + + Auto-detects the native graph format or delegates to registered graph + adapters, using extension and content sniffing as defined by the plugin + registry. Pass `format=` to override detection (CLI: `--graph-format`). + + `ctx` carries the run-derived parse knobs (see + :class:`~aiperf.dataset.graph.parse_context.GraphParseContext`) and is + passed opaquely to the selected adapter's ``parse(path, ctx)``; each + adapter maps the fields it consumes and ignores the rest. ``None`` (CLI + tooling / direct callers with no run config) keeps every adapter on its + protocol-default entry. + + ``**adapter_kwargs`` is format-agnostic plumbing for ADAPTER-SPECIFIC knobs + the build plane must inject as live objects the ``GraphParseContext`` (a + frozen bundle of run-derived VALUES) cannot carry -- currently only the + dynamo direct route's ``direct_store`` (``GraphStoreBuilder`` -> + :func:`~aiperf.dataset.graph.workload_detect.parse_graph_workload`). It + names no adapter: it is forwarded verbatim to ``adapter_cls.parse(path, + ctx, **adapter_kwargs)``, so a kwarg reaching an adapter whose ``parse`` + does not accept it fails loud with ``TypeError`` (not silently dropped). + The ``native`` path takes NO adapter kwargs (it has no adapter ``parse``), + so passing any with ``format="native"`` (or a file detected as native) + raises :class:`GraphParseError` up front rather than a confusing + downstream signature error. + """ + p = Path(path) + fmt = format or _detect(p) + if fmt == "native" and adapter_kwargs: + raise GraphParseError( + f"{p}: the native parse path accepts no adapter-specific kwargs, but " + f"received {sorted(adapter_kwargs)}; adapter_kwargs (e.g. dynamo's " + f"direct_store) are only valid for a registered graph adapter" + ) + try: + pb = ( + parse_native(p) + if fmt == "native" + else _parse_via_adapter(p, fmt, ctx, **adapter_kwargs) + ) + except (msgspec.ValidationError, GraphDecodeError) as e: + raise GraphParseError(_format_ir_error(p, e)) from e + return pb + + +def _format_ir_error( + path: Path, exc: msgspec.ValidationError | GraphDecodeError +) -> str: + """Render an IR (de)coding error as a single readable line. + + Without this wrapper, parse-time IR errors (e.g. ``max_iterations=0``, + out-of-range field values) escape ``parse_graph`` as + raw tracebacks. The CLI catches ``GraphParseError`` cleanly, so callers see + a friendly message instead of a stack trace. msgspec carries the field path + inline in its message (``... - at `$.nodes...```). + """ + return f"{path}: IR error: {exc}" + + +def _detect(path: Path) -> WorkloadFormat: + try: + return detect_format(path) + except WorkloadFormatError as e: + raise GraphParseError(str(e)) from e + + +def parse_native(path: Path) -> ParsedGraph: + """Native graph parse entry point. Used by `NativeGraphAdapter`. + + Native graph parsing is the canonical (non-adapter) path: JSONL or YAML + that already conforms to the graph schema. Auto-derive runs at the end + so callers get the same `ParsedGraph` shape regardless of input format. + """ + ext = path.suffix.lower() + records = _read_jsonl(path) if ext == ".jsonl" else _read_yaml(path) + pb = _assemble(records) + pb = _auto_inject_start_end(pb) + from aiperf.dataset.graph.auto_derive import auto_derive + from aiperf.dataset.graph.native_lowering import lower_native_to_unified + + return lower_native_to_unified(auto_derive(pb)) + + +def _parse_via_adapter( + path: Path, + fmt: WorkloadFormat, + ctx: GraphParseContext | None = None, + **adapter_kwargs: Any, +) -> ParsedGraph: + from aiperf.plugin import plugins + from aiperf.plugin.enums import PluginType + + try: + adapter_cls = plugins.get_class(PluginType.GRAPH_ADAPTER, fmt) + except Exception as e: + raise GraphParseError(f"{path}: unknown format {fmt!r}") from e + try: + return adapter_cls.parse(path, ctx, **adapter_kwargs) + except ValueError as e: + # Adapters raise their own ValueError subclasses. Re-wrap so callers + # only need to catch GraphParseError. A TypeError from an + # adapter_kwarg the adapter's parse does not accept is NOT caught here: + # it propagates as the documented fail-loud (a build-plane wiring bug, + # not a workload-file error). + raise GraphParseError(str(e)) from e + + +def _read_jsonl(path: Path) -> list[dict[str, Any]]: + out: list[dict[str, Any]] = [] + with path.open("rb") as f: + for lineno, raw in enumerate(f, start=1): + stripped = raw.strip() + if not stripped: + continue + try: + rec = orjson.loads(stripped) + except orjson.JSONDecodeError as e: + raise GraphParseError( + f"{path}: line {lineno}: invalid JSON: {e}" + ) from e + if not isinstance(rec, dict): + raise GraphParseError( + f"{path}: line {lineno}: record must be a JSON object" + ) + out.append(rec) + return out + + +def _read_yaml(path: Path) -> list[dict[str, Any]]: + import yaml + + try: + docs = list(yaml.safe_load_all(path.read_text())) + except yaml.YAMLError as e: + raise GraphParseError(f"{path}: invalid YAML: {e}") from e + docs = [d for d in docs if d is not None] + if len(docs) > 1 or (len(docs) == 1 and _looks_like_multi_doc(docs[0])): + return [_normalize_yaml_record(d, path) for d in docs] + if not docs: + return [] + return _expand_single_doc(docs[0], path) + + +def _looks_like_multi_doc(doc: Any) -> bool: + return isinstance(doc, dict) and "kind" in doc + + +def _normalize_yaml_record(doc: Any, path: Path) -> dict[str, Any]: + if not isinstance(doc, dict): + raise GraphParseError(f"{path}: YAML document must be a mapping") + return doc + + +_SINGLE_DOC_KEYS = ("graph", "traces") + + +def _expand_single_doc(doc: Any, path: Path) -> list[dict[str, Any]]: + if not isinstance(doc, dict): + raise GraphParseError(f"{path}: top-level YAML must be a mapping") + unknown = [k for k in doc if k not in _SINGLE_DOC_KEYS] + if unknown: + import difflib + + hints: list[str] = [] + for key in unknown: + close = difflib.get_close_matches(str(key), _SINGLE_DOC_KEYS, n=1) + hints.append( + f"{key!r} (did you mean {close[0]!r}?)" if close else f"{key!r}" + ) + raise GraphParseError( + f"{path}: unknown top-level key(s) {', '.join(hints)}; a " + f"single-document graph workload may only contain " + f"{list(_SINGLE_DOC_KEYS)} (node/edge topology goes under 'graph:')" + ) + if not doc: + raise GraphParseError( + f"{path}: top-level mapping contains none of {list(_SINGLE_DOC_KEYS)}; " + f"author the workload as a 'graph:' block plus a 'traces:' list" + ) + out: list[dict[str, Any]] = [] + if "graph" in doc: + graph_body = doc["graph"] or {} + if not isinstance(graph_body, dict): + raise GraphParseError(f"{path}: 'graph' must be a mapping") + out.append({"kind": "graph", **graph_body}) + if "traces" in doc: + traces = doc["traces"] or [] + if not isinstance(traces, list): + raise GraphParseError(f"{path}: 'traces' must be a list") + for t in traces: + if not isinstance(t, dict): + raise GraphParseError(f"{path}: each trace must be a mapping") + out.append({"kind": "trace", **t}) + return out + + +def _auto_inject_start_end(pb: ParsedGraph) -> ParsedGraph: + """Inject ``START -> `` and `` -> END`` edges when missing. + + Mirrors the foreign-adapter read paths' sentinel-edge auto-injection so a + native workload that declares only inter-node edges produces an equivalent + :class:`ParsedGraph` shape. Behavior-preserving: existing explicit + START/END edges are left untouched. + """ + if not pb.graph.nodes: + return pb + existing_sources: set[str] = set() + existing_targets: set[str] = set() + for e in pb.graph.edges: + if isinstance(e, StaticEdge): + existing_sources.add(e.source) + existing_targets.add(e.target) + has_start = any( + isinstance(e, StaticEdge) and e.source == START_NODE_ID for e in pb.graph.edges + ) + has_end = any( + isinstance(e, StaticEdge) and e.target == END_NODE_ID for e in pb.graph.edges + ) + if has_start and has_end: + return pb + new_edges: list[Any] = list(pb.graph.edges) + if not has_start: + roots = [nid for nid in pb.graph.nodes if nid not in existing_targets] + for r in roots: + new_edges.insert(0, StaticEdge(source=START_NODE_ID, target=r)) + if not has_end: + leaves = [nid for nid in pb.graph.nodes if nid not in existing_sources] + for leaf in leaves: + new_edges.append(StaticEdge(source=leaf, target=END_NODE_ID)) + if len(new_edges) == len(pb.graph.edges): + return pb + return msgspec.structs.replace( + pb, graph=msgspec.structs.replace(pb.graph, edges=new_edges) + ) + + +def _assemble(records: list[dict[str, Any]]) -> ParsedGraph: + graph: GraphRecord | None = None + traces: list[TraceRecord] = [] + any_trace_seen = False + for idx, rec in enumerate(records): + kind = rec.get("kind", "trace") + body = {k: v for k, v in rec.items() if k != "kind"} + if kind == "graph": + if any_trace_seen: + raise GraphParseError( + f"record {idx}: graph record must precede trace records (rule-19)" + ) + if graph is not None: + raise GraphParseError( + f"record {idx}: more than one 'kind: graph' record (rule-20)" + ) + graph = decode_graph(body) + elif kind == "trace": + any_trace_seen = True + try: + traces.append(msgspec.convert(body, type=TraceRecord)) + except msgspec.ValidationError as e: + raise GraphParseError(f"record {idx}: trace decode failed: {e}") from e + else: + raise GraphParseError(f"record {idx}: unknown kind {kind!r}") + return ParsedGraph( + graph=graph if graph is not None else GraphRecord(), + traces=traces, + ) + + +# --------------------------------------------------------------------------- +# Workload-format detection +# +# Walks the `graph_adapter` plugin registry and asks each adapter's +# `can_load(path)` method whether it recognizes the file. Ties between +# multiple matches are broken by `detection_priority` from each plugin's +# metadata (higher wins). Folded in from the adapters package because `parser` +# is its only importer (no adapter ever consumed it). +# --------------------------------------------------------------------------- + +_logger = AIPerfLogger(__name__) + +WorkloadFormat: TypeAlias = str +"""String alias for graph adapter names (e.g. "weka_trace", "native"). + +The dynamic `GraphAdapterType` enum is intentionally not referenced here: the +`graph_adapter` plugin category is not registered on this base, and the IR +ingestion detection path only needs the string alias at import time.""" + + +class WorkloadFormatError(ValueError): + """Raised when a file's workload format cannot be determined.""" + + +def detect_format(path: str | Path) -> WorkloadFormat: + """Return the workload format of `path`, walking the graph_adapter registry. + + Each registered adapter's `can_load(path)` is queried; ties between + multiple matches are broken by `detection_priority` from the plugin + metadata (higher wins). Raises `WorkloadFormatError` if no adapter + recognizes the file. + """ + from aiperf.plugin import plugins + from aiperf.plugin.enums import PluginType + from aiperf.plugin.schema.schemas import GraphAdapterMetadata + + p = Path(path) + matches: list[tuple[int, str]] = [] + for entry, cls in plugins.iter_all(PluginType.GRAPH_ADAPTER): + try: + claimed = cls.can_load(p) + except Exception as e: + # One adapter's sniff crashing on corrupt candidate bytes (or an + # adapter bug) must not abort detection for every other adapter; + # treat it as "does not claim". Real errors in the winning + # adapter still surface from its parse(). + _logger.debug( + f"graph adapter {entry.name!r} can_load({p}) raised {e!r}; " + "treating as does-not-claim" + ) + continue + if claimed: + meta = entry.get_typed_metadata(GraphAdapterMetadata) + matches.append((meta.detection_priority, entry.name)) + if not matches: + raise WorkloadFormatError( + f"{p}: no graph adapter recognizes this file. " + f"Registered adapters: " + f"{[e.name for e in plugins.iter_entries(PluginType.GRAPH_ADAPTER)]}" + ) + # max() picks the highest priority; ties resolved by lexicographic name + # which is fine for tie-equal-priority cases since priorities should be + # distinct for adapters whose can_load might overlap. + return max(matches)[1] diff --git a/src/aiperf/dataset/graph/segment_ir/__init__.py b/src/aiperf/dataset/graph/segment_ir/__init__.py new file mode 100644 index 0000000000..16cdaddec6 --- /dev/null +++ b/src/aiperf/dataset/graph/segment_ir/__init__.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Format-agnostic segment-trie IR: the content-addressed prefix-segment store +and its interned builder + ordinal scheme. + +Any adapter can target this IR by returning a ``ParsedGraph`` whose dispatchable +``LlmNode``s carry ``metadata["trie"]["prompt_segment_ids"]`` (an ordered path +into a ``SegmentPool``) with ``ParsedGraph.segment_pool`` set. Everything here is +independent of any specific recorded trace format (weka, dynamo, ...). +""" diff --git a/src/aiperf/dataset/graph/segment_ir/envelope.py b/src/aiperf/dataset/graph/segment_ir/envelope.py new file mode 100644 index 0000000000..a747ef9e24 --- /dev/null +++ b/src/aiperf/dataset/graph/segment_ir/envelope.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The segment-trie IR node envelope contract -- the adapter-facing seam. + +An adapter opts a dispatchable ``LlmNode`` into the segment-trie IR by stamping +``metadata["trie"]["prompt_segment_ids"]`` (an ordered path into the run's +``SegmentPool``) via :func:`stamp_prompt_segment_ids`. The build plane reads that +path back with :func:`read_prompt_segment_ids`. + +The build-plane envelope (see ``store_builder._trie_envelope``) re-derives the +wire-body overrides / ``stream`` from the node's own fields, and the sidecar +strips ``metadata["trie"]`` contents to ``{}`` (keeping the key as a routing +marker). The trace adapters (weka, dynamo) stamp only the path; the native and +dag_jsonl lowerings additionally stamp ``assembly``/``capture`` extras for +dynamic-slot nodes, which the eager interned drain resolves and persists into +the manifest (store-addressed ``items`` + ``capture``). +""" + +from __future__ import annotations + +from typing import Any + +import msgspec + +from aiperf.dataset.graph.models import LlmNode + +# The metadata key under which the segment-trie envelope lives. Internal to this +# module -- other modules do NOT import this constant; they key off the literal +# ``"trie"`` string directly (``graph_ir_replay._is_trie_graph`` routes on it, the +# graph_meta sidecar strips-but-keeps it). The string VALUE is the stable +# cross-module contract; the name is not part of the public surface. +TRIE_META_KEY = "trie" + + +def stamp_prompt_segment_ids( + node: LlmNode, prompt_segment_ids: list[str], *, extra: dict[str, Any] | None = None +) -> LlmNode: + """Return ``node`` with ``metadata["trie"]`` carrying ``prompt_segment_ids``. + + ``prompt_segment_ids`` is the ordered ``SegmentPool`` path the worker walks to + materialize the prompt. ``extra`` supplies companion keys (the native and + dag_jsonl lowerings' ``assembly``/``capture`` dynamic-slot keys, which the + eager interned drain persists), merged after ``prompt_segment_ids`` so + the ``metadata["trie"]`` dict has a deterministic key order. Any pre-existing + ``node.metadata`` is preserved. + ``LlmNode`` is frozen, so this returns a replaced copy. + """ + trie: dict[str, Any] = {"prompt_segment_ids": prompt_segment_ids} + if extra: + trie.update(extra) + new_meta = {**(node.metadata or {}), TRIE_META_KEY: trie} + return msgspec.structs.replace(node, metadata=new_meta) + + +def read_prompt_segment_ids(node: LlmNode) -> list[str] | None: + """Read a node's ``metadata["trie"]["prompt_segment_ids"]`` path, or ``None``. + + ``None`` means the node is not part of the segment-trie IR (no envelope, or a + malformed one) -- the build plane skips it (mints no manifest). + """ + trie_meta = (node.metadata or {}).get(TRIE_META_KEY) + if not isinstance(trie_meta, dict): + return None + path = trie_meta.get("prompt_segment_ids") + return path if isinstance(path, list) else None + + +__all__ = ["read_prompt_segment_ids", "stamp_prompt_segment_ids"] diff --git a/src/aiperf/dataset/graph/segment_ir/interval_order.py b/src/aiperf/dataset/graph/segment_ir/interval_order.py new file mode 100644 index 0000000000..d877cf9530 --- /dev/null +++ b/src/aiperf/dataset/graph/segment_ir/interval_order.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Interval-order dependency-edge derivation for the segment-trie IR. + +Format-agnostic: given a set of nodes that each expose a recorded interval +(``raw_start`` / ``raw_end``), a warped ``start`` / ``end`` (the clock the +runtime replays), a time-consistent ``rank``, and a set of enclosing async +subtree ids, derive each node's incoming ``StaticEdge``s by the finished-before + +transitive-reduction frontier rule. Any adapter whose nodes expose that duck-typed +surface -- a ``node_id`` str, an int ``rank``, float ``start`` / ``end`` (the +idle-gap-warped clock the runtime replays), float ``raw_start`` / ``raw_end`` (the +RAW recorded who-finished-before-whom interval), and a ``frozenset[str]`` +``async_ancestors`` (enclosing fire-and-forget subtree ids) -- can reuse this +(weka is the first caller); the rule reads only those attributes, never a +recorded-trace type. +""" + +from __future__ import annotations + +from typing import Any + +from aiperf.common.constants import MICROS_PER_SECOND +from aiperf.dataset.graph.models import START_NODE_ID, StaticEdge + + +def compute_ranks(nodes: list) -> None: + """Stamp each node's time-consistent global ``rank``. + + The rank is the node's index in the total order sorted by + ``(start, end, node_id)`` -- a linear extension of raw finished-before + (idle-warp monotonicity), so the interval-order edge rule is a strict partial + order and its transitive reduction is always a DAG. Mutates ``node.rank``. + """ + for i, node in enumerate(sorted(nodes, key=lambda n: (n.start, n.end, n.node_id))): + node.rank = i + + +def _excluded_async(cand: Any, target: Any) -> bool: + """``True`` when ``cand`` sits under an async boundary ``target`` does not share. + + A fire-and-forget (``async_launched``) subtree never AND-joins the scope that + launched it: a candidate whose enclosing async-subtree ids are NOT a subset of + the target's is excluded from the target's predecessor set before the frontier + filter runs. + """ + return not cand.async_ancestors <= target.async_ancestors + + +def build_interval_edges(nodes: list) -> dict[str, list[StaticEdge]]: + """Derive every node's incoming interval-order edges globally. + + Semantics: + + * ``A -> B`` iff ``A`` finished before ``B`` started (RAW clock, + ``A.raw_end <= B.raw_start``) AND ``rank(A) < rank(B)``. + * Async exclusion (:func:`_excluded_async`) drops fire-and-forget children + from the candidate set BEFORE the frontier filter. + * Frontier (transitive reduction): keep only the MAXIMAL finished-before + candidates -- drop ``c`` when another candidate ``d`` has ``c`` finished + before ``d`` AND the covering edge ``c -> d`` actually exists, i.e. ``d`` + does not async-exclude ``c`` (``c -> d -> node``, so ``c`` is transitively + covered). Without the exclusion check, a main-chain ``d`` outside ``c``'s + async subtree would drop ``c`` while carrying no ``c -> d`` edge itself, + losing the recorded ``c``-before-``node`` ordering inside the subtree. + * Binding-cause delay: the latest-ending frontier predecessor (``max by + .end``) carries the warped end-to-start gap; every other frontier + predecessor is an AND-join wait (delay ``0.0``). + * Empty frontier: the node roots at ``START`` at its own warped arrival offset. + + Complexity: the per-node frontier filter is ``O(|candidates|^2)`` worst case, + so total is ``Theta(n^2)`` per node up to ``Theta(n^3)`` for a pathological + wide fan-in (rare). A dropped ``c`` is always transitively covered: the + dominating ``d`` carries a real ``c -> d`` edge (by induction on ``d``'s own + candidate set, which includes ``c`` because ``d`` does not async-exclude it). + """ + by_rank = sorted(nodes, key=lambda n: n.rank) + out: dict[str, list[StaticEdge]] = {} + for node in nodes: + cands = [ + c + for c in by_rank + if c is not node + and c.rank < node.rank + and c.raw_end <= node.raw_start + and not _excluded_async(c, node) + ] + if not cands: + out[node.node_id] = [ + StaticEdge( + source=START_NODE_ID, + target=node.node_id, + min_start_delay_us=node.start * MICROS_PER_SECOND, + ) + ] + continue + # Frontier = maximal finished-before candidates (drop c if the edge + # c -> d exists for some later-ranked d in the candidate set). The + # edge exists only when d does not async-exclude c -- a main-chain d + # outside c's async subtree carries no c -> d edge, so it cannot cover c. + frontier = [ + c + for i, c in enumerate(cands) + if not any( + c.raw_end <= d.raw_start and not _excluded_async(c, d) + for d in cands[i + 1 :] + ) + ] + binding = max(frontier, key=lambda c: c.end) + out[node.node_id] = [ + StaticEdge( + source=c.node_id, + target=node.node_id, + delay_after_predecessor_us=( + max(0.0, node.start - c.end) * MICROS_PER_SECOND + if c is binding + else 0.0 + ), + ) + for c in frontier + ] + return out + + +def apply_start_anchors( + nodes: list, edges_by_node: dict[str, list[StaticEdge]] +) -> None: + """Replace an overlapped node's incoming edges with one start-anchored edge. + + For each node whose ``causal_parent_id`` names another node in the set + and whose recorded start falls INSIDE that parent's recorded interval + (``parent.raw_start <= node.raw_start < parent.raw_end``), the + interval-order edges are replaced with a single + ``StaticEdge(parent -> node, delay_after_predecessor_start_us=D)`` + where ``D`` is the warped start-to-start gap. The runtime schedules + such a node at its parent's DISPATCH and gates it ``D`` later, so + recorded mid-flight concurrency (subagent spawns, aux tool calls) + tracks the parent causally instead of freezing to the recorded wall + clock. Nodes whose causal parent had already finished keep their + interval-order edges -- end-anchoring is correct there by construction. + + When the parent is a streaming request (``parent.request.ttft`` set) and the + child was recorded at/after that first token (``node.raw_start - + parent.raw_start >= ttft``), the edge additionally carries + ``delay_after_predecessor_first_token_us = D - ttft*1e6`` so the runtime can + re-anchor onto the parent's OBSERVED first token, falling back to + dispatch + ``D`` when the parent terminates without streaming one. A + pre-TTFT child (started before the recorded first token) or a non-streaming + parent (``ttft is None``) keeps the pure dispatch anchor. + """ + by_id = {n.node_id: n for n in nodes} + for node in nodes: + pid = getattr(node, "causal_parent_id", None) + if pid is None: + continue + parent = by_id.get(pid) + if parent is None or parent is node: + continue + if not (parent.raw_start <= node.raw_start < parent.raw_end): + continue + delay_us = max(0.0, node.start - parent.start) * MICROS_PER_SECOND + ttft_s = getattr(parent.request, "ttft", None) + first_token_delay_us = None + if ttft_s is not None and (node.raw_start - parent.raw_start) >= ttft_s: + first_token_delay_us = max(0.0, delay_us - ttft_s * MICROS_PER_SECOND) + edges_by_node[node.node_id] = [ + StaticEdge( + source=parent.node_id, + target=node.node_id, + delay_after_predecessor_start_us=delay_us, + delay_after_predecessor_first_token_us=first_token_delay_us, + ) + ] + + +__all__ = [ + "apply_start_anchors", + "build_interval_edges", + "compute_ranks", +] diff --git a/src/aiperf/dataset/graph/segment_ir/pool.py b/src/aiperf/dataset/graph/segment_ir/pool.py new file mode 100644 index 0000000000..f786537334 --- /dev/null +++ b/src/aiperf/dataset/graph/segment_ir/pool.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Content-addressed segment pool primitives (neutral stdlib-only leaf). + +These dataclasses live in a dependency-free leaf module so ``models.py`` can +type ``ParsedGraph.segment_pool`` as ``SegmentPool | None`` without pulling in +``adapters/__init__`` (which eagerly imports ``weka_trace`` -> ``models`` and +would form a partial-init import cycle). +""" + +from __future__ import annotations + +import hashlib +from array import array +from dataclasses import dataclass, field +from typing import Any + +import orjson + + +@dataclass(slots=True, frozen=True) +class Segment: + id: str + """Content-addressed, prefix-dependent id (opaque to consumers).""" + role: str + """Message role (e.g. ``"user"``); ``message.get("role", "")`` for raw segments.""" + content: str + """Message content; ``""`` for raw segments (the wire blob is authoritative).""" + parent_id: str | None + """Prefix segment id this one extends, or ``None`` at the path root.""" + wire_json: str | None = None + """Verbatim ``orjson.dumps(message)`` for a raw-authored segment (key order and + extra keys preserved). ``None`` for a role/content segment, whose blob is derived + as ``{"role", "content"}`` at persist time (the existing normalized behavior).""" + + +def segment_id(parent_id: str | None, role: str, tokens: list[int]) -> str: + """Content-addressed, prefix-dependent id for one role/content segment. + + The id is OPAQUE (consumers treat it as a name, not a value): only its + determinism, prefix-dependence (``parent_id`` framing), role-framing, and + dedup-consistency matter. The token list is hashed via a bulk C-level + ``array("q", tokens).tobytes()`` int image -- a deterministic, injective + byte encoding of the int sequence -- instead of the per-token + ``str().encode()`` join that dominated the cold-build profile (~44M + per-token calls on a corpus-scale trace). The ``\\x00`` delimiters keep the + ``parent_id`` / ``role`` / tokens fields framed so distinct fields cannot + alias, and the fixed-width 8-byte little-endian image means two distinct + token lists cannot collide via concatenation. + """ + h = hashlib.blake2b(digest_size=16) + h.update((parent_id or "").encode()) + h.update(b"\x00") + h.update(role.encode()) + h.update(b"\x00") + # token ids are the canonical, tokenization-1:1 content key; the bulk + # fixed-width int image is a deterministic injective byte encoding. + h.update(array("q", tokens).tobytes()) + return h.hexdigest() + + +def text_segment_id(parent_id: str | None, role: str, content: str) -> str: + """Content-addressed id for a text-authored segment (no tokenizer needed). + + Same ``parent_id`` / ``role`` framing as :func:`segment_id`, with a distinct + domain tag before the payload so a text-derived id can never alias a + token-derived id for the same prefix and role. Ids are opaque to every + consumer; only determinism and dedup-consistency matter (the native + lowering has no tokenizer at parse time, so it keys on UTF-8 content + bytes instead of token ids). + """ + h = hashlib.blake2b(digest_size=16) + h.update((parent_id or "").encode()) + h.update(b"\x00") + h.update(role.encode()) + h.update(b"\x00text\x00") + h.update(content.encode()) + return h.hexdigest() + + +def raw_segment_id(parent_id: str | None, wire_json: str) -> str: + """Content-addressed id for a raw-wire-JSON segment (verbatim authored message). + + Same ``parent_id`` framing as :func:`text_segment_id`, with a distinct ``raw`` + domain tag before the payload so a raw-derived id can never alias a text- or + token-derived id for the same prefix. The payload is the verbatim + ``orjson.dumps(message)`` bytes -- so two messages differing only in key order + or an extra key hash to distinct ids, preserving the byte-verbatim contract. + Ids are opaque; only determinism, prefix-dependence, and dedup-consistency matter. + """ + h = hashlib.blake2b(digest_size=16) + h.update((parent_id or "").encode()) + h.update(b"\x00raw\x00") + h.update(wire_json.encode()) + return h.hexdigest() + + +@dataclass(slots=True) +class SegmentPool: + _by_id: dict[str, Segment] = field(default_factory=dict) + + @property + def by_id(self) -> dict[str, Segment]: + """Content-addressed id -> Segment map (read as-is; do not mutate).""" + return self._by_id + + def add( + self, + *, + role: str, + content: str, + tokens: list[int], + parent_id: str | None, + ) -> str: + sid = segment_id(parent_id, role, tokens) + if sid not in self._by_id: + self._by_id[sid] = Segment( + id=sid, role=role, content=content, parent_id=parent_id + ) + return sid + + def add_text( + self, + *, + role: str, + content: str, + parent_id: str | None, + ) -> str: + sid = text_segment_id(parent_id, role, content) + if sid not in self._by_id: + self._by_id[sid] = Segment( + id=sid, role=role, content=content, parent_id=parent_id + ) + return sid + + def add_raw_message(self, *, message: dict[str, Any], parent_id: str | None) -> str: + """Intern a raw-authored message verbatim (key order and extra keys kept). + + The message is serialized ONCE via ``orjson.dumps`` and that blob is the + segment's authoritative wire form: :func:`raw_segment_id` keys on it, so + two messages differing only in key order or an extra key dedup distinctly. + ``role`` records ``message.get("role", "")`` for envelope framing and + ``content`` is ``""`` (the wire blob, not the derived dict, is persisted). + """ + wire_json = orjson.dumps(message).decode() + sid = raw_segment_id(parent_id, wire_json) + if sid not in self._by_id: + self._by_id[sid] = Segment( + id=sid, + role=message.get("role", ""), + content="", + parent_id=parent_id, + wire_json=wire_json, + ) + return sid + + def get(self, sid: str) -> Segment: + return self._by_id[sid] + + def materialize(self, path_ids: list[str]) -> list[dict]: + out: list[dict] = [] + for i in path_ids: + s = self._by_id[i] + if s.wire_json is not None: + out.append(orjson.loads(s.wire_json)) + else: + out.append({"role": s.role, "content": s.content}) + return out diff --git a/src/aiperf/dataset/graph/segment_ir/prefix_cache.py b/src/aiperf/dataset/graph/segment_ir/prefix_cache.py new file mode 100644 index 0000000000..71d4e85a7d --- /dev/null +++ b/src/aiperf/dataset/graph/segment_ir/prefix_cache.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared theoretical prefix-cache stamping for the segment-trie IR. + +``hash_id_scope: "local"`` means one hash namespace per trace file, so a block +first sent by ANY conversation of a trace (root, subagent child, or flat chain) +is a cache hit when any other conversation of the same trace re-sends it. The +counts are computed over ONE shared per-trace seen-set consumed in recorded +global time order, then stamped per node onto the NATIVE +``LlmNode.theoretical_prefix_cache_hit_blocks`` / ``_total_blocks`` fields +(``Turn`` naming) where the ``theoretical_prefix_cache`` results accumulator +reads them back (:func:`extract_prefix_cache_by_node` via the dataset manager). + +Adapter-agnostic: any adapter that lowers through the trie IR (weka, dynamo) +gets the stamping by calling :func:`stamp_theoretical_prefix_cache` on its +assembled node map. Nodes whose requests carry no hash blocks are left +unstamped (the accumulator treats a zero total as absent). + +The seen-set stays per-trace even for ``hash_id_scope: "global"`` corpora: +recorded ``t`` is conversation-relative, so no cross-file clock exists to order +one trace's sends against another's. Under global scope the stamped hits are +therefore a LOWER BOUND -- a block another trace already sent still counts as a +miss here even though the server (infinite cache) would hit it. +""" + +from __future__ import annotations + +from collections.abc import Iterable, MutableMapping + +import msgspec + +from aiperf.dataset.graph.models import GraphRecord, LlmNode +from aiperf.dataset.graph.segment_ir.trie_content import TrieNode + + +def compute_shared_prefix_cache_counts( + trie_nodes: Iterable[TrieNode], +) -> dict[str, tuple[int, int]]: + """``{node_id: (hit_blocks, total_blocks)}`` over one shared seen-set. + + Nodes are consumed in recorded global time order (``request.t``, with the + flattened recorded-order index as the deterministic tiebreak — the same + ordering the interval-order rank uses). ``hit_blocks`` is the LEADING run of + the node's ``hash_ids`` already in the cache (stop at the first miss); + ``total_blocks`` is the full hash-id count. Every block of every request + enters the infinite cache regardless of hit position. + """ + out: dict[str, tuple[int, int]] = {} + seen: set[int] = set() + for node in sorted(trie_nodes, key=lambda n: (n.request.t, n.order)): + hashes = node.request.hash_ids + hits = 0 + for hid in hashes: + if hid not in seen: + break + hits += 1 + out[node.node_id] = (hits, len(hashes)) + seen.update(hashes) + return out + + +def stamp_theoretical_prefix_cache( + llm_nodes: MutableMapping[str, LlmNode], + trie_nodes: Iterable[TrieNode], +) -> None: + """Stamp per-node counts onto the native ``LlmNode`` fields in ``llm_nodes``. + + ``LlmNode`` is frozen, so each stamped node is REPLACED in the mapping + (``msgspec.structs.replace``) with the counts on + ``theoretical_prefix_cache_hit_blocks`` / ``_total_blocks``. Nodes with + zero hash blocks are skipped; trie nodes without an assembled LlmNode + (never the case today) are ignored. + """ + for node_id, (hit_blocks, total_blocks) in compute_shared_prefix_cache_counts( + trie_nodes + ).items(): + if total_blocks <= 0: + continue + llm = llm_nodes.get(node_id) + if llm is None: + continue + llm_nodes[node_id] = msgspec.structs.replace( + llm, + theoretical_prefix_cache_hit_blocks=hit_blocks, + theoretical_prefix_cache_total_blocks=total_blocks, + ) + + +def extract_prefix_cache_by_node( + top_graph: GraphRecord, +) -> dict[str, list[int]]: + """Collect every stamped node's prefix-cache counts as ``{node_id: [hit, total]}``. + + Reads the native ``theoretical_prefix_cache_hit_blocks`` / ``_total_blocks`` + fields off the graph. The key is the node id the graph dispatch path + reports on the wire, which is what the ``theoretical_prefix_cache`` + accumulator recovers from each record's ``x_correlation_id``. Returns ``{}`` + when no node was stamped (a non-trie graph, or a trace whose requests + carried no hash blocks). + """ + out: dict[str, list[int]] = {} + for node_id, node in top_graph.nodes.items(): + if not isinstance(node, LlmNode): + continue + if ( + node.theoretical_prefix_cache_hit_blocks is None + or node.theoretical_prefix_cache_total_blocks is None + ): + continue + out[node_id] = [ + node.theoretical_prefix_cache_hit_blocks, + node.theoretical_prefix_cache_total_blocks, + ] + return out + + +__all__ = [ + "compute_shared_prefix_cache_counts", + "extract_prefix_cache_by_node", + "stamp_theoretical_prefix_cache", +] diff --git a/src/aiperf/dataset/graph/segment_ir/store_builder.py b/src/aiperf/dataset/graph/segment_ir/store_builder.py new file mode 100644 index 0000000000..5d5a1f95a6 --- /dev/null +++ b/src/aiperf/dataset/graph/segment_ir/store_builder.py @@ -0,0 +1,521 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Build-time persistence for the segment-trie IR. + +The segment-trie IR (produced by any adapter, e.g. +:func:`aiperf.dataset.graph.adapters.weka.trie_build.build_trie_graph` or +the dynamo build-replay path) realizes a trace as a flat ``LlmNode`` + +``StaticEdge`` graph whose nodes each carry +``metadata["trie"]["prompt_segment_ids"]`` -- a path into the content-addressed +:class:`~aiperf.dataset.graph.segment_ir.pool.SegmentPool`. + +The trie graph persists as ONE :class:`GraphSegmentUnifiedBackingStore` -- a +content pool (every pool ``Segment`` written as ``id -> {role, content}``) +plus a per-node manifest region carrying each node's interned int-handle path, +``dispatch_overrides``, and ``stream``: + +* :func:`build_unified_trie_store_interned` drains a whole-graph parse into the + unified store eagerly. It is the in-process drain for EVERY non-weka format + (dynamo / native / dag_jsonl), and the byte-parity oracle for the streaming + split below. +* :func:`iter_trace_segment_payloads` + :func:`build_unified_trie_store_from_payloads` + are the streaming split for weka's worker-pool build (corpus-scale HF and + local weka sources): per-row workers emit :class:`TraceSegmentPayload`\\ s and + the parent drains them into the SAME unified store shape, so the worker opens + one reader either way. + +Ordinals are assigned densely over the trie graph's ``LlmNode``s in their +recorded ``arrival_offset_us`` order (ties broken by node id), which is the +node-creation order :func:`build_trie_graph` emits -- a stable, deterministic +contract both the build plane (here) and the schedule plane resolve from the +same parsed graph. + +Schedule-plane agreement: :func:`graph_path_catalog._catalog_for_trace` builds +the per-trace catalog by reusing THIS module's :func:`flat_trie_ordinals` over +the trace's flat +``LlmNode``s, so the live ``GraphIRReplayStrategy`` dispatches every trie node at +the SAME ordinal this builder wrote its envelope at. The build->persist-> +worker-materialize half (byte-correct prompts on dispatch) and the schedule half +(catalog/strategy enumerate + dispatch trie ``LlmNode``s) now share one ordinal +contract, so a full ``aiperf profile`` run dispatches the trie graph end-to-end. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Iterator +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +import orjson + +from aiperf.dataset.graph.models import LlmNode, ParsedGraph, TraceRecord +from aiperf.dataset.graph.segment_ir.envelope import read_prompt_segment_ids +from aiperf.dataset.graph_segment_unified_store import NodeEnvelope + +if TYPE_CHECKING: + from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + ) + +_PROFILING_VARIANT = "profiling" + + +@dataclass(slots=True) +class TraceSegmentPayload: + """Worker-serialized trie artifacts for ONE trace, shipped across the pool. + + A trie row crosses the worker boundary as content-addressed segments + + ``prompt_segment_ids`` envelopes. This is the parallel payload the per-row + worker emits when ``pg.segment_pool is not None`` -- the consumer + (:func:`build_unified_trie_store_from_payloads`) drains ``segments`` into + the unified content pool (idempotent dedup on the content-addressed id + bounds RAM) and writes ``envelopes`` as per-node manifests, mirroring the + eager :func:`build_unified_trie_store_interned` output. All fields are + plain (``str`` / ``int`` / ``bytes`` / ``tuple``) so the dataclass pickles + cleanly through ``multiprocessing.Pool``. + """ + + trace_id: str + """Trace identifier.""" + + node_ordinals: dict[str, int] + """Node id -> dense ordinal map for the trace (the addressing catalog).""" + + envelopes: list[NodeEnvelope] + """Per-node ``prompt_segment_ids`` envelopes (profiling variant).""" + + segments: list[tuple[str, str, str, str | None]] = field(default_factory=list) + """The row's pool segments as ``(id, role, content, wire_json)`` tuples to ``put``. + + ``wire_json`` is the verbatim ``orjson.dumps(message)`` for a raw-authored dag + segment (persisted byte-for-byte) or ``None`` for a role/content segment (the + store derives the ``{"role", "content"}`` blob). Plain (``str``/``None``) so the + dataclass still pickles cleanly through ``multiprocessing.Pool``.""" + + structural_graph: bytes = b"" + """Content-free structural ``ParsedGraph`` (msgpack) for this row: topology + + timing with ``replay_outputs`` emptied and the segment pool emptied (but kept + non-None so the loaded graph still reads as a segment-store IR). The streaming consumer + merges these across rows into the corpus structural graph and writes the + ``graph_meta`` sidecar, so the TimingManager loads it instead of the + whole-corpus eager re-parse (which balloons to tens of GB). Empty (``b""``) + for non-first payloads of a multi-trace parse (attached once, like + ``segments``).""" + + +def trie_node_ordinals(trace_graph_nodes: dict[str, LlmNode]) -> dict[str, int]: + """Assign a dense ordinal to each trie ``LlmNode`` in deterministic order. + + The order is the node's recorded ``arrival_offset_us`` (its warped recorded + ``t``), ties broken by node id -- identical to the recorded-order list + :func:`build_trie_graph` walks, so the build-plane ordinal here matches the + schedule-plane ordinal resolved from the same parsed graph. Ordinals are dense + (``0..N-1``) and unique within the trace. + """ + ordered = sorted( + trace_graph_nodes.items(), + key=lambda kv: (kv[1].arrival_offset_us or 0, kv[0]), + ) + return {node_id: ordinal for ordinal, (node_id, _) in enumerate(ordered)} + + +def flat_trie_ordinals(parsed: ParsedGraph, trace: TraceRecord) -> dict[str, int]: + """THE shared build-plane/schedule-plane ordinal scheme. + + Both :func:`build_unified_trie_store_interned` (build) and + ``graph_path_catalog._catalog_for_trace`` (schedule) call this so a node's + build-time manifest ordinal equals its dispatch-time catalog ordinal. Every + live producer emits a flat LlmNode graph, so this is + :func:`trie_node_ordinals` over the trace's :func:`_trie_llm_nodes` keyed by + bare node id. + """ + return trie_node_ordinals(_trie_llm_nodes(parsed, trace)) + + +def _trie_llm_nodes(parsed: ParsedGraph, trace: TraceRecord) -> dict[str, LlmNode]: + """The trie graph's ``LlmNode``s for ``trace`` (the trivial IR has only these). + + The trie builder emits a single top graph with ALL nodes ``LlmNode`` and no + subgraphs; resolve through the trace's graph_ref so a multi-trace merge still + addresses the right topology. + """ + from aiperf.dataset.graph.models import resolve_trace_graph + + top = resolve_trace_graph(parsed, trace) + return {nid: n for nid, n in top.nodes.items() if isinstance(n, LlmNode)} + + +def _prompt_segment_ids(node: LlmNode) -> list[str] | None: + """Read a trie node's ``prompt_segment_ids`` path, or ``None`` if absent.""" + return read_prompt_segment_ids(node) + + +def _trie_envelope(node: LlmNode, prompt_segment_ids: list[str]) -> dict: + """Compose the per-node trie envelope the worker materializes from. + + Carries ``prompt_segment_ids`` (the segment-pool path the worker walks), + a ``dispatch_overrides`` wire-body dict assembled from the node's + ``extra_body`` (vendor keys) with the native ``LlmNode.model`` / + ``max_tokens`` / ``raw_tools`` FOLDED IN (``model`` and ``tools`` verbatim; + the cap as the worker-mapped ``max_output_tokens`` entry; a hand-authored + ``extra_body`` entry wins over any fold), the ``stream`` flag (from the + native ``LlmNode.streaming``), and -- when the adapter stamped per-node + HTTP headers on the native ``LlmNode.extra_headers`` field + (dynamo session identity) -- an ``extra_headers`` map the worker attaches + to the request HEADERS, never the body. The trie path is self-contained: + the worker walks the pool path directly, no ancestor accumulation. + + When the adapter stamped ``metadata["dispatch"]["endpoint_extra_applied"] = + True`` (it already folded the run's ``--extra-inputs`` into + ``dispatch_overrides`` at parse), the envelope carries + ``endpoint_extra_applied: True`` so the worker skips re-merging + ``endpoint.extra`` and the adapter-owned values win. Both keys are OMITTED + when unset, so header-less / flag-less corpora envelopes stay byte-identical. + """ + overrides = dict(node.extra_body or {}) + if node.model is not None and "model" not in overrides: + overrides["model"] = node.model + if node.max_tokens is not None and "max_output_tokens" not in overrides: + overrides["max_output_tokens"] = node.max_tokens + if node.raw_tools is not None and "tools" not in overrides: + overrides["tools"] = node.raw_tools + envelope = { + "prompt_segment_ids": prompt_segment_ids, + "dispatch_overrides": overrides, + "stream": bool(node.streaming), + } + if node.extra_headers: + envelope["extra_headers"] = dict(node.extra_headers) + dispatch_meta = (node.metadata or {}).get("dispatch", {}) + if dispatch_meta.get("endpoint_extra_applied"): + envelope["endpoint_extra_applied"] = True + return envelope + + +def _trace_trie_envelopes( + llm_nodes: dict[str, LlmNode], ordinals: dict[str, int] +) -> list[NodeEnvelope]: + """Build the per-node ``prompt_segment_ids`` envelopes for one trie trace. + + Used by the streaming :func:`iter_trace_segment_payloads`; the envelope + bytes are ``orjson.dumps(_trie_envelope(...))`` -- byte-identical to what + the eager unified builders derive from the same nodes (the byte-equality + contract). Nodes without a ``prompt_segment_ids`` path are skipped. + + Slot-carrying nodes (assembly ``items`` / ``capture`` in + ``metadata["trie"]``) are REJECTED loudly: this streaming envelope carries + neither, so silently continuing would persist a manifest missing the node's + dynamic slots (the eager :func:`build_unified_trie_store_interned` is the + only path that persists them today). + """ + envelopes: list[NodeEnvelope] = [] + for node_id, ordinal in sorted(ordinals.items(), key=lambda kv: kv[1]): + node = llm_nodes[node_id] + trie_meta = (node.metadata or {}).get("trie") or {} + if trie_meta.get("assembly") or trie_meta.get("capture"): + raise NotImplementedError( + f"node {node_id!r}: carries trie assembly items/capture, which " + "the streaming segment-store split does not support; " + "slot-carrying graphs must build through the eager store path" + ) + prompt_segment_ids = _prompt_segment_ids(node) + if prompt_segment_ids is None: + continue + envelopes.append( + NodeEnvelope( + node_ordinal=ordinal, + phase_variant=_PROFILING_VARIANT, + envelope_bytes=orjson.dumps(_trie_envelope(node, prompt_segment_ids)), + ) + ) + return envelopes + + +def graph_carries_assembly_slots(parsed: ParsedGraph) -> bool: + """True when any ``LlmNode`` carries trie assembly items/capture metadata. + + The schedule-plane t*-gate predicate: ``workload_detect._gate_dynamic_slots_vs_tstar`` + is now its only production caller, using it to reject a graph that both + carries dynamic slots AND engages a t* snapshot window (a slot producer + chopped into warmup would leave its consumer's pool value undefined). It no + longer routes the store build: every non-weka format takes the in-process + interned drain regardless of slots (that drain is the only one that persists + slot envelopes anyway), and weka never carries them -- so this is the ONE + definition of "carries dynamic slots" for the gate, not a store-route fork. + + It still scans the SAME ``trie_meta.get("assembly") or + trie_meta.get("capture")`` condition the streaming envelope + (:func:`_trace_trie_envelopes`) raises ``NotImplementedError`` on, so the + two agree on what "slots" means; but because slot-carrying graphs never take + the weka streaming route, that rejection is unreachable armor rather than a + routing dependency. + + Why the ``capture`` clause is safe for the t*-gate: at graph level + ``capture`` and ``assembly`` are equivalent. ``capture: true`` is only ever + stamped on producers referenced by some slot's assembly program (see + native lowering), so a capture-without-assembly node cannot exist in a + lowered graph. The ``capture`` clause is therefore redundant-but-cheap armor + here, which is why ``_gate_dynamic_slots_vs_tstar`` can resolve this same + union predicate without any behavior difference from an assembly-only check. + """ + graphs = [parsed.graph, *parsed.graphs.values()] + for graph in graphs: + for node in graph.nodes.values(): + if not isinstance(node, LlmNode): + continue + trie_meta = (node.metadata or {}).get("trie") or {} + if trie_meta.get("assembly") or trie_meta.get("capture"): + return True + return False + + +def iter_trace_segment_payloads(parsed: ParsedGraph) -> Iterator[TraceSegmentPayload]: + """Yield per-trace trie payloads (envelopes + segments) for ``parsed``. + + The worker-side counterpart to the eager unified builders: for each trace it + builds the same ``prompt_segment_ids`` envelopes AND carries the trace's + pool segments as ``(id, role, content, wire_json)`` tuples so the parent + (:func:`build_unified_trie_store_from_payloads`) can drain the unified store + without re-parsing the row. ``parsed.segment_pool`` MUST be present (trie + IR); the pool's content-addressed entries dedup across rows in the consumer. + + After the in-process-drain flip this has NO main-process production caller: + the non-weka route (dynamo / native / dag_jsonl) drains the whole parse + eagerly through :func:`build_unified_trie_store_interned`, never through + payloads. Its remaining production use is inside the weka pool workers + (:func:`weka.trace_parallel._parse_item_to_segment_payloads`), for which it + is the reference shape of the ``TraceSegmentPayload``s those workers ship to + :func:`build_unified_trie_store_from_payloads`. It is also the + streaming-drain oracle for the byte-parity suites + (``test_dynamo_streaming_store_parity`` / ``test_dag_jsonl_streaming_store_parity``), + which feed a whole parse through it to prove the streamed store equals the + interned store byte-for-byte. + """ + from aiperf.dataset.graph.codecs import encode_parsed_graph_msgpack + from aiperf.dataset.graph.graph_meta_sidecar import strip_replay_text + + pool = parsed.segment_pool + segments: list[tuple[str, str, str, str | None]] = ( + [(s.id, s.role, s.content, s.wire_json) for s in pool.by_id.values()] + if pool is not None + else [] + ) + # Content-free structural graph for the sidecar: the canonical + # ``strip_replay_text`` empties ``replay_outputs``, the segment pool (kept + # non-None so the loaded graph still reads as a segment-store IR), AND + # every ``LlmNode.prompt`` (the + # dominant inline-content field for the trie IR). Shipped once (like + # ``segments``); the streaming consumer merges these across rows and writes the + # ``graph_meta`` sidecar so the TimingManager skips the whole-corpus re-parse. + # Topology, edges, arrival offsets, and node metadata (incl. + # ``prompt_segment_ids`` refs) are preserved verbatim. + structural_bytes = encode_parsed_graph_msgpack(strip_replay_text(parsed)) + # Segments are shipped once per trace; the consumer dedups by id. A + # single-trace ParsedGraph (the streaming worker case) carries only its own + # trace's pool, so this stays memory-bounded. + for index, trace in enumerate(parsed.traces): + llm_nodes = _trie_llm_nodes(parsed, trace) + ordinals = trie_node_ordinals(llm_nodes) + envelopes = _trace_trie_envelopes(llm_nodes, ordinals) + yield TraceSegmentPayload( + trace_id=trace.id, + node_ordinals=ordinals, + envelopes=envelopes, + # Attach the pool only to the first trace's payload to avoid shipping + # the same single-trace pool N times; multi-trace ParsedGraphs are + # not produced by the streaming worker (one row == one trace). + segments=segments if index == 0 else [], + structural_graph=structural_bytes if index == 0 else b"", + ) + + +async def build_unified_trie_store_from_payloads( + payloads: Iterable[TraceSegmentPayload], + store: GraphSegmentUnifiedBackingStore, + *, + structural_sink: list[bytes] | None = None, +) -> dict[str, dict[str, int]]: + """Drain STREAMED trie payloads into the ONE unified store; return the catalog. + + The streaming counterpart to the eager + :func:`build_unified_trie_store_interned`, and -- after the in-process-drain + flip -- the WEKA-only production drain (`GraphStoreBuilder._build_graph_store_streaming_trie` + feeds it the worker-pool payload stream; the non-weka route drains the whole + parse eagerly instead). Each payload's segments are + ``put_segment``'d into the unified content pool (idempotent dedup on the + content-addressed id bounds RAM -- the streaming property) and each node's + ``prompt_segment_ids`` envelope is resolved to int handles and written as + an interned manifest. This + is what makes the unified store the ACTUAL store on the corpus-scale weka + path, not just the eager path -- the worker opens the SAME unified store it + opens for an eager run. Segments are put before the manifests that reference + them (within a payload, and earlier payloads' segments are already resident), + so every handle resolves. Returns ``{trace_id: {node_id: node_ordinal}}``. + + When ``structural_sink`` is provided, each payload's non-empty + ``structural_graph`` (content-free per-trace topology, msgpack) is appended + to it as the stream drains. The caller merges these into the corpus + structural graph and writes the ``graph_meta`` sidecar so the TimingManager + skips its whole-corpus re-parse. The bytes are content-free (empty pool + + empty ``replay_outputs``), so accumulating all of them stays bounded. + """ + catalog: dict[str, dict[str, int]] = {} + for payload in payloads: + for segment_id, role, content, wire_json in payload.segments: + store.put_segment(segment_id, role, content, wire_json=wire_json) + for rec in payload.envelopes: + envelope = orjson.loads(rec.envelope_bytes) + handles: list[int] = [] + for sid in envelope["prompt_segment_ids"]: + handle = store.segment_handle(sid) + if handle is None: + raise ValueError( + f"trace {payload.trace_id} node ordinal " + f"{rec.node_ordinal} references segment {sid!r} absent " + "from the unified pool" + ) + handles.append(handle) + store.add_node_manifest_interned( + payload.trace_id, + rec.node_ordinal, + rec.phase_variant, + handles, + envelope.get("dispatch_overrides", {}), + bool(envelope.get("stream", False)), + extra_headers=envelope.get("extra_headers"), + endpoint_extra_applied=bool(envelope.get("endpoint_extra_applied")), + ) + catalog[payload.trace_id] = payload.node_ordinals + if structural_sink is not None and payload.structural_graph: + structural_sink.append(payload.structural_graph) + await store.finalize() + return catalog + + +def _resolve_assembly_items( + assembly: list[dict] | None, + catalog_key: str, + ordinals: dict[str, int], + handle_for, +) -> list[dict] | None: + """Resolve a lowering assembly program to store-addressed ``items``. + + The lowering stamps tokens with hex segment ids and producer NODE IDS + (``metadata["trie"]["assembly"]``); the persisted + envelope carries int handles and node ORDINALS — the worker's native keys. + A producer absent from the trace's ordinal map is build-time corruption. + """ + if not assembly: + return None + + def _ordinal_for(node_id: str) -> int: + ordinal = ordinals.get(node_id) + if ordinal is None: + raise ValueError( + f"trie node {catalog_key} slot references producer " + f"{node_id!r} absent from the trace's ordinal map" + ) + return ordinal + + items: list[dict] = [] + for token in assembly: + if "seg" in token: + items.append({"h": handle_for(token["seg"])}) + elif "s" in token: + items.append({"s": {"src": _ordinal_for(token["s"]["src"])}}) + elif "m" in token: + parts = [ + {"sv": _ordinal_for(part["sv"])} if "sv" in part else part + for part in token["m"]["parts"] + ] + items.append({"m": {"role": token["m"]["role"], "parts": parts}}) + else: + raise ValueError( + f"trie node {catalog_key} assembly carries unknown token {token!r}" + ) + return items + + +async def build_unified_trie_store_interned( + parsed: ParsedGraph, store: GraphSegmentUnifiedBackingStore +) -> dict[str, dict[str, int]]: + """Drain an eager trie ``ParsedGraph`` into ONE interned (A2) unified store: + assign int handles during pool drain, then write each node's manifest as a + handle path. + + Every pool ``Segment`` is written via :meth:`put_segment`, which assigns a + dense insertion-index handle (the ``i``-th segment drained gets handle ``i``). + Each trie ``LlmNode``'s hex ``prompt_segment_ids`` path is then resolved to + those int handles via :meth:`segment_handle`, and the profiling manifest is + written as a handle path via :meth:`add_node_manifest_interned`. The + hex->handle map lives ONLY in ``store`` (build-time): workers read the handle + path directly and never see hex ids. A node referencing a segment absent from + the pool is a build-time corruption and raises ``ValueError``. + + Node set + ordinal source: iterates :func:`flat_trie_ordinals` so it + covers the top graph, keying each + manifest by the SINGLE monotonic ordinal that helper assigns (the SAME ordinal + the schedule plane resolves). This + reduces to :func:`trie_node_ordinals` over the top graph, so weka manifests are + keyed by bare-id ordinals exactly as before. Returns the + ``{trace_id: {node_id: node_ordinal}}`` catalog keyed by bare node id (the IR + is flat). + """ + + pool_segments = parsed.segment_pool.by_id + for segment in pool_segments.values(): + store.put_segment( + segment.id, segment.role, segment.content, wire_json=segment.wire_json + ) + + catalog: dict[str, dict[str, int]] = {} + for trace in parsed.traces: + nodes = _trie_llm_nodes(parsed, trace) + ordinals = flat_trie_ordinals(parsed, trace) + for catalog_key, ordinal in ordinals.items(): + node = nodes[catalog_key] + prompt_segment_ids = _prompt_segment_ids(node) + if prompt_segment_ids is None: + continue + + def _handle_for(sid: str, catalog_key: str = catalog_key) -> int: + handle = store.segment_handle(sid) + if handle is None: + raise ValueError( + f"trie node {catalog_key} references segment {sid!r} " + "absent from pool" + ) + return handle + + handles = [_handle_for(sid) for sid in prompt_segment_ids] + trie_meta = (node.metadata or {}).get("trie") or {} + items = _resolve_assembly_items( + trie_meta.get("assembly"), catalog_key, ordinals, _handle_for + ) + envelope = _trie_envelope(node, prompt_segment_ids) + store.add_node_manifest_interned( + trace.id, + ordinal, + _PROFILING_VARIANT, + handles, + envelope.get("dispatch_overrides") or {}, + bool(envelope.get("stream", False)), + items=items, + capture=bool(trie_meta.get("capture")), + extra_headers=envelope.get("extra_headers"), + endpoint_extra_applied=bool(envelope.get("endpoint_extra_applied")), + ) + catalog[trace.id] = ordinals + await store.finalize() + return catalog + + +__all__ = [ + "TraceSegmentPayload", + "build_unified_trie_store_from_payloads", + "build_unified_trie_store_interned", + "graph_carries_assembly_slots", + "iter_trace_segment_payloads", + "flat_trie_ordinals", + "trie_node_ordinals", +] diff --git a/src/aiperf/dataset/graph/segment_ir/trie_content.py b/src/aiperf/dataset/graph/segment_ir/trie_content.py new file mode 100644 index 0000000000..a0e85faaf0 --- /dev/null +++ b/src/aiperf/dataset/graph/segment_ir/trie_content.py @@ -0,0 +1,1149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Format-agnostic LCP-trie content lowering for the segment-trie IR. + +The shared core behind :func:`~aiperf.dataset.graph.adapters.weka.trie_build.build_trie_graph`, +moved verbatim from ``adapters/weka/trie_build.py`` so any adapter that can +normalize its recorded requests into :class:`TrieRequest`/:class:`TrieNode` +reuses ONE pipeline (weka and dynamo both lower through it). Any behavior +change here breaks weka byte-exactness -- treat edits as frozen. + +The pipeline (:func:`build_trie_ir`), in exactly the weka build order: + +1. **Content parents** (:func:`resolve_content_parents`): a node's + content-parent is the earlier node whose ``hash_ids`` is the longest full + prefix (tie-break most recent), else the longest partial-LCP branch point + (tie-break earliest), else a fresh root. +2. **Idle warp + ranks + interval edges**: the active-interval idle-gap warp + (:func:`apply_idle_gap_warp`), + the time-consistent global rank, and the finished-before interval-order + edges (:mod:`~aiperf.dataset.graph.segment_ir.interval_order`). +3. **Frozen block tags** (:func:`compute_asst_caps` + :func:`assign_block_tags`): + every covered block gets one ``(role, starts_new_message)`` tag, assigned by + the node that first materializes it and then FROZEN, so a shared block + prefix yields identical tags -> identical messages -> a real cache prefix. +4. **Per-node emission** (:func:`_assemble_messages_from` + covered-count ISL + gate :func:`assert_covered_isl` + :func:`emit_response_segment`): group the + tags into messages, emit one content-addressed pool entry per message chained + root->tip, gate the reconstructed token count, and append the recorded + assistant output as one trailing pool segment. Because the content-parent's + tags are copied VERBATIM into the inherited prefix, every WHOLE message + strictly inside a node's ``inherited`` block count re-derives the exact sid + the parent already emitted; the driver therefore SPLICES the parent's sid + chain for those messages (from its :class:`_EmissionRecord`) and emits fresh + only the straddling fragment plus the new-region messages -- byte-identical + output (a spliced ``pool.add`` would have been a dedup no-op) without the + quadratic re-decode+re-hash of the shared prefix. :func:`assemble_messages` + stays the whole-prefix (``start_block=0``) wrapper for callers/tests. +""" + +from __future__ import annotations + +import math +from array import array +from bisect import bisect_right +from collections.abc import Callable, Iterable +from dataclasses import dataclass, field +from typing import Any + +import msgspec + +from aiperf.dataset.graph.models import ( + START_NODE_ID, + ChannelRequirement, + ChannelSpec, + GraphRecord, + LlmNode, + ProvenanceSpec, + StaticEdge, +) +from aiperf.dataset.graph.segment_ir.interval_order import ( + apply_start_anchors, + build_interval_edges, + compute_ranks, +) +from aiperf.dataset.graph.segment_ir.pool import SegmentPool + +# --- block geometry (merged verbatim from segment_ir/block_geometry.py) ----- +# Format-agnostic block-aligned LCP geometry: operates on plain +# ``(hash_ids, token_count, block_size)`` values, originally moved verbatim from +# ``adapters/shared/content.py`` so weka and dynamo share ONE geometry source +# (any behavior change here breaks weka byte-exactness -- treat edits as frozen). + + +def longest_common_prefix(prev_hash_ids: list[int], curr_hash_ids: list[int]) -> int: + """Return the index of the first differing element of the two sequences.""" + n = min(len(prev_hash_ids), len(curr_hash_ids)) + for i in range(n): + if prev_hash_ids[i] != curr_hash_ids[i]: + return i + return n + + +@dataclass(slots=True) +class TurnBlockGeometry: + """Block-aligned geometry of one turn relative to its content parent. + + All four fields are WHOLE-block counts/token-counts on the block-aligned + grid; the ``curr_in_tokens % bs`` partial tail is deliberately excluded from + every field (it is synthesized separately at message assembly, never frozen + into a block role). ``synth_tail_n`` counts ONLY the tokens of missing WHOLE + blocks (hash_ids truncated below ``curr_in_tokens // bs``), not the partial + remainder. + """ + + lcp: int + m_curr_covered: int + new_blocks_count: int + synth_tail_n: int + + +def compute_turn_block_geometry( + prev_hash_ids: list[int], + curr_hash_ids: list[int], + curr_in_tokens: int, + block_size: int, +) -> TurnBlockGeometry: + """Compute the block-aligned geometry of ``curr`` relative to ``prev``. + + The SINGLE geometry source shared by :func:`block_role_split` and the + trailing-user cap planner (``_compute_asst_caps``). ``m_curr_covered`` is the + number of covered WHOLE blocks (``min(len(curr_hash_ids), in // bs)``); + ``lcp`` is the longest common prefix of the two hash-id lists; + ``new_blocks_count`` is the covered blocks past the LCP; ``synth_tail_n`` is + the token count of missing WHOLE blocks only (never the ``in % bs`` partial + tail). + """ + bs = block_size + m_full = curr_in_tokens // bs + m_covered = min(len(curr_hash_ids), m_full) + lcp = longest_common_prefix(prev_hash_ids, curr_hash_ids) + return TurnBlockGeometry( + lcp=lcp, + m_curr_covered=m_covered, + new_blocks_count=max(0, m_covered - lcp), + synth_tail_n=max(0, (m_full - len(curr_hash_ids)) * bs), + ) + + +def block_role_split( + *, + prev_hash_ids: list[int], + curr_hash_ids: list[int], + curr_in_tokens: int, + prev_out_tokens: int, + block_size: int, + max_asst_blocks: int | None, + parent_has_user: bool, + parent_covered_blocks: int | None = None, +) -> tuple[int, list[str]]: + """Split a turn's new blocks into per-block ``assistant``/``user`` roles. + + Returns ``(inherited, roles)`` where ``inherited`` is the block count carried + over from the content parent and ``roles`` is the creation-time role tag for + each new block. ``inherited`` is ``min(lcp, len(prev_hash_ids), + m_curr_covered)`` -- clamped to this node's OWN covered-block count so a child + that shares a full block prefix but declares a smaller ``in`` (``in // bs < + lcp``) never inherits more tags than it emits (which would trip the ISL gate + on an otherwise-legitimate trace). When ``parent_covered_blocks`` is provided + it is additionally clamped to the parent's actual covered-block count, so a + parent that UNDER-COVERS (fewer covered blocks than the child's LCP) yields + the covered-but-uninherited blocks as NEW blocks the child materializes and + tags itself. Passing ``None`` preserves the original behavior for every + existing caller. + ``ceil(prev_out / bs)`` leading new blocks are attributed to the assistant + (the previous turn's response), clamped to the available new-block width and + to ``max_asst_blocks`` when set; when the parent has no user context (a + context-loss branch) every new block becomes ``user``. + + Trailing-user is guaranteed at BLOCK CREATION (frozen): a node whose new + region is all-assistant flips its OWN last new block to user. Because the + tag is inherited verbatim by every inheritor, the flip is consistent across + the whole subtree and cache-safe (turn boundaries align framing). + """ + geo = compute_turn_block_geometry( + prev_hash_ids, curr_hash_ids, curr_in_tokens, block_size + ) + inherited = min(geo.lcp, len(prev_hash_ids), geo.m_curr_covered) + if parent_covered_blocks is not None: + inherited = min(inherited, parent_covered_blocks) + new_n = max(0, geo.m_curr_covered - inherited) + asst = math.ceil(prev_out_tokens / block_size) if prev_out_tokens > 0 else 0 + if not parent_has_user: + asst = 0 + asst = min(asst, new_n) + if max_asst_blocks is not None: + asst = min(asst, max_asst_blocks) + if asst == new_n and asst > 0: + asst -= 1 # trailing-user (frozen at creation): a node whose new region is all-assistant + # flips its OWN last new block to user; inherited verbatim, so consistent + # across all inheritors and cache-safe (turn boundaries align framing). + roles = ["assistant"] * asst + ["user"] * (new_n - asst) + return inherited, roles + + +# --- idle-gap warp (merged verbatim from segment_ir/idle_warp.py) ----------- +# Active-interval idle-gap warp: operates on plain ``(raw_start, raw_end)`` +# intervals and the trie-node protocol, originally moved verbatim from +# ``adapters/weka/trie_build.py`` so weka and dynamo share ONE idle-warp source +# (any behavior change here breaks weka byte-exactness -- treat edits as frozen). + + +class ActiveIdleWarp: + """Idle-gap warp over the UNION of request ACTIVE INTERVALS, not their starts. + + Plot every request in the trace as its active interval ``[raw_start, + raw_end]`` (``raw_end = t + api_time``) on one line. A true IDLE gap is a + stretch where NOTHING is running: between the latest end of everything + started so far (the running max end) and the next request's start. Any idle + gap longer than ``cap`` is collapsed to ``cap``; every later timestamp shifts + left by the excess. + + This is the crucial difference from capping START-to-START gaps: an active + stretch -- a single long request's processing, OR overlapping subagents -- + is NEVER cut, so every request keeps its EXACT temporal shape (durations and + overlaps preserved) and only the dead air between requests is removed. + Capping start-to-start gaps instead eats into a long request's own + ``api_time`` (warping its end PAST the next request's start), which distorts + the shape and manufactures false overlaps. Because no cut ever falls inside + an active interval, ``warped_end == warped_start + api_time`` always holds and + a request that genuinely finished before another (raw) still does so warped. + """ + + def __init__(self, intervals: list[tuple[float, float]], cap: float) -> None: + # ``_cuts`` is an ascending list of ``(next_start, cumulative_excess)``: + # every timestamp at/after ``next_start`` shifts left by that cumulative + # excess. Built by a sweep over intervals sorted by start. + self._cuts: list[tuple[float, float]] = [] + if not intervals: + return + ordered = sorted(intervals) + running_end = ordered[0][1] + cumulative = 0.0 + for start, end in ordered[1:]: + if start > running_end: # nothing active in (running_end, start) + idle = start - running_end + if idle > cap: + cumulative += idle - cap + self._cuts.append((start, cumulative)) + if end > running_end: + running_end = end + + def map(self, t: float) -> float: + shift = 0.0 + for next_start, cumulative in self._cuts: + if t < next_start: + break + shift = cumulative + return t - shift + + +def apply_idle_gap_warp(nodes: list, idle_gap_cap_seconds: float | None) -> None: + """Stamp each node's ``warped_start`` from the active-interval idle warp. + + ``nodes`` must expose ``raw_start`` / ``raw_end`` / ``request.t`` and a + writable ``warped_start`` (the trie-node protocol). ``None`` cap disables + the warp (raw ``request.t`` passthrough). + """ + if idle_gap_cap_seconds is None: + for node in nodes: + node.warped_start = node.request.t + return + warp = ActiveIdleWarp( + [(n.raw_start, n.raw_end) for n in nodes], idle_gap_cap_seconds + ) + for node in nodes: + node.warped_start = warp.map(node.request.t) + + +# --- content-addressed message chaining (merged verbatim from segment_ir/message_addressing.py) --- +# The shared root->tip discipline both adapters use to turn a per-turn message +# sequence into a ``SegmentPool`` path: add one segment per message, threading +# each segment's ``parent_id`` onto the previous segment's content-addressed id +# (the first is a root, ``parent_id=None``). Adapters differ only in how they +# DERIVE each message's ``(role, content, tokens)``; the chaining loop -- and +# therefore the pool insertion order and resulting segment ids -- is identical. + + +# One message to content-address: (role, content, tokens). +MessageUnit = tuple[str, str, list[int]] + + +def add_message_chain(pool: SegmentPool, messages: Iterable[MessageUnit]) -> list[str]: + """Add each message as one pool segment, threading ``parent_id`` root->tip. + + Returns the ordered per-message segment ids (the ``prompt_segment_ids`` path). + The first message is a root (``parent_id=None``); each subsequent segment is + parented on the prior segment's id. + """ + ids: list[str] = [] + prev_id: str | None = None + for role, content, tokens in messages: + prev_id = pool.add(role=role, content=content, tokens=tokens, parent_id=prev_id) + ids.append(prev_id) + return ids + + +@dataclass(frozen=True) +class ReconCallbacks: + """The three deterministic content callbacks the message-unit builder needs. + + Injectable so unit tests can drive the builder with collision-free stub + decoders (no tokenizer / corpus build), while the production default wires a + :class:`~aiperf.dataset.graph.adapters.shared.content.CorpusContentSynthesizer`'s + byte-faithful callbacks (see + :func:`~aiperf.dataset.graph.adapters.weka.trie_build._default_callbacks`). + + DETERMINISM CONTRACT (relied on by prefix-path reuse in :func:`build_trie_ir`): + ``decode_block_tokens`` MUST be a pure per-build function of its hash-id + arguments -- the SAME hash id decodes to the SAME tokens for every call + within one :func:`build_trie_ir` invocation. Every production callback + satisfies this by construction (each caches per ``(root seed, trace_id, + hash_id, block_size)``; see ``_default_callbacks``). The reuse path decodes + each shared block ONCE, at first materialization, and splices the resulting + sid into every inheritor; a callback that drifted between the parent's and a + child's emission (a contract violation) would no longer be re-caught at the + child by the ISL gate. Successful-build store bytes are unaffected either + way (a re-decoded duplicate segment was already discarded by pool dedup). + """ + + decode_block_tokens: Callable[[list[int]], list[int]] + sample_partial_tail_tokens: Callable[[int, str], list[int]] + decode_tokens_to_text: Callable[[list[int]], str] + block_exact: bool = True + """``decode_block_tokens`` returns EXACTLY ``block_size`` tokens per block. + + Test-only escape hatch; all production callbacks leave this True (the + assembled-token ISL gate hard-aborts the build on drift). Set False ONLY by + unit-test stub callbacks whose tiny block runs would otherwise trip the gate + on every node.""" + + +@dataclass(slots=True) +class TrieRequest: + """Normalized per-request view of one recorded LLM call.""" + + hash_ids: list[int] + """Full-prompt block-hash list (recorded or virtual; opaque position keys).""" + input_length: int + """Recorded prompt token count.""" + output_length: int + """Recorded completion token count.""" + t: float + """Trace-relative request start, seconds.""" + api_time: float + """Recorded request duration, seconds.""" + model: str | None = None + """Recorded model name when present.""" + streaming: bool = False + """Whether the recorded request streamed.""" + ttft: float | None = None + """Recorded time-to-first-token, seconds; None for non-streaming requests.""" + + +@dataclass +class TrieNode: + """One recorded leaf request plus its derived structural context.""" + + node_id: str + request: TrieRequest + # Index into the flattened recorded-order list (also the recorded ``t`` + # order tiebreak key). + order: int + # Enclosing async_launched subtree-root ids (transitive). Timing async-exclusion only. + async_ancestors: frozenset[str] = field(default_factory=frozenset) + # Global time-consistent rank (stamped by ``interval_order.compute_ranks``). + rank: int = 0 + # Resolved content-parent (longest hash-id prefix / branch point), else + # ``None`` for a fresh root. Filled in a second pass. CONTENT/PROMPT ONLY: + # selects which segment-pool prefix this turn materializes; it is NOT a + # timing/dependency cause (a branch point can be arbitrarily far back, which + # would make the firing delay the cumulative warped distance -- the + # aggregate-timestamp bug). Timing anchors on the interval-order finished- + # before frontier (:func:`~aiperf.dataset.graph.segment_ir.interval_order.build_interval_edges`) instead. + content_parent: TrieNode | None = field(default=None) + # Recorded raw start ``request.t`` mapped onto the idle-gap-warped clock + # (see :func:`apply_idle_gap_warp`). + # Equals ``request.t`` when no cap is active. Stamped in a pass after + # content-parent resolution, before edge building. + warped_start: float = field(default=0.0) + # Node id of this request's CAUSAL predecessor: the spawning request for a + # subagent's first inner request, else the previous request in its own + # chain; None for chain roots. Consumed by apply_start_anchors -- when the + # causal parent is still IN FLIGHT at this node's recorded start, the + # node's incoming edges are replaced with one start-anchored edge. + causal_parent_id: str | None = field(default=None) + # Dynamo-only recorded metadata attached during trie lowering; None for + # other adapters. + dynamo_meta: dict[str, Any] | None = None + + @property + def start(self) -> float: + """Node start on the idle-gap-warped clock (raw ``request.t`` when uncapped).""" + return self.warped_start + + @property + def raw_start(self) -> float: + """Recorded start on the RAW clock (``request.t``); interval-order input.""" + return self.request.t + + @property + def raw_end(self) -> float: + """Recorded completion on the RAW clock (``request.t + api_time``). + + The who-finished-before-whom ground truth for the interval-order edge rule + and the active-interval idle-gap warp -- both read raw timestamps. + """ + return self.request.t + (self.request.api_time or 0.0) + + @property + def end(self) -> float: + """Completion time on the warped clock = warped start + raw ``api_time``. + + ``api_time`` is the request's own server-processing duration, not an + inter-request idle gap, so it is NOT warped -- it is added raw to the + warped start. End-to-start delays therefore subtract a RAW ``api_time`` + from a warped start-to-start gap. + """ + return self.warped_start + (self.request.api_time or 0.0) + + +@dataclass(slots=True) +class TrieNodeBuild: + """Per-node output of the shared trie build.""" + + prompt_path: list[str] + """Ordered per-message SegmentPool ids for the node's prompt.""" + response_id: str + """Pool id of the node's synthesized assistant response segment.""" + small_prompt: bool = False + """True when the covered-count was 0 and the tiny-prompt fallback fired.""" + + +@dataclass(slots=True) +class TrieBuild: + """Whole-graph output of the shared trie build.""" + + builds: dict[str, TrieNodeBuild] + """node_id -> per-node build artifacts.""" + edges_by_node: dict[str, list[StaticEdge]] + """node_id -> incoming interval-order StaticEdges.""" + + +@dataclass(slots=True) +class _EmissionRecord: + """Per-node prompt-message bookkeeping for prefix-path reuse. + + Local to ONE :func:`build_trie_ir` call (one trace per worker task); dies + with the build. ``array("q")`` int arrays keep the per-message state compact + on deep, wide traces (~16 MB worst single trace vs ~70 MB with int lists). + """ + + msg_end_blocks: array + """Exclusive end block index of each emitted prompt message (absolute).""" + cum_token_counts: array + """Cumulative ACTUAL decoded token count through each message (inclusive).""" + prompt_path: list[str] + """This node's ordered per-message pool sids (the same list object as its + :attr:`TrieNodeBuild.prompt_path` -- referenced, never duplicated).""" + + +# --- content-parent resolution -------------------------------------------- + + +def resolve_content_parents(nodes: list[TrieNode]) -> None: + """Fill each node's content-parent from the hash-id prefix tree. + + For node R, the content-parent is the earlier node (lower ``order``) whose + ``hash_ids`` is the longest FULL prefix of R's, tie-broken toward the most + recent (highest ``order``). When no earlier node is a full prefix, it is the + earlier node with the longest partial ``hash_ids`` LCP (the branch point). + With no overlap at all (LCP 0 and no full prefix), R stays a fresh root. + + This is an O(sum of ``hash_ids`` lengths) incremental prefix-automaton pass + that yields byte-for-byte the SAME selection as scanning all earlier nodes + pairwise with a full-prefix / longest-common-prefix comparison (see the + brute-force oracle in test_weka_trie_build_resolution.py), but without the + O(n^2 * m) double loop. Each node is resolved against the automaton built from all + strictly-earlier nodes, then inserted. Empty-``hash_ids`` nodes are never + inserted (they can never be a full prefix and contribute LCP 0, so skipping + them cannot change the selection). + + Representation: a flat int-state automaton, NOT a tree of node objects (that + tree cost 312 B/position; this costs 175 B). State 0 is the root; + ``transitions[(state, h)]`` is the state reached by consuming hash ``h`` from + ``state``, created EXACTLY where the node-object trie created a child + (``children.get(h) is None`` <-> ``transitions.get((state, h)) is None``), so + the reachable state graph IS that trie. Parallel ``terminal``/``passer`` lists + hold, per state, the MOST-RECENT full-prefix owner (overwrite-always: the + full-prefix tie-break favors the most recent) and the EARLIEST pass-through + owner (set-once-if-None: the partial-LCP tie-break favors the earliest). The + walk visits the same depths with the same ``matched > best_full_len`` / + ``matched > best_partial_lcp`` comparisons, so the ``content_parent`` + assignment is identical as a theorem about dict semantics. Tuple keys + ``(int, int)`` are hashed/compared by value, so arbitrary ids -- negative + virtual (dynamo) or >64-bit weka JSON -- key identically to the old nested + ``dict[int, ...]`` (the combined single-int-key variant was REJECTED because + weka ids need not fit 64 bits). Nothing iterates a dict anywhere (point + get/set only), so no ordering sensitivity exists to disturb. + """ + transitions: dict[tuple[int, int], int] = {} + terminal: list[TrieNode | None] = [None] + passer: list[TrieNode | None] = [None] + # Hoist the transition probe to a local: the dict object is mutated in place + # by ``_insert_flat`` (never rebound), so the bound method stays valid across + # nodes and this drops one attribute lookup per walk step (pure speedup, no + # semantic change). + tget = transitions.get + for r in nodes: + r_hashes = r.request.hash_ids + best_full: TrieNode | None = None + best_full_len = -1 + best_partial: TrieNode | None = None + best_partial_lcp = 0 + + state = 0 + matched = 0 # hashes consumed so far == the trie depth reached + for h in r_hashes: + nxt = tget((state, h)) + if nxt is None: + break + state = nxt + matched += 1 + # A node terminating here is a full prefix of R of length ``matched``. + term = terminal[state] + if term is not None and matched > best_full_len: + best_full_len = matched + best_full = term + # Any node walking through here shares an LCP of ``matched`` with R. + pas = passer[state] + if pas is not None and matched > best_partial_lcp: + best_partial_lcp = matched + best_partial = pas + + if best_full is not None: + r.content_parent = best_full + elif best_partial is not None: + r.content_parent = best_partial + + if r_hashes: + # Hand off the state the walk already reached: ``r_hashes[:matched]`` + # follows existing transitions whose ``passer`` an earlier inserter + # already stamped, so re-walking that prefix is a no-op -- resuming + # the insert at ``(state, matched)`` is byte-identical and skips the + # redundant probes on deep shared-prefix chains. + _insert_flat( + transitions, + terminal, + passer, + r_hashes, + r, + start_state=state, + start_index=matched, + ) + + +def _insert_flat( + transitions: dict[tuple[int, int], int], + terminal: list[TrieNode | None], + passer: list[TrieNode | None], + hashes: list[int], + node: TrieNode, + *, + start_state: int = 0, + start_index: int = 0, +) -> None: + """Insert ``node``'s ``hash_ids`` into the flat automaton, stamping passer/terminal. + + Insertion is in ``order`` (ascending). A missing transition ``(state, h)`` + creates a fresh state id ``len(terminal)`` -- appending ``None`` to BOTH the + ``terminal`` and ``passer`` lists (kept in lockstep, so ``len`` is the next + id) -- exactly where the node-object trie created a child. ``passer`` is set + only on its FIRST inserter at each state (the earliest/lowest-order node) -- + the partial-LCP tie-break favors the earliest node. ``terminal`` is + overwritten on every inserter (the latest wins) -- the full-prefix tie-break + favors the most recent. + + ``start_state`` / ``start_index`` resume the insert from a prefix + :func:`resolve_content_parents`' walk already traversed. The contract (relied + on ONLY by that caller): ``hashes[:start_index]`` leads from the root to + ``start_state`` via EXISTING transitions, and every existing state was + created by an earlier inserter that stamped its ``passer`` -- so the + set-once-if-None checks on that prefix are all no-ops and skipping them is + byte-identical. Called with the defaults it is a plain full insert from the + root (its standalone/measurement use). + """ + tget = transitions.get + t_append = terminal.append + p_append = passer.append + state = start_state + suffix = hashes if start_index == 0 else hashes[start_index:] + for h in suffix: + key = (state, h) + nxt = tget(key) + if nxt is None: + nxt = len(terminal) + transitions[key] = nxt + t_append(None) + p_append(None) + state = nxt + if passer[state] is None: + passer[state] = node + terminal[state] = node + + +# --- frozen block tags ------------------------------------------------------ + + +def compute_asst_caps(nodes: list[TrieNode], block_size: int) -> dict[str, int | None]: + """Pass-1 trailing-user planner over the GLOBAL ``content_parent`` tree. + + A degenerate pull-back (``new_blocks_count == 0`` and ``synth_tail_n == 0``) + at ``eff_lcp >= 1`` re-exposes block ``eff_lcp - 1`` of the parent lineage; to + keep the frozen boundary landing on a user block, the owner ancestor of that + block is capped. Owners are tracked as a per-node "tile" list (block index -> + owning ``node_id``), rebuilt off :func:`compute_turn_block_geometry` so the + role split and this planner share the SINGLE geometry source. The inherited + count uses the SAME three-way clamp as :func:`assign_block_tags` + (``min(lcp, len(parent_tiles), m_curr_covered)``) so on an over-share row + (``in // bs < lcp``) the re-exposed boundary block -- and therefore the + capped owner -- matches the frozen tags. + + Skips capping a root owner (a node whose ``content_parent is None``), + matching agentx's ``owner != 0`` guard. Requires :func:`resolve_content_parents` + to have run so each node's ``content_parent`` is set. + """ + is_root = {n.node_id: (n.content_parent is None) for n in nodes} + caps: dict[str, int | None] = {} + tiles: dict[str, list[str]] = {} + eff: dict[str, int] = {} + for node in nodes: + parent = node.content_parent + curr = node.request.hash_ids + if parent is None: + g = compute_turn_block_geometry( + [], curr, node.request.input_length, block_size + ) + tiles[node.node_id] = [node.node_id] * g.m_curr_covered + eff[node.node_id] = 0 + continue + pt = tiles.get(parent.node_id, []) + g = compute_turn_block_geometry( + parent.request.hash_ids, curr, node.request.input_length, block_size + ) + e = min(g.lcp, len(pt), g.m_curr_covered) + eff[node.node_id] = e + new_n = max(0, g.m_curr_covered - e) + if new_n == 0 and g.synth_tail_n == 0 and e >= 1: + owner = pt[e - 1] + if not is_root.get(owner, True): + bound = (e - 1) - eff.get(owner, 0) + if bound >= 0: + caps[owner] = ( + bound if caps.get(owner) is None else min(caps[owner], bound) + ) + tiles[node.node_id] = pt[:e] + [node.node_id] * new_n + return caps + + +def _assign_block_tags_and_inheritance( + nodes: list[TrieNode], + block_size: int, + caps: dict[str, int | None], +) -> tuple[dict[str, list[tuple[str, bool]]], dict[str, int]]: + """Freeze block tags AND return each node's inherited-block count. + + The single source of the frozen ``(role, starts_new_message)`` tags (see + :func:`assign_block_tags`) and, alongside, the geometric ``inherited`` count + each node's tags were built from. :func:`build_trie_ir`'s prefix-path reuse + consumes this EXACT ``inherited`` value (never a recomputed one) so the reuse + boundary can never drift from the frozen tags -- the two are computed once, + together. + + Returns ``(tags, inherited_by_node)`` where ``tags`` is + ``node_id -> [(role, starts_new_message), ...]`` (one entry per covered + block) and ``inherited_by_node`` is ``node_id -> inherited`` (the count of + leading blocks whose tags were copied verbatim from the content-parent). + Requires :func:`resolve_content_parents` to have run. + """ + tags: dict[str, list[tuple[str, bool]]] = {} + inherited_by_node: dict[str, int] = {} + for node in nodes: + parent = node.content_parent + parent_tags = tags.get(parent.node_id, []) if parent is not None else [] + prev_hash = list(parent.request.hash_ids) if parent is not None else [] + prev_out = parent.request.output_length if parent is not None else 0 + geo = compute_turn_block_geometry( + prev_hash, + list(node.request.hash_ids), + node.request.input_length, + block_size, + ) + # Clamp inherited to BOTH ends: a content-parent that UNDER-COVERS holds + # fewer frozen tags than its LCP with this child (inherit only what the + # parent tagged, covered-but-uninherited blocks become NEW blocks here); + # and this child may OVER-share -- a full block prefix (lcp) but a smaller + # declared in (in // bs < lcp), so clamp to its own covered count too or + # it would carry more tags than it emits and trip the ISL gate. + inherited = min(geo.lcp, len(parent_tags), geo.m_curr_covered) + parent_has_user = any(role == "user" for role, _ in parent_tags[:inherited]) + inh2, new_roles = block_role_split( + prev_hash_ids=prev_hash, + curr_hash_ids=list(node.request.hash_ids), + curr_in_tokens=node.request.input_length, + prev_out_tokens=prev_out, + block_size=block_size, + max_asst_blocks=caps.get(node.node_id), + parent_has_user=parent_has_user, + parent_covered_blocks=len(parent_tags), + ) + # Geometry is the single source of the inherited-block count; the split + # helper must agree or the frozen prefix would drift from the tags. + if inh2 != inherited: + raise ValueError( + f"node {node.node_id}: block-tag/geometry disagreement " + f"{inh2} != {inherited}" + ) + node_tags: list[tuple[str, bool]] = list(parent_tags[:inherited]) + for j, role in enumerate(new_roles): + starts = (j == 0) or (role != new_roles[j - 1]) + node_tags.append((role, starts)) + tags[node.node_id] = node_tags + inherited_by_node[node.node_id] = inherited + return tags, inherited_by_node + + +def assign_block_tags( + nodes: list[TrieNode], + block_size: int, + caps: dict[str, int | None], +) -> dict[str, list[tuple[str, bool]]]: + """Freeze a ``(role, starts_new_message)`` tag per COVERED block at creation. + + The CONTENT invariant's enforcement core: every covered block position of + every node gets exactly one ``(role, starts_new_message)`` tag, assigned by + the node that FIRST materializes it and then FROZEN. Two requests sharing a + block prefix therefore read the SAME tags on that prefix -> identical + messages -> cache-safe. Re-tagging a shared block on a later turn would + silently break that prefix identity, which is why tags freeze at creation. + + Nodes are processed in recorded (``order``) order so a content-parent is + tagged before any child reads its frozen tags. For each node: + + * Inherited blocks ``[0, inherited)`` reuse the content-parent's frozen tags + VERBATIM (``inherited = min(lcp, len(parent.hash_ids))`` from the shared + geometry). A fresh root (``content_parent is None``) inherits nothing. + * ``parent_has_user`` for the context-loss rule is computed over the + INHERITED PREFIX ONLY -- the blocks this node actually carries forward. + * New blocks ``[inherited, m_covered)`` get roles from :func:`block_role_split` + (assistant run then user run, capped, context-loss aware). + * ``starts_new_message`` is True for the node's FIRST new block (a new + recorded turn always opens a message, even when it continues the parent's + tail role -- this preserves contiguous same-role turns) OR at a role + transition within the new region; otherwise False. + + Returns ``node_id -> [(role, starts_new_message), ...]`` (one entry per + covered block, ``min(len(hash_ids), in // bs)`` of them). Requires + :func:`resolve_content_parents` to have run. Thin wrapper over + :func:`_assign_block_tags_and_inheritance` (drops the inheritance map) so the + public tag output and every existing caller/oracle stay unchanged. + """ + tags, _inherited_by_node = _assign_block_tags_and_inheritance( + nodes, block_size, caps + ) + return tags + + +# --- message + segment emission -------------------------------------------- + + +def _assemble_messages_from( + hash_ids: list[int], + block_tags: list[tuple[str, bool]], + start_block: int, + seed_parent_id: str | None, + pool: SegmentPool, + decode_block_tokens: Callable[[list[int]], list[int]], + decode_tokens_to_text: Callable[[list[int]], str], +) -> tuple[list[str], int, list[int], list[int]]: + """Emit the messages of ``block_tags[start_block:]``, chained onto a seed. + + The reuse-aware emission core: identical grouping and content-addressing to + the whole-prefix emission, but starting at ``start_block`` and threading the + first fresh message's ``parent_id`` onto ``seed_parent_id`` (the spliced + content-parent tip; ``None`` at a chain root). A new message begins at + ``start_block`` and at every block whose frozen ``starts_new_message`` is + True; consecutive False blocks join it. Each message's tokens are the + block-aligned concatenation of its blocks (NO partial tail). + + Returns ``(fresh_segment_ids, fresh_token_count, group_end_blocks, + group_token_counts)``: the ordered per-message pool ids emitted here, their + ACTUAL decoded token total, and per-message ABSOLUTE exclusive end block + indices + actual token counts (the driver folds these into the node's + :class:`_EmissionRecord`). Absolute end blocks are valid because a child's + inherited-prefix tags ARE the parent's, copied verbatim. + """ + groups: list[tuple[str, list[int]]] = [] + for k in range(start_block, len(block_tags)): + role, starts = block_tags[k] + if starts or not groups: + groups.append((role, [k])) + else: + groups[-1][1].append(k) + fresh_ids: list[str] = [] + fresh_token_count = 0 + group_end_blocks: list[int] = [] + group_token_counts: list[int] = [] + prev_id = seed_parent_id + for role, idxs in groups: + toks: list[int] = [] + for k in idxs: + toks.extend(decode_block_tokens([hash_ids[k]])) + prev_id = pool.add( + role=role, + content=decode_tokens_to_text(toks), + tokens=toks, + parent_id=prev_id, + ) + fresh_ids.append(prev_id) + fresh_token_count += len(toks) + group_end_blocks.append(idxs[-1] + 1) + group_token_counts.append(len(toks)) + return fresh_ids, fresh_token_count, group_end_blocks, group_token_counts + + +def assemble_messages( + hash_ids: list[int], + block_tags: list[tuple[str, bool]], + pool: SegmentPool, + decode_block_tokens: Callable[[list[int]], list[int]], + decode_tokens_to_text: Callable[[list[int]], str], +) -> tuple[list[str], int]: + """Group frozen per-block tags into messages and emit one content-addressed + pool entry per message, chained root->tip. + + A new message begins at the first covered block and at every block whose + frozen ``starts_new_message`` is True; consecutive ``starts_new_message=False`` + blocks join the current message. Each message's tokens are the block-aligned + concatenation of its blocks (NO partial tail), and it is emitted as one + ``pool.add`` parented at the previous message's id (first message's + ``parent_id`` is ``None``). + + Because the tags are frozen per trie position, a shared block prefix yields + identical message grouping -> identical ``(parent_id, role, tokens)`` chain + -> identical content-addressed ids -> a real cache prefix. There is NO role + coercion and NO relabeling here: trailing-user is already frozen in the tags. + + Returns ``(prompt_segment_ids, assembled_token_count)`` -- the ordered + per-message pool ids plus the ACTUAL decoded token total across all + messages, so the ISL gate can compare what was assembled (not a value + re-derived from the tag count, which would be tautological). The whole-prefix + (``start_block=0``, no seed) wrapper over :func:`_assemble_messages_from`; + behavior-identical to the pre-reuse emission, retained for callers and the + differential oracle. + """ + fresh_ids, token_count, _ends, _counts = _assemble_messages_from( + hash_ids, block_tags, 0, None, pool, decode_block_tokens, decode_tokens_to_text + ) + return fresh_ids, token_count + + +def emit_response_segment( + node: TrieNode, + *, + pool: SegmentPool, + parent_id: str | None, + callbacks: ReconCallbacks, +) -> str: + """Append the node's assistant output as one pool segment; return its id. + + The recorded ``out`` token count is synthesized into deterministic tokens + keyed on the node id so the response is content-addressed and stable. An + ``out == 0`` response still emits an (empty) segment so ``response_id`` + always names a real pool entry. + """ + n_out = max(0, node.request.output_length) + tokens = callbacks.sample_partial_tail_tokens(n_out, f"{node.node_id}:response") + content = callbacks.decode_tokens_to_text(tokens) + return pool.add( + role="assistant", content=content, tokens=tokens, parent_id=parent_id + ) + + +# --- ISL gate + fan-in ------------------------------------------------------ + + +class TrieISLMismatchError(ValueError): + """Reconstructed prompt token count did not equal the block-aligned covered-count target.""" + + +def assert_covered_isl( + node: TrieNode, prompt_token_count: int, block_size: int +) -> None: + """Hard build-abort when a reconstructed prompt misses its covered-count target. + + The target is the COVERED-count: only the blocks actually emitted (message-unit + emission covers ``min(recorded hash blocks, in // block_size)`` blocks, block- + aligned, no partial tail, no synthesis of missing whole blocks). It is NOT + ``(in // block_size) * block_size`` -- a recorded request may store fewer hash + blocks than ``in // block_size``, and demanding that would hard-abort the build + on a legitimate trace. :func:`build_trie_ir` wires the call site. + """ + expected = ( + min(len(node.request.hash_ids), node.request.input_length // block_size) + * block_size + ) + if prompt_token_count != expected: + raise TrieISLMismatchError( + f"node {node.node_id}: {prompt_token_count} tokens != covered-count " + f"{expected} (in={node.request.input_length}, bs={block_size}, " + f"{len(node.request.hash_ids)} blocks)" + ) + + +def with_fan_in_inputs(llm: LlmNode, node_edges: list[StaticEdge]) -> LlmNode: + """Attach one AND-fan-in input requirement per non-START predecessor edge. + + A multi-predecessor trie node declares NO input channels otherwise, so the + executor's ``await_inputs`` gate is a no-op and the node fires on its FIRST + completing predecessor (then a later predecessor re-schedules it past the + cycle guard). Reading each non-START predecessor's already-write-declared + ``{source}_out`` channel turns ``await_inputs`` into the recorded AND-join. + ``count=1``: every predecessor writes its ``_out`` exactly once. This adds + only the WAIT; the LlmNode prompt still comes from the segment pool, not the + channel value. START-rooted nodes (no non-START edge) keep empty ``inputs``. + """ + inputs = [ + ChannelRequirement(channel=f"{e.source}_out", count=1) + for e in node_edges + if e.source != START_NODE_ID and e.delay_after_predecessor_start_us is None + ] + if not inputs: + return llm + return msgspec.structs.replace(llm, inputs=inputs) + + +# --- driver ----------------------------------------------------------------- + + +def build_trie_ir( + nodes: list[TrieNode], + *, + block_size: int, + callbacks: ReconCallbacks, + pool: SegmentPool, + idle_gap_cap_seconds: float | None, + small_prompt_fallback: bool = False, +) -> TrieBuild: + """Run the shared trie-content pipeline over normalized nodes. + + Exactly the weka build order: content parents -> idle warp -> ranks -> + interval edges -> assistant caps -> frozen block tags -> per-node message + assembly + covered-count ISL gate + response segment. The gate compares + the ACTUAL assembled token count (skipped when + ``callbacks.block_exact`` is False -- deliberate placeholder content). + + Recorded inter-request delays always replay on the edges, warped through + :func:`apply_idle_gap_warp` (``idle_gap_cap_seconds=None`` disables the + warp, never the delays). + ``small_prompt_fallback`` emits a single user message sized to + ``input_length`` for covered-count-0 nodes instead of an empty prompt (the + ISL gate is skipped for those nodes); both recorded adapters pass True (a + recorded sub-block prompt was a real prompt -- an empty messages array is + unreplayable). False remains for callers that must preserve empty prompts. + """ + resolve_content_parents(nodes) + apply_idle_gap_warp(nodes, idle_gap_cap_seconds) + compute_ranks(nodes) + edges_by_node = build_interval_edges(nodes) + apply_start_anchors(nodes, edges_by_node) + caps = compute_asst_caps(nodes, block_size) + tags, inherited_by_node = _assign_block_tags_and_inheritance( + nodes, block_size, caps + ) + + builds: dict[str, TrieNodeBuild] = {} + # node_id -> emission record; local to this build (one trace per worker task). + records: dict[str, _EmissionRecord] = {} + for node in nodes: + node_tags = tags[node.node_id] + covered = len(node_tags) + small = False + if covered == 0 and small_prompt_fallback and node.request.input_length > 0: + toks = callbacks.sample_partial_tail_tokens( + node.request.input_length, f"{node.node_id}:tiny" + ) + prompt_path = [ + pool.add( + role="user", + content=callbacks.decode_tokens_to_text(toks), + tokens=toks, + parent_id=None, + ) + ] + small = True + # Small-prompt nodes publish NO record: their single message is not + # tag-derived, and any child structurally inherits 0 (empty tags), so + # nothing can reuse from them. + else: + parent = node.content_parent + # The reuse boundary MUST come from the geometric ``inherited`` + # (single-sourced from the tag pass above), never a tag-prefix + # comparison: tags can coincide beyond the lcp while hash ids differ, + # and splicing there would ship a wrong content-addressed sid. + inherited = inherited_by_node[node.node_id] + # Missing-parent-record guard: a resolved content_parent may publish + # no record (small-prompt parent), so gate on BOTH the record + # existing AND inherited > 0 before attempting any splice. + rec = records.get(parent.node_id) if parent is not None else None + j = 0 + splice: list[str] = [] + reused_tokens = 0 + start_block = 0 + seed_parent_id: str | None = None + if rec is not None and inherited > 0: + # Count whole parent messages ending at or before ``inherited``; + # bisect_right includes a message ending exactly at inherited. + j = bisect_right(rec.msg_end_blocks, inherited) + if j > 0: + splice = rec.prompt_path[:j] + reused_tokens = rec.cum_token_counts[j - 1] + start_block = rec.msg_end_blocks[j - 1] + seed_parent_id = splice[-1] + # The resume block must OPEN a message: it is a parent message + # boundary copied verbatim into this child's frozen tags, so a + # False here means the grouping drifted in this frozen file. + if start_block < covered: + assert node_tags[start_block][1], ( + f"node {node.node_id}: reuse resume block " + f"{start_block} does not start a message (grouping drift)" + ) + fresh_ids, fresh_tokens, group_ends, group_counts = _assemble_messages_from( + node.request.hash_ids, + node_tags, + start_block, + seed_parent_id, + pool, + callbacks.decode_block_tokens, + callbacks.decode_tokens_to_text, + ) + prompt_path = splice + fresh_ids + # Gate on the ACTUAL assembled token count (reused + fresh) so decode + # drift (blocks not sized block_size) aborts the build; placeholder + # callbacks (block_exact=False, test-only stubs) skip it. Reused + # blocks contribute the parent's actual decoded counts: they are + # ISL-verified ONCE, at first materialization, and re-verifying here + # would be identical only under the ReconCallbacks determinism + # contract (all production callbacks satisfy it). + if callbacks.block_exact: + assert_covered_isl(node, reused_tokens + fresh_tokens, block_size) + # Publish this node's record: reused ends/counts (cumulative) copied + # from the parent, then fresh groups appended with the running total + # continued from ``reused_tokens``. + end_blocks = rec.msg_end_blocks[:j] if j else array("q") + cum_counts = rec.cum_token_counts[:j] if j else array("q") + running = reused_tokens + for end_block, count in zip(group_ends, group_counts, strict=True): + end_blocks.append(end_block) + running += count + cum_counts.append(running) + records[node.node_id] = _EmissionRecord( + msg_end_blocks=end_blocks, + cum_token_counts=cum_counts, + prompt_path=prompt_path, + ) + response_id = emit_response_segment( + node, + pool=pool, + parent_id=prompt_path[-1] if prompt_path else None, + callbacks=callbacks, + ) + builds[node.node_id] = TrieNodeBuild( + prompt_path=prompt_path, response_id=response_id, small_prompt=small + ) + return TrieBuild(builds=builds, edges_by_node=edges_by_node) + + +# --- graph assembly --------------------------------------------------------- + + +def assemble_trie_graph( + nodes: list[TrieNode], + result: TrieBuild, + build_node: Callable[[TrieNode], LlmNode], + provenance: ProvenanceSpec, +) -> GraphRecord: + """Shared assembly epilogue for trie adapters: build LlmNodes, attach + fan-in edges, stamp the theoretical prefix cache, declare output channels, + and wrap in a ``version="2.0"`` :class:`GraphRecord`. + + Each node is turned into a pre-fan-in :class:`LlmNode` by ``build_node`` -- + the per-adapter builder closed over ``result.builds`` and the adapter's own + dispatch knobs (e.g. weka's per-node ``max_osl`` + cap). Its interval-order edges are appended to the + graph edge list and turned into AND-fan-in input requirements + (:func:`with_fan_in_inputs`). After every node is assembled the shared + per-trace theoretical prefix cache is stamped across the node map, then + each node's ``{node_id}_out`` scratch channel is declared and the whole + topology is wrapped under ``provenance``. The assembly order and every + constructed value are byte-identical across adapters -- only ``build_node`` + and ``provenance`` differ. + + Every ``LlmNode`` writes its response into a per-node ``{node_id}_out`` + channel; the executor's channel store rejects writes to any channel not + declared in ``graph.state`` (``UnknownChannelError``), so each output + channel is declared with the default TEXT / overwrite spec. Successors read + these channels ONLY to enforce the AND-fan-in WAIT (one ``count=1`` + requirement per predecessor) -- the prompt itself comes from the segment + pool, never the channel value -- so the default single-producer overwrite + spec is contractually sufficient. + """ + # call-scoped: prefix_cache imports TrieNode from this module + from aiperf.dataset.graph.segment_ir.prefix_cache import ( + stamp_theoretical_prefix_cache, + ) + + llm_nodes: dict[str, LlmNode] = {} + edges: list[StaticEdge] = [] + for node in nodes: + llm = build_node(node) + node_edges = result.edges_by_node[node.node_id] + edges.extend(node_edges) + llm_nodes[node.node_id] = with_fan_in_inputs(llm, node_edges) + stamp_theoretical_prefix_cache(llm_nodes, nodes) + + state = {f"{nid}_out": ChannelSpec() for nid in llm_nodes} + return GraphRecord( + version="2.0", + provenance=provenance, + state=state, + nodes=dict(llm_nodes), + edges=edges, + ) + + +__all__ = [ + "ActiveIdleWarp", + "MessageUnit", + "ReconCallbacks", + "TrieBuild", + "TrieISLMismatchError", + "TrieNode", + "TrieNodeBuild", + "TrieRequest", + "TurnBlockGeometry", + "add_message_chain", + "apply_idle_gap_warp", + "assemble_messages", + "assemble_trie_graph", + "assert_covered_isl", + "assign_block_tags", + "block_role_split", + "build_trie_ir", + "compute_asst_caps", + "compute_turn_block_geometry", + "emit_response_segment", + "longest_common_prefix", + "resolve_content_parents", + "with_fan_in_inputs", +] diff --git a/src/aiperf/dataset/graph/store_build.py b/src/aiperf/dataset/graph/store_build.py new file mode 100644 index 0000000000..c9bf8557bb --- /dev/null +++ b/src/aiperf/dataset/graph/store_build.py @@ -0,0 +1,524 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Graph unified-store build, extracted from the DatasetManager. + +:class:`GraphStoreBuilder` owns the ONE store-build pipeline for every graph +workload (weka / dynamo / native / dag_jsonl): parse or stream the workload, +drain it into the unified segment store (content pool + per-node manifests) +rooted where the worker reads, write the mandatory content-free graph_meta +sidecar, and hand back the graph facet plus both locations as a +:class:`GraphStoreBuildResult`. The DatasetManager is the only production +caller; it broadcasts the result's locations in the graph-typed dataset +notification and never rebuilds. +""" + +from __future__ import annotations + +import asyncio +import shutil +import tempfile +from collections.abc import Iterable +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +from aiperf.common.environment import Environment +from aiperf.common.exceptions import DatasetError +from aiperf.common.mixins.aiperf_logger_mixin import AIPerfLoggerMixin +from aiperf.common.models.dataset_models import GraphDatasetMetadata +from aiperf.dataset.graph.workload_detect import publish_graph_loader_tokenizer_env + +if TYPE_CHECKING: + from aiperf.config.resolution.plan import BenchmarkRun + from aiperf.dataset.graph.models import ParsedGraph + from aiperf.dataset.graph.segment_ir.store_builder import TraceSegmentPayload + from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphStoreBuildStats, + ) + +__all__ = [ + "GraphStoreBuildResult", + "GraphStoreBuilder", +] + + +def _format_store_build_stats(stats: GraphStoreBuildStats | None) -> str: + """Format a store build snapshot as a compact ``key=value`` log suffix. + + Module-level (never a method) so the ``_build_graph_store_streaming*`` build + methods -- which existing tests invoke unbound with stub selves -- can call + it without a ``self.`` attribute that those stubs lack. Total over ``None`` + (``build_stats`` is ``None`` until finalize) so a pre-finalize store logs a + marker instead of crashing the build-success line. + """ + if stats is None: + return "build_stats=unavailable" + rss = "n/a" if stats.peak_rss_mib is None else f"{stats.peak_rss_mib:.1f}" + return ( + f"segments={stats.segment_count} " + f"content_bytes={stats.content_bytes} " + f"node_manifests={stats.node_manifest_count} " + f"manifest_bytes={stats.manifest_bytes} " + f"traces={stats.trace_count} " + f"peak_rss_mib={rss}" + ) + + +@dataclass(slots=True) +class GraphStoreBuildResult: + """Everything the DatasetManager needs from one graph store build.""" + + facet: GraphDatasetMetadata + """Trace universe + per-node prefix-cache map for the broadcast.""" + + sidecar_path: Path + """Where the mandatory graph_meta sidecar was written.""" + + base_path: Path + """Store root the worker opens (derived ONCE here).""" + + +class GraphStoreBuilder(AIPerfLoggerMixin): + """Builds the unified graph segment store + sidecar for one run. + + One instance per build: :meth:`build` streams or parses the workload into + the unified store rooted at the SAME location the worker reads from + (``Environment.DATASET.MMAP_BASE_PATH`` or the system temp dir, plus the + run ``benchmark_id``), writes the mandatory graph_meta sidecar, and + returns the :class:`GraphStoreBuildResult` the DatasetManager broadcasts. + """ + + def __init__(self, run: BenchmarkRun, **kwargs) -> None: + super().__init__(**kwargs) + self.run = run + self._sidecar_path: Path | None = None + + async def build(self, graph_path: Path) -> GraphStoreBuildResult: + """Build the graph store mmap for a graph workload; return facet + locations. + + Serves every graph workload (weka / dynamo / native / dag_jsonl) + through the ONE streaming store build + (:meth:`_build_graph_store_streaming`), which parses into a + ``ParsedGraph`` (threading ``run.random_seed`` so the synthesized + corpus / node ordinals stay deterministic for the run), serializes + every node's request envelope into the graph store mmap, and recovers + the graph facet (:class:`GraphDatasetMetadata`: the trace universe + plus the per-node theoretical prefix-cache map). The worker + materializes the real per-node payloads from the graph store. + + The sidecar is mandatory for graph runs (the TimingManager only + ingests the sidecar from the broadcast path; it never re-parses), so a + drain that completed without recording its written path is a hard + build failure, not a degraded result. + """ + publish_graph_loader_tokenizer_env(self.run) + self._prestart_loader_forkserver() + + from aiperf.dataset.graph.workload_detect import resolve_graph_workload + + base_path = Environment.DATASET.MMAP_BASE_PATH or Path(tempfile.gettempdir()) + + ref = resolve_graph_workload(self.run) + fmt = ref.format if ref is not None else None + # ONE store-build pipeline for every graph workload: weka streams + # worker-pool payloads through the trie drain; every other format + # parses once in-process and drains that single parse through the + # eager interned builder. Both drains write their own sidecar. + catalog, prefix_source = await self._build_graph_store_streaming( + graph_path, base_path, fmt + ) + + node_count = sum(len(nodes) for nodes in catalog.values()) + self.info( + f"graph unified store built: {len(catalog)} traces, " + f"{node_count} node manifests at {base_path} (benchmark_id=" + f"{self.run.benchmark_id})" + ) + + # The weka payload drain drops each per-trace ParsedGraph, but the + # returned prefix source (merged structural graph on the weka drain, or + # the full parse in-process) preserves the native + # ``theoretical_prefix_cache_*`` node fields, so the per-node + # prefix-cache map is recovered from it. Trace universe + prefix-cache + # map are first-class on the graph facet -- no stub conversations. + prefix_cache_by_trace = await asyncio.to_thread( + self._build_graph_prefix_cache_by_trace, prefix_source + ) + if self._sidecar_path is None: + raise DatasetError( + "graph store build completed without recording a graph_meta " + "sidecar path; the sidecar is mandatory for graph runs" + ) + return GraphStoreBuildResult( + facet=GraphDatasetMetadata( + trace_ids=list(catalog), + prefix_cache_by_trace=prefix_cache_by_trace, + ), + sidecar_path=self._sidecar_path, + base_path=base_path, + ) + + def _prestart_loader_forkserver(self) -> None: + """Start the trace-loader forkserver before the build is offloaded. + + ``_eagerly_start_forkserver`` dup2-redirects the PROCESS-WIDE stdio fds + around the helper spawn. The graph store builds run in + ``asyncio.to_thread`` workers, so a lazily started helper (first pool + open inside the offloaded parse) would race that swap against the + event loop's live logging -- console lines emitted during the window + silently go to ``/dev/null``. Starting here, on the loop at a + known-quiet point (right after the env publish, before any parse + thread), makes the pool's later ``get_loader_mp_context`` call a + cached no-op (the context is built once and reused). Threads the SAME + preload tokenizer the pool path resolves + (``_loader_pool_context`` <- ``resolve_graph_content_tokenizer``) so + the helper preloads the right tokenizer. + """ + from aiperf.common.tokenizer import BUILTIN_TOKENIZER_NAME + from aiperf.dataset._mp_context import get_loader_mp_context + from aiperf.timing.config import resolve_graph_content_tokenizer + + # Same fallback the pool applies to its non-scheduling tokenizer_name. + get_loader_mp_context( + preload_tokenizer=resolve_graph_content_tokenizer(self.run) + or BUILTIN_TOKENIZER_NAME + ) + + async def _build_graph_store_streaming( + self, + graph_path: Path, + base_path: Path, + fmt: str | None, + ) -> tuple[dict[str, dict[str, int]], ParsedGraph]: + """Build the unified store for EVERY graph workload; two drains. + + The one store-build pipeline, chosen by ``fmt``: + + * ``weka_trace`` (worker-pool payload stream): resolves the SAME + run-derived content knobs as the in-process + :func:`parse_graph_workload` -- ONE resolution, + :func:`~aiperf.dataset.graph.workload_detect.resolve_graph_parse_context`, + whose fields (seed, tokenizer, corpus, ``max_osl``, idle-gap cap) + spread verbatim into the stream entry -- so build-plane topology, + catalog, and node ordinals stay deterministic for the run (and + content bytes additionally match when an explicit ``--random-seed`` + is set), then streams worker-pool-built trie payloads + (:func:`stream_weka_trace_segment_payloads` emits + :class:`TraceSegmentPayload`s) into the unified store via + :meth:`_build_graph_store_streaming_trie`. Each worker serializes its + trace's envelopes before returning, so the parent does not decode a + full ``ParsedGraph`` only to re-serialize the same real content. + ``idle_gap_cap_seconds`` forwards AS-IS: the resolver always yields + a resolved value (never ``UNSET``), and an explicit ``None`` means + warping DISABLED and must arrive as ``None``. + * Every other format (dynamo / native / dag_jsonl / undetected -- + ``fmt=None`` fails inside ``parse_graph_workload``'s own detection): + parse ONCE in-process (whole-graph lowering) and drain that single + parse through the eager :meth:`_build_interned_unified_store` -- the + in-process interned drain. In-process there is no worker pool to fan + out to, so the payload round trip the weka path needs is pure + overhead; the interned drain also persists dynamic-slot envelopes + (native ``@channel`` assembly items/capture, dag_jsonl live-reply + lineage) that the streaming payload envelope cannot carry. + ``dynamo_trace`` additionally takes the DIRECT write-through route: + the store is constructed BEFORE the parse and passed + as ``direct_store`` so ``build_trie_ir``'s ``pool.add`` interns each + segment straight into the store during the parse (no second RAM pool + copy); the interned drain then no-ops over the empty returned pool. + Byte-identical to the eager drain by construction -- both intern in + ``build_trie_ir``'s content-loop first-occurrence order. + + Both drains build the SAME on-disk unified store (content pool + + per-node manifests; the byte-parity suites prove it) and each writes + its own mandatory content-free graph_meta sidecar; no caller sidecar + pass exists. + + Returns ``(catalog, prefix_source)`` where the second element is the + caller's prefix-cache source: the content-free merged structural graph + (weka payload drain) or the full parse (in-process interned drain). + """ + if fmt == "weka_trace": + from aiperf.dataset.graph.adapters.weka.trace import ( + stream_weka_trace_segment_payloads, + ) + from aiperf.dataset.graph.workload_detect import ( + resolve_graph_parse_context, + ) + + ctx = resolve_graph_parse_context(self.run) + payloads = stream_weka_trace_segment_payloads( + str(graph_path), + idle_gap_cap_seconds=ctx.idle_gap_cap_seconds, + content_root_seed=ctx.content_root_seed, + content_tokenizer=ctx.content_tokenizer, + prompt_corpus=ctx.prompt_corpus, + max_osl=ctx.max_osl, + num_dataset_entries=ctx.num_dataset_entries, + max_context_length=ctx.max_context_length, + ) + return await self._build_graph_store_streaming_trie(payloads, base_path) + + # Every non-weka format: the adapter's parse returns ONE ParsedGraph + # at this layer (dynamo fans out per session-tree INSIDE its own parse + # and lowers each tree independently; dag_jsonl expands whole trees), + # so there is a single parsed result to drain. Parse once off-loop, + # then drain it through the eager interned builder: in-process there + # is no worker pool to fan out to, so the payload round trip is pure + # overhead, and the interned drain is the only one that persists + # dynamic-slot envelopes (native @channel assembly items/capture, + # dag_jsonl live-reply lineage). + from aiperf.dataset.graph.workload_detect import parse_graph_workload + from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + ) + + pool_missing_msg = ( + f"graph workload {graph_path} parsed without a segment_pool; " + "every graph parse lowers onto the unified segment store, so " + "a pool-less parse is a lowering bug" + ) + if fmt == "dynamo_trace": + # Dynamo direct write-through route: construct the + # store BEFORE the off-loop parse and thread it as ``direct_store`` + # so build_trie_ir's ``pool.add`` interns each segment STRAIGHT INTO + # the store during the parse (no second RAM pool copy). The store is + # live before the parse, so the abort()+rmtree cleanup MUST cover the + # parse too: content.blob spills incrementally, so a mid-parse + # DynamoISLMismatchError (or any failure) must leave no partial file. + unified_store = GraphSegmentUnifiedBackingStore( + base_path=base_path, + benchmark_id=self.run.benchmark_id, + ) + try: + parsed = await asyncio.to_thread( + parse_graph_workload, + self.run, + graph_path, + direct_store=unified_store, + ) + if parsed.segment_pool is None: + raise ValueError(pool_missing_msg) + catalog = await self._build_interned_unified_store( + parsed, unified_store + ) + except BaseException: + unified_store.abort() + shutil.rmtree(unified_store.data_dir, ignore_errors=True) + raise + else: + parsed = await asyncio.to_thread(parse_graph_workload, self.run, graph_path) + if parsed.segment_pool is None: + raise ValueError(pool_missing_msg) + + unified_store = GraphSegmentUnifiedBackingStore( + base_path=base_path, + benchmark_id=self.run.benchmark_id, + ) + # The store spills content.blob incrementally, so a drain that raises + # before finalize would leave a partial blob on disk; abort() + rmtree + # remove it so a later store open never trips on the half-written file. + try: + catalog = await self._build_interned_unified_store( + parsed, unified_store + ) + except BaseException: + unified_store.abort() + shutil.rmtree(unified_store.data_dir, ignore_errors=True) + raise + self.info( + f"GRAPH_SEGMENT UNIFIED store built (in-process interned drain): " + f"{len(catalog)} traces at {base_path} " + f"(benchmark_id={self.run.benchmark_id}) " + f"{_format_store_build_stats(unified_store.build_stats)}" + ) + await asyncio.to_thread(self._write_graph_sidecar, parsed, catalog, base_path) + return catalog, parsed + + async def _build_graph_store_streaming_trie( + self, + payloads: Iterable[TraceSegmentPayload], + base_path: Path, + ) -> tuple[dict[str, dict[str, int]], ParsedGraph]: + """Drain the weka worker-pool payload STREAM into the ONE unified store. + + The weka drain: :meth:`_build_graph_store_streaming` routes only + ``weka_trace`` here (every other format takes the in-process + interned drain). Each pool worker serializes its trace's trie payloads, + and this method drains that stream into the interned unified store + (content pool + per-node manifests) via + :func:`build_unified_trie_store_from_payloads`, so the corpus-scale + path builds the SAME unified store the in-process interned drain does + (the worker opens the unified reader either way). Streaming + ``put_segment`` dedup on the content-addressed id bounds RAM. Each + streamed row's content-free structural graph is collected and merged + ONCE (:meth:`_merge_structural_graphs`); the merged graph feeds both + the mandatory :meth:`_write_graph_sidecar` (the TimingManager loads the + sidecar instead of re-parsing the whole corpus) and the returned + prefix-cache source. Returns ``(catalog, merged_structural)``. + """ + from aiperf.dataset.graph.segment_ir.store_builder import ( + build_unified_trie_store_from_payloads, + ) + from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + ) + + structural_sink: list[bytes] = [] + unified = GraphSegmentUnifiedBackingStore( + base_path=base_path, + benchmark_id=self.run.benchmark_id, + ) + + def _drain_stream() -> dict[str, dict[str, int]]: + # The payload iterator blocks in multiprocessing result.get() per + # trace, so the whole drain runs in a worker thread (its own loop + # covers the store's async finalize) and this service's event loop + # keeps serving heartbeats during corpus-scale builds. + return asyncio.run( + build_unified_trie_store_from_payloads( + payloads, unified, structural_sink=structural_sink + ) + ) + + # As with the interned drain, a mid-stream failure leaves a partially + # spilled content.blob; abort() + rmtree remove the store so no + # half-written file survives for a later open. + try: + catalog = await asyncio.to_thread(_drain_stream) + except BaseException: + unified.abort() + shutil.rmtree(unified.data_dir, ignore_errors=True) + raise + self.info( + f"GRAPH_SEGMENT UNIFIED store built (streaming): " + f"{len(catalog)} traces at {base_path} " + f"(benchmark_id={self.run.benchmark_id}) " + f"{_format_store_build_stats(unified.build_stats)}" + ) + # Structural decode/merge + catalog cross-check + write are pure + # CPU/sync-IO on corpus-scale inputs; off-loop like the drain above. + merged = await asyncio.to_thread(self._merge_structural_graphs, structural_sink) + await asyncio.to_thread(self._write_graph_sidecar, merged, catalog, base_path) + return catalog, merged + + async def _build_interned_unified_store( + self, + parsed: ParsedGraph, + unified: GraphSegmentUnifiedBackingStore, + ) -> dict[str, dict[str, int]]: + """Build the interned unified segment-trie store from a whole-graph parse. + + Drains the parse's content-addressed pool into ``unified`` and writes the + per-node interned manifests via the node-typed + :func:`build_unified_trie_store_interned`. The in-process drain for + EVERY non-weka format inside :meth:`_build_graph_store_streaming`, and + the only drain that persists dynamic-slot (assembly items/capture) + envelopes. + """ + from aiperf.dataset.graph.segment_ir.store_builder import ( + build_unified_trie_store_interned, + ) + + def _build() -> dict[str, dict[str, int]]: + # The interned build is a zero-yield CPU drain (orjson-encode + + # pool copy of ALL synthesized content) until its trailing + # ``store.finalize()`` await, so it runs in a worker thread whose + # own loop covers that finalize and this service's event loop + # keeps serving heartbeats during corpus-scale builds. + return asyncio.run(build_unified_trie_store_interned(parsed, unified)) + + return await asyncio.to_thread(_build) + + def _merge_structural_graphs(self, structural_sink: list[bytes]) -> ParsedGraph: + """Merge the weka payload drain's per-trace structural graphs; hard fail on any gap. + + Only the weka payload drain (:meth:`_build_graph_store_streaming_trie`) + collects a per-trace structural stream to merge here; the in-process + interned drain keeps the whole parse and never calls this. The merged + graph is content-free but preserves node metadata, so it feeds BOTH the + mandatory graph_meta sidecar and the per-node prefix-cache map. A + missing or unmergeable structural stream is a build failure, not a + degradation. + """ + from aiperf.dataset.graph.codecs import decode_parsed_graph_msgpack + from aiperf.dataset.graph.merge import merge_parsed_graphs + + if not structural_sink: + raise DatasetError( + "graph store build produced no structural graphs: the " + "streaming drain yielded zero traces, so there is nothing to " + "benchmark and no graph_meta sidecar can be written" + ) + try: + return merge_parsed_graphs( + decode_parsed_graph_msgpack(b) for b in structural_sink + ) + except Exception as e: + raise DatasetError( + f"streaming structural graphs failed to merge: {e!r}; the " + "graph_meta sidecar is mandatory for graph runs" + ) from e + + def _write_graph_sidecar( + self, + parsed: ParsedGraph, + catalog: dict[str, dict[str, int]], + base_path: Path, + ) -> None: + """Write the mandatory content-free graph_meta sidecar for this build. + + Every graph route MUST land this file: the TimingManager only ingests + the sidecar, from the exact path recorded here for the graph-typed + dataset broadcast (it never re-parses), so a skipped or failed write + would leave the run unschedulable. A structural catalog that diverges + from the store's build catalog means the sidecar would describe a + DIFFERENT topology than the envelopes the worker reads -- hard fail. + """ + from aiperf.dataset.graph.codecs import GRAPH_META_SCHEMA_VERSION + from aiperf.dataset.graph.graph_meta_sidecar import ( + catalogs_match, + write_graph_meta_sidecar, + ) + + if not catalogs_match(parsed, catalog): + raise DatasetError( + "graph_meta sidecar catalog mismatch: the structural graph's " + "node ordinals diverged from the unified store's build " + "catalog; the TimingManager cannot schedule this run" + ) + out = write_graph_meta_sidecar( + parsed, + base_path=base_path, + benchmark_id=self.run.benchmark_id, + source_fingerprint={}, + schema_version=GRAPH_META_SCHEMA_VERSION, + ) + self._sidecar_path = out + self.info(f"graph_meta sidecar written: {out}") + + @staticmethod + def _build_graph_prefix_cache_by_trace( + parsed: ParsedGraph, + ) -> dict[str, dict[str, list[int]]]: + """Per-trace ``{node_id: [hit, total]}`` from a parsed trie graph. + + Resolves each trace's own graph (single-graph and heterogeneous + multi-graph corpora), then reads the prefix-cache counts the shared + trie build stamped on each ``LlmNode``'s native + ``theoretical_prefix_cache_hit_blocks`` / ``_total_blocks`` fields. + Empty when nothing is stamped (a non-trie graph or hash-id-free + requests). + """ + from aiperf.dataset.graph.models import resolve_trace_graph + from aiperf.dataset.graph.segment_ir.prefix_cache import ( + extract_prefix_cache_by_node, + ) + + out: dict[str, dict[str, list[int]]] = {} + for trace in parsed.traces: + trace_graph = resolve_trace_graph(parsed, trace) + by_node = extract_prefix_cache_by_node(trace_graph) + if by_node: + out[trace.id] = by_node + return out diff --git a/src/aiperf/dataset/graph/validator.py b/src/aiperf/dataset/graph/validator.py new file mode 100644 index 0000000000..c2e4c09793 --- /dev/null +++ b/src/aiperf/dataset/graph/validator.py @@ -0,0 +1,581 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Graph validator framework: shared helpers, rule implementations, and entry points. + +See `docs/reference/graph-ir-validation.md` for the full rule list. Rules 19 +and 20 are enforced at parse time (see `parser.py`). +""" + +from __future__ import annotations + +import math +from collections import defaultdict, deque +from collections.abc import Iterator +from enum import Enum + +from pydantic import Field + +from aiperf.common.models.base_models import AIPerfBaseModel +from aiperf.dataset.graph.models import ( + END_NODE_ID, + START_NODE_ID, + GraphRecord, + LlmNode, + ParsedGraph, + ReducerName, + StaticEdge, +) + + +class ValidationSeverity(str, Enum): + ERROR = "error" + WARNING = "warning" + + +class ValidationIssue(AIPerfBaseModel): + """A single validator finding.""" + + rule_id: str = Field(description="Rule identifier, e.g. 'rule-1'.") + severity: ValidationSeverity = Field( + description="ERROR (blocks execution) or WARNING (informational)." + ) + location: str = Field( + description="Human-readable pointer (e.g. 'graph.nodes.plan')." + ) + message: str = Field(description="Plain-English description of the problem.") + suggested_fix: str | None = Field( + default=None, description="Optional remediation hint." + ) + + +_RESERVED_NAMES = {START_NODE_ID, END_NODE_ID} +_KNOWN_VERSIONS = {"2.0"} + + +def validate(parsed: ParsedGraph) -> list[ValidationIssue]: + """Run every implemented rule, return all issues. Caller decides whether warnings block. + + Rules run over the main ``parsed.graph`` AND over every per-trace graph in + ``parsed.graphs`` (multi-graph workloads: weka heterogeneous directories, + per-trace native lowering). Issues found in a ``parsed.graphs`` entry are + re-located under ``graphs[]`` so the offending graph is identifiable. + """ + issues = _validate_graph_record(parsed.graph) + for name, graph in parsed.graphs.items(): + if graph is parsed.graph: + # The native lowering aliases parsed.graph to the first graphs + # entry; skip the duplicate pass. + continue + for issue in _validate_graph_record(graph): + issues.append( + issue.model_copy( + update={"location": _graphs_location(name, issue.location)} + ) + ) + return issues + + +def _validate_graph_record(graph: GraphRecord) -> list[ValidationIssue]: + """Run every implemented rule over one :class:`GraphRecord`.""" + issues: list[ValidationIssue] = [] + issues.extend(_rule_01_cycles(graph)) + issues.extend(_rule_09_overwrite_conflict(graph)) + issues.extend(_rule_11_reserved_names(graph)) + issues.extend(_rule_12_version(graph)) + issues.extend(_rule_13_provenance_tool(graph)) + issues.extend(_rule_15_expected_cached_tokens(graph)) + issues.extend(_rule_21_unreachable(graph)) + issues.extend(_rule_54_edge_delay_exclusivity(graph)) + issues.extend(_rule_55_first_token_anchor_shape(graph)) + issues.extend(_rule_56_edge_endpoints(graph)) + issues.extend(_rule_57_finite_delays(graph)) + return issues + + +def _graphs_location(graph_name: str, location: str) -> str: + """Re-root a ``graph.*`` issue location under ``graphs[]``.""" + prefix = "graph." + if location.startswith(prefix): + return f"graphs[{graph_name}].{location[len(prefix) :]}" + return f"graphs[{graph_name}]: {location}" + + +# ---- helpers --------------------------------------------------------------- + + +def _build_adjacency(graph: GraphRecord) -> dict[str, set[str]]: + """Map src -> set of declared static successors.""" + adj: dict[str, set[str]] = defaultdict(set) + for e in graph.edges: + if isinstance(e, StaticEdge): + adj[e.source].add(e.target) + return adj + + +def _node_writers(graph: GraphRecord) -> dict[str, list[str]]: + """Map channel name to the list of node ids that write to it.""" + writers: dict[str, list[str]] = defaultdict(list) + for nid, node in graph.nodes.items(): + if isinstance(node, LlmNode): + writers[node.output].append(nid) + return writers + + +# ---- rules ----------------------------------------------------------------- + + +def _rule_01_cycles(graph: GraphRecord) -> list[ValidationIssue]: + # Iterative DFS with an explicit stack: recorded corpora reach 100k+-node + # chains, far past Python's recursion limit. + adj = _build_adjacency(graph) + color: dict[str, int] = {} + GREY, BLACK = 1, 2 + + for root in [*graph.nodes.keys(), START_NODE_ID]: + if color.get(root, 0) != 0: + continue + color[root] = GREY + stack: list[tuple[str, Iterator[str]]] = [(root, iter(adj.get(root, ())))] + while stack: + node, successors = stack[-1] + for succ in successors: + c = color.get(succ, 0) + if c == GREY: + return [ + ValidationIssue( + rule_id="rule-1", + severity=ValidationSeverity.ERROR, + location="graph.edges", + message="Graph contains a cycle in the static edges.", + suggested_fix="Pre-unroll cycles in the trace topology (cycles are a future graph feature).", + ) + ] + if c == 0: + color[succ] = GREY + stack.append((succ, iter(adj.get(succ, ())))) + break + else: + color[node] = BLACK + stack.pop() + return [] + + +def _rule_09_overwrite_conflict(graph: GraphRecord) -> list[ValidationIssue]: + """Two or more nodes write the same overwrite-reducer channel.""" + writers = _node_writers(graph) + issues: list[ValidationIssue] = [] + for ch, nodes in writers.items(): + if len(nodes) <= 1: + continue + spec = graph.state.get(ch) + if spec is not None and spec.reducer is not ReducerName.OVERWRITE: + continue + issues.append( + ValidationIssue( + rule_id="rule-9", + severity=ValidationSeverity.ERROR, + location=f"graph.state.{ch}", + message=( + f"Channel '{ch}' has overwrite reducer but is written by " + f"multiple nodes: {sorted(nodes)}." + ), + suggested_fix=( + f"Either give '{ch}' an add_messages reducer (only valid for " + f"type=messages) or rename one of the writer outputs." + ), + ) + ) + return issues + + +def _rule_11_reserved_names(graph: GraphRecord) -> list[ValidationIssue]: + issues: list[ValidationIssue] = [] + for nid in graph.nodes: + if nid in _RESERVED_NAMES or nid.startswith("_aiperf"): + issues.append( + ValidationIssue( + rule_id="rule-11", + severity=ValidationSeverity.ERROR, + location=f"graph.nodes.{nid}", + message=( + f"Node id '{nid}' is reserved (matches START/END or " + f"begins with '_aiperf')." + ), + suggested_fix="Rename the node.", + ) + ) + return issues + + +def _rule_12_version(graph: GraphRecord) -> list[ValidationIssue]: + if graph.version in _KNOWN_VERSIONS: + return [] + return [ + ValidationIssue( + rule_id="rule-12", + severity=ValidationSeverity.ERROR, + location="graph.version", + message=( + f"Unknown major version {graph.version!r}; known: " + f"{sorted(_KNOWN_VERSIONS)}." + ), + suggested_fix=f"Set version to one of {sorted(_KNOWN_VERSIONS)}.", + ) + ] + + +def _rule_13_provenance_tool(graph: GraphRecord) -> list[ValidationIssue]: + p = graph.provenance + if p.source == "hand-authored": + return [] + tool = (p.tool or "").strip() + if not tool: + return [ + ValidationIssue( + rule_id="rule-13", + severity=ValidationSeverity.ERROR, + location="graph.provenance.tool", + message=( + f"provenance.source={p.source!r} requires a non-empty " + f"provenance.tool (origin tool plus version)." + ), + suggested_fix=( + "Set provenance.tool to '/', or change " + "source to 'hand-authored'." + ), + ) + ] + if tool == "manual": + # "manual" is the field default; an adapter-emitted graph still + # carrying it never stamped its generating tool. + return [ + ValidationIssue( + rule_id="rule-13", + severity=ValidationSeverity.WARNING, + location="graph.provenance.tool", + message=( + f"provenance.source={p.source!r} carries the default " + f"provenance.tool 'manual'; an adapter-emitted graph " + f"should stamp its generating tool plus version." + ), + suggested_fix=( + "Have the generating adapter set provenance.tool to " + "'/', or change source to " + "'hand-authored'." + ), + ) + ] + return [] + + +def _rule_15_expected_cached_tokens(graph: GraphRecord) -> list[ValidationIssue]: + """WARNING: emit when an LLM node sets expected.cache_read_tokens or + expected.cache_creation_tokens. The engine may not report cache fields; we + cannot know at file-load time whether the configured engine supports them. + Surface as a warning so users see it.""" + issues: list[ValidationIssue] = [] + cache_fields = ("cache_read_tokens", "cache_creation_tokens") + for nid, node in graph.nodes.items(): + if not isinstance(node, LlmNode) or node.expected is None: + continue + for field in cache_fields: + if getattr(node.expected, field) is None: + continue + issues.append( + ValidationIssue( + rule_id="rule-15", + severity=ValidationSeverity.WARNING, + location=f"graph.nodes.{nid}.expected.{field}", + message=( + f"Node '{nid}' sets expected.{field} but not all engines " + f"report {field}; this assertion will be skipped if the " + f"engine doesn't surface that field." + ), + ) + ) + return issues + + +def _rule_21_unreachable(graph: GraphRecord) -> list[ValidationIssue]: + adj = _build_adjacency(graph) + seen: set[str] = set() + q: deque[str] = deque([START_NODE_ID]) + while q: + u = q.popleft() + if u in seen: + continue + seen.add(u) + q.extend(adj.get(u, ())) + seen.discard(START_NODE_ID) + seen.discard(END_NODE_ID) + declared = set(graph.nodes.keys()) + unreachable = declared - seen + issues: list[ValidationIssue] = [] + for nid in sorted(unreachable): + issues.append( + ValidationIssue( + rule_id="rule-21", + severity=ValidationSeverity.ERROR, + location=f"graph.nodes.{nid}", + message=f"Node '{nid}' is unreachable from START.", + suggested_fix=( + f"Add an edge `{{from: , to: {nid}}}` or remove the node." + ), + ) + ) + return issues + + +def _rule_54_edge_delay_exclusivity(graph: GraphRecord) -> list[ValidationIssue]: + """Start-anchored edge legality (two checks, one rule). + + 1. An edge must not set both ``delay_after_predecessor_us`` and + ``delay_after_predecessor_start_us``: the first anchors the successor to + the predecessor's COMPLETION, the second to its DISPATCH; honoring both + on one edge has no recorded-trace meaning. + 2. A ``START``-sourced edge must not be start-anchored. The ``START`` + pseudo-node never dispatches, so the runtime routes such a target to + ``Scheduler._start_anchored_succ`` where ``entry_nodes()`` never sees it + and ``start_anchored_successors("START")`` is never consulted -- the + target would be silently orphaned. Fail loudly instead of scheduling a + node that can never fire. + """ + issues: list[ValidationIssue] = [] + for edge in graph.edges: + if not isinstance(edge, StaticEdge): + continue + if ( + edge.delay_after_predecessor_us is not None + and edge.delay_after_predecessor_start_us is not None + ): + issues.append( + ValidationIssue( + rule_id="rule-54", + severity=ValidationSeverity.ERROR, + location=f"graph.edges[{edge.source}->{edge.target}]", + message=( + f"edge {edge.source!r} -> {edge.target!r} sets both " + "delay_after_predecessor_us and " + "delay_after_predecessor_start_us; they are mutually " + "exclusive." + ), + suggested_fix=( + "Keep exactly one anchor: end-anchored " + "(delay_after_predecessor_us) or start-anchored " + "(delay_after_predecessor_start_us)." + ), + ) + ) + if ( + edge.source == START_NODE_ID + and edge.delay_after_predecessor_start_us is not None + ): + issues.append( + ValidationIssue( + rule_id="rule-54", + severity=ValidationSeverity.ERROR, + location=f"graph.edges[{edge.source}->{edge.target}]", + message=( + f"edge {edge.source!r} -> {edge.target!r} is " + "start-anchored (delay_after_predecessor_start_us), but " + "START-sourced edges cannot be start-anchored; the START " + "pseudo-node never dispatches, so the target would never " + "be scheduled." + ), + suggested_fix=( + "Use min_start_delay_us for an absolute offset from trace " + "start instead of delay_after_predecessor_start_us." + ), + ) + ) + return issues + + +def _rule_55_first_token_anchor_shape(graph: GraphRecord) -> list[ValidationIssue]: + """A first-token-anchored edge must carry its dispatch fallback and a real + source: ``delay_after_predecessor_first_token_us`` requires + ``delay_after_predecessor_start_us`` on the same edge (the runtime falls + back to dispatch + start delay when the predecessor never streams a first + token), must not combine with the completion anchor + ``delay_after_predecessor_us``, and cannot source at START (the START + pseudo-node never dispatches or streams).""" + issues: list[ValidationIssue] = [] + for edge in graph.edges: + if not isinstance(edge, StaticEdge): + continue + if edge.delay_after_predecessor_first_token_us is None: + continue + if edge.delay_after_predecessor_start_us is None: + issues.append( + ValidationIssue( + rule_id="rule-55", + severity=ValidationSeverity.ERROR, + location=f"graph.edges[{edge.source}->{edge.target}]", + message=( + f"edge {edge.source!r} -> {edge.target!r} is " + "first-token-anchored (delay_after_predecessor_first_token_us) " + "but sets no delay_after_predecessor_start_us; the runtime " + "needs the start delay as the dispatch fallback when the " + "predecessor terminates without a first token." + ), + suggested_fix=( + "Add delay_after_predecessor_start_us to this edge as the " + "dispatch-anchored fallback for delay_after_predecessor_first_token_us." + ), + ) + ) + if edge.delay_after_predecessor_us is not None: + issues.append( + ValidationIssue( + rule_id="rule-55", + severity=ValidationSeverity.ERROR, + location=f"graph.edges[{edge.source}->{edge.target}]", + message=( + f"edge {edge.source!r} -> {edge.target!r} sets both " + "delay_after_predecessor_us and " + "delay_after_predecessor_first_token_us; the completion " + "anchor and the first-token anchor are mutually exclusive." + ), + suggested_fix=( + "Drop delay_after_predecessor_us and keep " + "delay_after_predecessor_first_token_us with its " + "delay_after_predecessor_start_us fallback." + ), + ) + ) + if edge.source == START_NODE_ID: + issues.append( + ValidationIssue( + rule_id="rule-55", + severity=ValidationSeverity.ERROR, + location=f"graph.edges[{edge.source}->{edge.target}]", + message=( + f"edge {edge.source!r} -> {edge.target!r} is " + "first-token-anchored (delay_after_predecessor_first_token_us), " + "but START-sourced edges cannot be first-token-anchored; the " + "START pseudo-node never dispatches or streams a first token." + ), + suggested_fix=( + "Use min_start_delay_us for an absolute offset from trace " + "start instead of delay_after_predecessor_first_token_us." + ), + ) + ) + return issues + + +def _rule_56_edge_endpoints(graph: GraphRecord) -> list[ValidationIssue]: + """Every edge endpoint must be a declared node or the matching sentinel. + + A dangling endpoint (typo'd node id, `START` as a target, `END` as a + source) otherwise validates clean and produces a node that never fires or + an edge the scheduler silently ignores. + """ + declared = set(graph.nodes) + issues: list[ValidationIssue] = [] + for edge in graph.edges: + if not isinstance(edge, StaticEdge): + continue + loc = f"graph.edges[{edge.source}->{edge.target}]" + if edge.source not in declared and edge.source != START_NODE_ID: + issues.append( + ValidationIssue( + rule_id="rule-56", + severity=ValidationSeverity.ERROR, + location=loc, + message=( + f"edge source '{edge.source}' is neither a declared node " + f"nor the START sentinel." + ), + suggested_fix=( + "Declare the node in graph.nodes or fix the edge " + "source (only 'START' is a valid non-node source)." + ), + ) + ) + if edge.target not in declared and edge.target != END_NODE_ID: + issues.append( + ValidationIssue( + rule_id="rule-56", + severity=ValidationSeverity.ERROR, + location=loc, + message=( + f"edge target '{edge.target}' is neither a declared node " + f"nor the END sentinel." + ), + suggested_fix=( + "Declare the node in graph.nodes or fix the edge " + "target (only 'END' is a valid non-node target)." + ), + ) + ) + return issues + + +_EDGE_DELAY_FIELDS = ( + "delay_after_predecessor_us", + "min_start_delay_us", + "delay_after_predecessor_start_us", + "delay_after_predecessor_first_token_us", +) + + +def _rule_57_finite_delays(graph: GraphRecord) -> list[ValidationIssue]: + """Every delay value must be finite. + + ``msgspec.Meta(ge=0)`` admits ``+inf``, and the executor gates a successor + at ``finish + delay`` — an infinite delay hangs the trace for the whole + benchmark. The loose decoder rejects non-finite values at decode time + (see ``decode.py``); this rule catches graphs constructed directly from + typed structs (adapter output, programmatic builders). + """ + issues: list[ValidationIssue] = [] + for edge in graph.edges: + if not isinstance(edge, StaticEdge): + continue + for field in _EDGE_DELAY_FIELDS: + value = getattr(edge, field) + if value is None or math.isfinite(value): + continue + issues.append( + ValidationIssue( + rule_id="rule-57", + severity=ValidationSeverity.ERROR, + location=f"graph.edges[{edge.source}->{edge.target}].{field}", + message=( + f"edge {edge.source!r} -> {edge.target!r} sets " + f"{field}={value!r}; delay values must be finite (a " + f"non-finite gate never clears and hangs the trace)." + ), + suggested_fix=( + f"Set {field} to a finite microsecond value or drop it." + ), + ) + ) + for nid, node in graph.nodes.items(): + value = node.min_start_delay_us + if value is None or math.isfinite(value): + continue + issues.append( + ValidationIssue( + rule_id="rule-57", + severity=ValidationSeverity.ERROR, + location=f"graph.nodes.{nid}.min_start_delay_us", + message=( + f"node '{nid}' sets min_start_delay_us={value!r}; delay " + f"values must be finite (a non-finite gate never clears " + f"and hangs the trace)." + ), + suggested_fix=( + "Set min_start_delay_us to a finite microsecond value or drop it." + ), + ) + ) + return issues + + +__all__ = [ + "ValidationIssue", + "ValidationSeverity", + "validate", +] diff --git a/src/aiperf/dataset/graph/workload_detect.py b/src/aiperf/dataset/graph/workload_detect.py new file mode 100644 index 0000000000..9f06b48c90 --- /dev/null +++ b/src/aiperf/dataset/graph/workload_detect.py @@ -0,0 +1,494 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Graph-workload detection + parsing seam for the dataset plane. + +Graph traces are NOT a ``custom_dataset_loader`` plugin -- they are ingested via +the ``graph_adapter`` plugin registry (``weka_trace`` / ``dynamo_trace`` / ...) +into a :class:`~aiperf.dataset.graph.models.ParsedGraph`. Auto-detection +walks that registry (excluding ``native``, which is explicit-``--graph-format`` only) and +parse dispatch is registry-driven too -- :func:`parse_graph_workload` resolves +every run-derived parse knob into ONE :class:`GraphParseContext` +(:func:`resolve_graph_parse_context`) and every format goes through the generic +:func:`~aiperf.dataset.graph.parser.parse_graph`, whose adapters map the ctx +fields they consume via the uniform ``parse(path, ctx)`` protocol. + +The **DatasetManager** (build plane) is the ONLY caller that parses: it builds +the ``ParsedGraph``, serializes the per-node envelopes into the graph store +mmap, and writes + broadcasts the content-free graph_meta sidecar. The +**TimingManager** (schedule plane) still calls :func:`resolve_graph_workload` +(detection only, from the run config alone) to decide whether the input is a +graph workload and to flip the profiling phase's timing mode to ``GRAPH_IR``, +but it ingests the broadcast sidecar rather than parsing. Deriving both the +build-time and the dispatch-time node ordinals (catalog) from that single +parse is what keeps the worker reading the right envelope. + +Detection runs AT MOST ONCE per process: :func:`resolve_graph_workload` is the +single accessor every consumer calls; it reads the +``run.resolved.graph_workload`` memo (populated eagerly by the config resolver +chain in single-run mode) or derives-and-memoizes on first access for runs +that never pass the chain (``aiperf service`` processes, test-built runs). +""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from aiperf.dataset.graph.parse_context import GraphParseContext + +if TYPE_CHECKING: + from aiperf.config.resolution.plan import BenchmarkRun, GraphWorkloadRef + from aiperf.dataset.graph.models import ParsedGraph + +__all__ = [ + "GraphEndpointUnsupportedError", + "is_graph_workload_path", + "parse_graph_workload", + "publish_graph_loader_tokenizer_env", + "resolve_graph_parse_context", + "resolve_graph_workload", + "validate_graph_endpoint_type", +] + + +class GraphEndpointUnsupportedError(ValueError): + """The run targets an endpoint type the graph workload dispatch cannot serve. + + The graph dispatch path materializes a chat-completions body + (``{"messages": [...], "max_completion_tokens": N, "stream": bool}``) and + sends it verbatim, bypassing the endpoint's ``format_payload`` reshaping. + Pointing that body at a non-chat endpoint (``completions`` expects + ``prompt``; ``embeddings`` expects ``input``; etc.) makes every request fail + with a server-side 422 and no actionable up-front error. This is raised at + configure time so the run rejects cleanly before any request is dispatched. + """ + + +def validate_graph_endpoint_type(run: BenchmarkRun) -> None: + """Reject non-chat endpoint types for a graph workload run. + + The graph dispatch path materializes a chat-completions body + (``{"messages": [...], "max_completion_tokens": N, "stream": bool}``) and + sends it verbatim, bypassing the endpoint's ``format_payload`` reshaping. + Only an endpoint whose path is the chat-completions path can serve that + body; any other type (``completions`` expects ``prompt``, ``embeddings`` + expects ``input``, ...) would fail every request with a server-side 422. + Chat-compatibility is keyed on the endpoint metadata's ``endpoint_path`` + ending in ``/chat/completions`` so any future chat-completions endpoint + plugin passes without an allowlist edit. The DatasetManager runs this gate + at configure time, before the store is built, so the failure is a clean + configure-time rejection. + + Raises: + GraphEndpointUnsupportedError: When the run's endpoint type is not a + chat-completions endpoint. + """ + from aiperf.plugin import plugins + + endpoint_type = run.cfg.endpoint.type + endpoint_path = plugins.get_endpoint_metadata(endpoint_type).endpoint_path + if endpoint_path is None or not endpoint_path.endswith("/chat/completions"): + raise GraphEndpointUnsupportedError( + f"graph workload: --endpoint-type '{endpoint_type}' is not " + f"supported (endpoint_path={endpoint_path!r}). The graph " + f"replay path emits a chat-completions body and only supports a " + f"chat-completions endpoint (e.g. --endpoint-type chat). Re-run " + f"with a chat endpoint, or convert the trace for the target " + f"endpoint shape." + ) + + +# ``native`` is the explicit ``--graph-format`` / ``parse_graph`` fallback whose +# ``can_load`` matches any ``.yaml``/``.yml``/``.jsonl``; excluding it here keeps +# plain conversation datasets on the linear pipeline instead of hijacking them +# into graph mode. ``dag_jsonl`` is excluded so legacy +# ``--custom-dataset-type dag_jsonl`` runs stay on the linear pipeline -- the +# graph adapter is opt-in via ``--graph-format dag_jsonl`` only (its +# ``can_load`` stays implemented + tested for a future autodetect flip). +_AUTODETECT_EXCLUDED = frozenset({"native", "dag_jsonl"}) + + +def _detect_graph_workload_format(path: Path) -> str | None: + """Return the graph-adapter format name auto-detected for ``path``, else None. + + Walks the ``graph_adapter`` plugin registry, skipping the explicit-only + adapters in :data:`_AUTODETECT_EXCLUDED` (``native`` and ``dag_jsonl``), and + returns the highest-``detection_priority`` adapter whose ``can_load`` matches. + Returns ``None`` when no trace adapter recognizes the file -- the caller then + keeps the linear (non-graph) pipeline. + """ + from aiperf.plugin import plugins + from aiperf.plugin.enums import PluginType + from aiperf.plugin.schema.schemas import GraphAdapterMetadata + + matches: list[tuple[int, str]] = [] + for entry, cls in plugins.iter_all(PluginType.GRAPH_ADAPTER): + if entry.name in _AUTODETECT_EXCLUDED: + continue + if cls.can_load(path): + meta = entry.get_typed_metadata(GraphAdapterMetadata) + matches.append((meta.detection_priority, entry.name)) + return max(matches)[1] if matches else None + + +def _graph_format_override(run: BenchmarkRun) -> str | None: + """Return the explicit ``--graph-format`` override for ``run``, or None. + + Reads ``FileDataset.graph_format``. When set, the input is treated as a graph + workload of exactly this adapter regardless of auto-detection (this is the + only path that can select ``native``, which auto-detection excludes). + """ + try: + dataset = run.cfg.get_default_dataset() + except (IndexError, AttributeError): + return None + fmt = getattr(dataset, "graph_format", None) + return str(fmt) if fmt is not None else None + + +def resolve_graph_workload(run: BenchmarkRun) -> GraphWorkloadRef | None: + """Graph workload for this run, detected AT MOST ONCE per process. + + Reads the memoized resolution when present (populated by the config + resolver chain in single-run mode, or by a prior call); otherwise runs + override-or-autodetect ONCE and memoizes onto ``run.resolved``. Detection + failures degrade to None exactly like the per-consumer sniffs this + accessor replaced (``can_load`` returns False on unreadable inputs), so + ``aiperf service`` processes and test-built runs behave identically to the + status quo -- minus the repeated registry walks and file I/O. + """ + # STRICT ``is True``: a truthy check would auto-truthify on MagicMock runs + # (the repo's documented MagicMock-path-drift trap). + if run.resolved.graph_workload_resolved is True: + return run.resolved.graph_workload + ref = _derive_graph_workload(run) + # Two threads racing the first call can both derive; the double-detect is + # benign (idempotent same-value memoization, GIL-atomic assignments), so + # no lock is taken. + run.resolved.graph_workload = ref + run.resolved.graph_workload_resolved = True + return ref + + +def _derive_graph_workload(run: BenchmarkRun) -> GraphWorkloadRef | None: + """Run override-or-autodetect once; None when the input is not a graph. + + Reads the default dataset's ``path``. When ``--graph-format`` is set + (:func:`_graph_format_override`), the input is FORCED to be a graph + workload of that adapter -- including ``native``, which auto-detection + excludes. Otherwise asks the ``graph_adapter`` registry (via + :func:`_detect_graph_workload_format`) whether any trace adapter + recognizes the file. Returns ``None`` for any non-file dataset (synthetic + / public / inline), a file no adapter recognizes, AND any derivation + failure: the swallow consumers used to wrap around detection lives HERE + now, so callers read the accessor's result bare. + """ + from aiperf.config.resolution.plan import GraphWorkloadRef + + try: + try: + dataset = run.cfg.get_default_dataset() + except (IndexError, AttributeError): + return None + raw_path = getattr(dataset, "path", None) + if raw_path is None: + return None + path = Path(raw_path) + fmt = _graph_format_override(run) + if fmt is None: + fmt = _detect_graph_workload_format(path) + if fmt is None: + return None + return GraphWorkloadRef(path=path, format=fmt) + except Exception: + return None + + +def is_graph_workload_path(path: Path) -> bool: + """True when `path` is a graph workload recognized by a trace adapter. + + Path-level companion to :func:`resolve_graph_workload` (which takes a + run). Uses the SAME registry detection (`_detect_graph_workload_format`, + which excludes `native`), so callers skip exactly what the graph pipeline + will claim. + """ + return _detect_graph_workload_format(path) is not None + + +def publish_graph_loader_tokenizer_env(run: BenchmarkRun) -> None: + """Publish the run's tokenizer trust/revision for graph content synthesis. + + ``CorpusContentSynthesizer._build_generator`` (and the forkserver preload) + reads ``AIPERF_LOADER_PRELOAD_TRUST_REMOTE_CODE`` / ``_REVISION`` at + tokenizer-load time -- the parallel loader threads only the tokenizer NAME. + :func:`parse_graph_workload` is the one seam every graph parse goes through, + so it must publish the run's tokenizer trust/revision triple itself: a + direct caller (tooling, tests) has no DatasetManager configure step to do + it, and the forkserver helper snapshots the env once at spawn. Idempotent: + re-publishing the same run-derived values is a no-op. + """ + from aiperf.dataset._mp_context import configure_loader_tokenizer_env + + tokenizer = run.cfg.tokenizer + configure_loader_tokenizer_env( + trust_remote_code=tokenizer.trust_remote_code, + revision=tokenizer.revision, + ) + + +def resolve_graph_parse_context(run: BenchmarkRun) -> GraphParseContext: + """Resolve EVERY run-derived graph parse knob into one :class:`GraphParseContext`. + + The ONE knob spelling for the run -> parse seam: every field is populated + with the run's RESOLVED value verbatim, so a registry-dispatched + ``adapter.parse(path, ctx)`` is byte-identical to the run's own parse for + every format. Each adapter maps only the ctx fields it consumes and + ignores the rest (see :class:`GraphParseContext`), so resolving dag knobs + for a weka run (and vice versa) is harmless -- every resolver here is a + pure function of the run config, safe to evaluate for every format. + + ``idle_gap_cap_seconds`` carries :func:`_resolve_graph_idle_gap_cap`'s + result AS-IS: 60.0 default / explicit float / ``None`` for the user's + explicit ``synthesis.idle_gap_cap_seconds: null`` (warping DISABLED). It + is never ``UNSET`` -- a run always has a resolved answer -- and the weka + and dynamo adapters forward the tri-state verbatim. + + Lazy imports: ``timing.config`` avoids the ``timing`` <-> graph loader + import cycle; the rest keep this module light at import time. + """ + from aiperf.common.models.model_endpoint_info import ModelEndpointInfo + from aiperf.dataset.loader.dag_jsonl import _resolve_delay_cap_seconds + from aiperf.timing.config import ( + resolve_graph_content_seed, + resolve_graph_content_tokenizer, + ) + + tokenizer = run.cfg.tokenizer + endpoint_info = ModelEndpointInfo.from_run(run) + return GraphParseContext( + content_root_seed=resolve_graph_content_seed(run), + content_tokenizer=resolve_graph_content_tokenizer(run), + tokenizer_trust_remote_code=tokenizer.trust_remote_code, + tokenizer_revision=tokenizer.revision, + prompt_corpus=_resolve_graph_corpus(run), + max_osl=_resolve_graph_max_osl(run), + num_dataset_entries=_resolve_graph_num_entries(run), + max_context_length=_resolve_graph_max_context(run), + idle_gap_cap_seconds=_resolve_graph_idle_gap_cap(run), + trajectory_start_max_ratio=run.cfg.trajectory_start_max_ratio or 0.0, + default_model=endpoint_info.primary_model_name, + run_streaming=endpoint_info.endpoint.streaming, + delay_cap_seconds=_resolve_delay_cap_seconds(run), + endpoint_extra=endpoint_info.endpoint.extra, + ) + + +def parse_graph_workload( + run: BenchmarkRun, path: Path | str, **adapter_kwargs: object +) -> ParsedGraph: + """Parse a graph workload into a ``ParsedGraph`` for ``run``, via the registry. + + ONE dispatch for every format: resolve the run's parse knobs into a + :class:`GraphParseContext` (:func:`resolve_graph_parse_context`) and hand + it to :func:`~aiperf.dataset.graph.parser.parse_graph`, which routes to the + selected adapter's uniform ``parse(path, ctx)``. Each adapter maps only the + ctx fields it consumes (weka: seed / tokenizer / corpus / max_osl / + idle-gap cap / selection caps; dynamo: the same minus max_osl; + dag_jsonl: the four dispatch knobs; + native: none), so a parse is fully determined by the run config + file for + every format. The DatasetManager is the ONLY production caller; the + TimingManager ingests the graph_meta sidecar this build writes and + broadcasts, and never parses. + + ``**adapter_kwargs`` is the build-plane's seam for ADAPTER-SPECIFIC live + objects the run-derived :class:`GraphParseContext` cannot carry -- currently + only the dynamo direct route's ``direct_store`` (the ``GraphStoreBuilder`` + constructs the unified store BEFORE the parse and passes it here so + ``build_trie_ir``'s ``pool.add`` write-throughs intern straight into the + store). It is forwarded verbatim to :func:`parse_graph`, which fails loud + (``TypeError`` for an unknown-to-the-adapter kwarg, ``GraphParseError`` for + any kwarg with ``fmt == "native"``); a caller that passes no adapter kwargs + is byte-identical to the pre-seam call. + + Adapter failures surface uniformly as + :class:`~aiperf.dataset.graph.parser.GraphParseError` (the registry seam + re-wraps adapter ``ValueError`` subclasses, message text preserved), so + callers need exactly one except class. + + The t*/dynamic-slot gate runs uniformly on the parsed result; + all-explicit-zero arrival offsets (the shape dag_jsonl lowering stamps and + guards at its own parse seam) take the carve-out documented in + :func:`_gate_dynamic_slots_vs_tstar`. + """ + publish_graph_loader_tokenizer_env(run) + p = Path(path) + ref = resolve_graph_workload(run) + # The ref's format describes the RUN'S OWN dataset input; parsing a + # divergent path with that format would silently mis-route a future + # caller, so pin the congruence here. + if ref is not None and p != ref.path: + raise ValueError( + f"parse_graph_workload path {p} diverges from the run's resolved " + f"graph workload {ref.path}; pass the run's own dataset path" + ) + fmt = ref.format if ref is not None else None + + from aiperf.dataset.graph.parser import parse_graph + + ctx = resolve_graph_parse_context(run) + parsed = parse_graph(p, format=fmt, ctx=ctx, **adapter_kwargs) + _gate_dynamic_slots_vs_tstar(parsed, ctx.trajectory_start_max_ratio) + return parsed + + +def _gate_dynamic_slots_vs_tstar( + parsed: ParsedGraph, trajectory_start_max_ratio: float +) -> None: + """Reject dynamic slots + an engaged t* snapshot window (unsupported by design). + + The t* chop partitions nodes into warmup/profiled by recorded offsets; a + slot producer chopped into warmup would leave its consumer's pool value + undefined. The gate runs on the one build-plane parse seam (the + TimingManager ingests the sidecar of a build that already passed it). The + window is off by default; it engages only via + ``--scenario inferencex-agentx-mvp`` or explicit + ``--trajectory-start-min/max-ratio`` flags, hence the message names both + exits (drop the scenario or the flags) -- an explicit + ``--trajectory-start-max-ratio 0`` would collide with a scenario-applied + window and raise ``ScenarioLockError`` instead. + + Resolves :func:`graph_carries_assembly_slots` -- this t*-gate is the + predicate's sole production caller -- so + there is ONE definition of "carries dynamic slots"; the + predicate's union of ``assembly`` and ``capture`` is equivalent to an + assembly-only check at graph level (capture is only ever stamped on a + producer referenced by some assembly program). Lazy import matches this + module's local-import style and avoids the graph loader import cycle. + """ + from aiperf.dataset.graph.segment_ir.store_builder import ( + graph_carries_assembly_slots, + ) + + if not graph_carries_assembly_slots(parsed): + return + if trajectory_start_max_ratio <= 0.0: + return + # EXPLICIT-ZERO carve-out: skip the t*-rejection iff EVERY node's + # ``arrival_offset_us`` is explicitly ``0`` -- int zero; ``None`` (the + # un-stamped default on natively authored nodes) does NOT qualify and + # keeps gating exactly as before. The skip is load-bearing on an + # INVARIANT, not on cross-plane determinism: dag_jsonl lowering stamps + # ``arrival_offset_us=0`` on EVERY node, so such a trace's recorded + # duration is 0 and any sampled t* = ratio * 0 = 0. At t*=0 the snapshot + # chop drops nothing (no node's offset is < 0), so the chop -- and the + # "slot producer chopped into warmup leaves its consumer's pool value + # undefined" hazard this gate exists to reject -- is a structural no-op; + # rejecting would ONLY ever false-positive on the scenario-applied + # ``trajectory_start_max_ratio=1.0`` window. ``DagJsonlGraphAdapter.parse`` guards + # the all-zero invariant at its own seam + # (``_assert_dag_zero_arrival_offsets``); if dag ever emits recorded + # offsets, that guard raises AND this carve-out stops matching, so the + # gate re-engages by construction. Contract delta (c), intentional: a + # NATIVE graph AUTHORED with all-zero offsets now passes where it used to + # reject -- the t*-degeneracy invariant is identical there, so the old + # rejection was the same false positive. + if all( + node.arrival_offset_us == 0 + for record in (parsed.graph, *parsed.graphs.values()) + for node in record.nodes.values() + ): + return + raise ValueError( + "graph workload carries dynamic content slots (prompt refs to " + "LlmNode-written channels); the t* snapshot window is not " + "supported with dynamic slots. The window was engaged by " + "--scenario inferencex-agentx-mvp or explicit " + "--trajectory-start-min/max-ratio flags; drop the scenario (or the " + "flags) to run a full native replay." + ) + + +def _synthesis_attr(run: BenchmarkRun, name: str, default: Any) -> Any: + """Read an attribute off the run's default-dataset synthesis config. + + Tolerates an absent default dataset or an absent ``synthesis`` block: both + yield ``default``. A ``getattr`` default (not ``or``) is returned so a + present-but-``None`` attribute is preserved by the caller, not coalesced. + """ + dataset = run.cfg.get_default_dataset() + synthesis = getattr(dataset, "synthesis", None) if dataset else None + return getattr(synthesis, name, default) if synthesis else default + + +def _resolve_graph_max_osl(run: BenchmarkRun) -> int | None: + """Resolve the ``--synthesis-max-osl`` cap for the run, or ``None``. + + Reads ``synthesis.max_osl`` off the run's default dataset (where + ``base_trace_loader`` reads it for the linear path). Resolving from the run + config alone keeps the cap identical for every parse of the run + (in-process build or spawn-started pool worker). ``None`` (flag unset) + leaves the recorded ``out`` uncapped. + """ + return _synthesis_attr(run, "max_osl", None) + + +def _resolve_graph_num_entries(run: BenchmarkRun) -> int | None: + """Resolve the explicit ``entries`` cap on the run's default dataset, or ``None``. + + Reads ``entries`` off the run's default dataset, gated on + ``"entries" in ds.model_fields_set`` so ONLY a user-set value is a cap: a + dataset class's own default (``FileDataset.entries=None``, or the linear + ``SyntheticDataset.entries=100``) is not a graph-plane selection ceiling and + resolves to ``None`` (use all eligible traces). Resolving from the run + config alone keeps the cap identical for every parse of the run + (in-process build or spawn-started pool worker) -- and, unlike + ``DatasetResolver._resolve_one``, + this seam is never skipped by the weka-HF ``org/name`` or local-graph + early-returns. + """ + dataset = run.cfg.get_default_dataset() + if not dataset or "entries" not in getattr(dataset, "model_fields_set", set()): + return None + return getattr(dataset, "entries", None) + + +def _resolve_graph_max_context(run: BenchmarkRun) -> int | None: + """Resolve the ``--max-context-length`` per-trace context cap, or ``None``. + + Reads ``synthesis.max_context_length`` off the run's default dataset, + mirroring :func:`_resolve_graph_max_osl`. Resolving from the run config + alone keeps the cap identical for every parse of the run. ``None`` (flag + unset) applies no context-length filter. + """ + return _synthesis_attr(run, "max_context_length", None) + + +def _resolve_graph_corpus(run: BenchmarkRun) -> str: + """Resolve the weka content corpus (`--prompt-corpus`) for the run. + + Reads ``synthesis.corpus`` off the run's default dataset (where + ``--prompt-corpus`` lands via the dataset converter), defaulting to + ``"coding"`` -- the corpus the recorded weka workloads were captured against. + Resolving from the run config alone keeps the corpus identical for every + parse of the run. ``"sonnet"`` selects the Shakespeare + pool. The trailing ``or "coding"`` also coalesces a present-but-empty/``None`` + corpus to the default (a getattr default alone would not). + """ + return _synthesis_attr(run, "corpus", None) or "coding" + + +# Default per-trace idle-gap cap when no synthesis config is present (60s, the +# value the recorded weka workloads were captured/replayed against). +_GRAPH_IDLE_GAP_CAP_DEFAULT = 60.0 + + +def _resolve_graph_idle_gap_cap(run: BenchmarkRun) -> float | None: + """Resolve the per-trace idle-gap cap (`--synthesis-idle-gap-cap`) for the run. + + Reads ``synthesis.idle_gap_cap_seconds`` off the run's default dataset + (default 60.0 when the field is unset, ``None`` when explicitly set to null to + disable warping). Falls back to :data:`_GRAPH_IDLE_GAP_CAP_DEFAULT` when no + synthesis config is present. Resolving from the run config alone keeps the + cap identical for every parse of the run. + """ + return _synthesis_attr(run, "idle_gap_cap_seconds", _GRAPH_IDLE_GAP_CAP_DEFAULT) diff --git a/src/aiperf/dataset/graph_client_store.py b/src/aiperf/dataset/graph_client_store.py new file mode 100644 index 0000000000..e3981454af --- /dev/null +++ b/src/aiperf/dataset/graph_client_store.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Client-store plugin face for graph IR runs. + +Graph credits address the unified segment store by ``(trace_id, +node_ordinal)`` through the worker's ``GraphSegmentUnifiedClient`` -- there +are no conversations to serve. This class exists so the worker's generic +client-store construction (``plugins.get_class(DATASET_CLIENT_STORE, +client_metadata.client_type)``) has a registered target for +``GraphSegmentClientMetadata``; conversation lookups are an explicit hard +error, never a silent empty result. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from aiperf.common.exceptions import InvalidOperationError +from aiperf.common.mixins import AIPerfLifecycleMixin +from aiperf.common.models.dataset_models import GraphSegmentClientMetadata + +if TYPE_CHECKING: + from aiperf.common.models.dataset_models import Conversation + + +class GraphSegmentDatasetClientStore(AIPerfLifecycleMixin): + """No-op lifecycle wrapper carrying the graph store locations.""" + + def __init__(self, client_metadata: GraphSegmentClientMetadata, **kwargs) -> None: + super().__init__(**kwargs) + self.client_metadata = client_metadata + + async def get_conversation(self, conversation_id: str) -> Conversation: + """Graph runs have no conversation store; always raises.""" + raise InvalidOperationError( + f"graph runs have no conversation store (requested " + f"{conversation_id!r}); graph payloads are addressed by " + "(trace_id, node_ordinal) in the unified segment store" + ) diff --git a/src/aiperf/dataset/graph_segment_unified_store.py b/src/aiperf/dataset/graph_segment_unified_store.py new file mode 100644 index 0000000000..eb82685bdc --- /dev/null +++ b/src/aiperf/dataset/graph_segment_unified_store.py @@ -0,0 +1,534 @@ +from __future__ import annotations + +import array as _array +import contextlib +import mmap +from dataclasses import dataclass +from io import BufferedWriter +from pathlib import Path + +import aiofiles +import orjson + +from aiperf.common.constants import IS_MACOS, IS_WINDOWS +from aiperf.common.exceptions import ( + MemoryMapFileOperationError, + MemoryMapSerializationError, +) + +_UNIFIED_DIR_PREFIX = "aiperf_graph_segments_" +_CONTENT_BLOB = "content.blob" +_CONTENT_IDX = "content.idx" +_NODES_BLOB = "nodes.blob" +_NODES_IDX = "nodes.idx" + +_HEAD = b'{"messages":[' + +_IDX_TYPECODE = "Q" # 64-bit unsigned; offsets/sizes fit comfortably + + +def _unified_dir(base_path: Path, benchmark_id: str) -> Path: + return Path(base_path) / f"{_UNIFIED_DIR_PREFIX}{benchmark_id}" + + +def _encode_inner_key(node_ordinal: int, phase_variant: str) -> str: + """The store's (ordinal, variant) inner key, ':' -- + self-consistent between the writer and GraphSegmentUnifiedClient.""" + return f"{node_ordinal}:{phase_variant}" + + +@dataclass(slots=True) +class NodeEnvelope: + """A single per-(trace, node) manifest envelope to be persisted.""" + + node_ordinal: int + """Zero-based ordinal of the node within its trace.""" + + phase_variant: str + """Phase variant this envelope belongs to (e.g. ``"profiling"``).""" + + envelope_bytes: bytes + """Pre-serialized ``orjson.dumps`` envelope blob to store verbatim.""" + + +@dataclass(slots=True, frozen=True) +class GraphStoreBuildStats: + """One store-build memory snapshot, computed at :meth:`finalize` entry. + + A cheap (``O(traces) + O(1)``) derivation from the buffers the store already + holds -- the measurement baseline that makes pool/envelope-size regressions + visible in the build log instead of surfacing as mystery RSS. + """ + + segment_count: int + """Number of unique interned content segments (``len(self._spans)``).""" + + content_bytes: int + """Total bytes of the interned content blob. + + Post-spill this is the running ``_content_bytes_written`` counter (the + encoded content is streamed straight to ``content.blob`` at ``put`` time and + never accumulated in RAM), which equals the finalized ``content.blob`` size + and the sum of every span's length.""" + + node_manifest_count: int + """Total per-node manifest envelopes indexed across all traces.""" + + manifest_bytes: int + """Total bytes appended to the manifest region (``len(self._nodes_buf)``).""" + + trace_count: int + """Number of distinct traces carrying at least one manifest.""" + + peak_rss_mib: float | None + """Process peak resident set size in MiB at finalize entry, or ``None`` where + unavailable (Windows). ``RUSAGE_SELF`` only -- excludes weka pool workers.""" + + +def _peak_rss_mib() -> float | None: + """Process peak resident set size in MiB, or ``None`` where unavailable. + + Windows has no ``resource`` module, so return ``None`` there (the import is + lazy, after the early return, so this module still imports on Windows). + ``ru_maxrss`` units are platform-specific: macOS (Darwin) reports BYTES, + Linux reports KiB -- convert each to MiB with the right divisor. + """ + if IS_WINDOWS: + return None + import resource + + max_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + if IS_MACOS: + return max_rss / (1024.0 * 1024.0) + return max_rss / 1024.0 + + +class GraphSegmentUnifiedBackingStore: + """Build-time writer for the unified store: an interned (A2) content pool + (blob + packed span index; the hex->handle map lives only in this writer) + AND a per-node manifest region (concatenated envelope blobs + a + (trace, ordinal:variant) offset index). + + Content spills incrementally: :meth:`put_segment` streams each segment's + wire blob straight to an open ``content.blob`` write handle and advances a + running ``_content_bytes_written`` counter, so the encoded content is never + accumulated in a RAM ``bytearray`` and :meth:`finalize` no longer + materializes a ``bytes(self._content_buf)`` transient double. Only the + ``_ids`` map (dedup + handle resolution) and the ``_spans`` list + (``content.idx`` is written from it) stay resident; both are released at + finalize END. The handle is plain buffered binary file I/O with NO + event-loop coupling, so the store may be constructed on one thread and + drained / finalized inside ``asyncio.run`` on a worker thread (see + ``GraphStoreBuilder``). + + ``_nodes_buf`` is deliberately NOT spilled: the manifest region fills in the + drain window (below the parse peak) and is small (~0.2-0.3 GB even at 1M + nodes), so streaming it would add a second write handle for no meaningful + peak-RAM win. It stays a RAM ``bytearray`` flushed once at finalize. + + A store that errors before finalize would otherwise leave a half-written + ``content.blob`` on disk (the incremental spill writes as it goes); + :meth:`abort` closes the handle and unlinks the store files so a later + :class:`GraphSegmentUnifiedClient` open never trips on a partial blob.""" + + def __init__(self, base_path: Path | str, benchmark_id: str) -> None: + d = _unified_dir(Path(base_path), benchmark_id) + d.mkdir(parents=True, exist_ok=True) + self._dir = d + self._content_blob_path = d / _CONTENT_BLOB + self._content_idx_path = d / _CONTENT_IDX + self._nodes_blob_path = d / _NODES_BLOB + self._nodes_idx_path = d / _NODES_IDX + # content region: put_segment streams each blob to _content_blob_file and + # advances _content_bytes_written; the encoded content never sits in RAM. + self._ids: dict[str, int] = {} + self._spans: list[tuple[int, int]] = [] + self._content_blob_file: BufferedWriter | None = open( # noqa: SIM115 + self._content_blob_path, "wb" + ) + self._content_bytes_written = 0 + self._content_buf = bytearray() # vestigial: never appended to post-spill + # nodes region: trace_id -> inner_key -> (offset, size) + self._nodes_buf = bytearray() + self._node_offsets: dict[str, dict[str, list[int]]] = {} + self._finalized = False + self._build_stats: GraphStoreBuildStats | None = None + + @property + def data_dir(self) -> Path: + """The store directory (``aiperf_graph_segments_``).""" + return self._dir + + def put_segment( + self, segment_id: str, role: str, content: str, wire_json: str | None = None + ) -> int: + """Intern one segment's wire blob; return its dense insertion-index handle. + + When ``wire_json`` is provided (a raw-authored dag message) the persisted + blob is ``wire_json.encode()`` verbatim -- key order and extra keys preserved + byte-for-byte. Otherwise the blob is the derived ``{"role", "content"}`` dict + (the existing normalized behavior for token/text segments). + """ + if self._finalized: + raise RuntimeError("Cannot put_segment after finalize") + existing = self._ids.get(segment_id) + if existing is not None: + return existing + b = ( + wire_json.encode() + if wire_json is not None + else orjson.dumps({"role": role, "content": content}) + ) + # Spill: write the blob straight to disk and advance the running byte + # counter. off is the pre-write counter -- identical to the old + # len(self._content_buf) -- so content.blob bytes and content.idx spans + # are byte-identical to the accumulate-then-flush build. + off = self._content_bytes_written + assert self._content_blob_file is not None + self._content_blob_file.write(b) + self._content_bytes_written += len(b) + handle = len(self._spans) + self._ids[segment_id] = handle + self._spans.append((off, len(b))) + return handle + + def segment_handle(self, segment_id: str) -> int | None: + if self._finalized: + raise RuntimeError( + "Cannot segment_handle after finalize: the writer's id map is " + "released once finalize flushes the store; read handles via " + "GraphSegmentUnifiedClient instead." + ) + return self._ids.get(segment_id) + + def add_node_manifest( + self, + trace_id: str, + node_ordinal: int, + phase_variant: str, + envelope_bytes: bytes, + ) -> None: + if self._finalized: + raise RuntimeError("Cannot add_node_manifest after finalize") + off = len(self._nodes_buf) + self._nodes_buf += envelope_bytes + key = _encode_inner_key(node_ordinal, phase_variant) + self._node_offsets.setdefault(trace_id, {})[key] = [off, len(envelope_bytes)] + + def add_node_manifest_interned( + self, + trace_id: str, + node_ordinal: int, + phase_variant: str, + handles: list[int], + dispatch_overrides: dict, + stream: bool, + *, + items: list[dict] | None = None, + capture: bool = False, + extra_headers: dict[str, str] | None = None, + endpoint_extra_applied: bool = False, + ) -> None: + """Write one node's manifest envelope. + + ``items``/``capture`` are the dynamic-content additions: + ``items`` is the ordered assembly program for slot-carrying + nodes (``{"h": handle}`` / ``{"s": {"src": ordinal}}`` / + ``{"m": {"role", "parts"}}``), ``capture`` marks producer nodes whose + responses the worker pools. ``extra_headers`` carries per-node HTTP + headers (dynamo ``x-dynamo-*`` session identity) the worker attaches to + the request HEADERS, never the body. ``endpoint_extra_applied`` marks a + node whose adapter already folded the run's ``--extra-inputs`` into + ``dispatch_overrides`` at parse, so the worker must NOT re-merge + ``endpoint.extra`` (the adapter-owned values win). All four are OMITTED + when unset, so envelopes for header-less / flag-less corpora (weka, + static native, dynamo) stay byte-identical. + """ + envelope: dict = { + "handles": list(handles), + "dispatch_overrides": dispatch_overrides, + "stream": stream, + } + if items is not None: + envelope["items"] = items + if capture: + envelope["capture"] = True + if extra_headers: + envelope["extra_headers"] = dict(extra_headers) + if endpoint_extra_applied: + envelope["endpoint_extra_applied"] = True + self.add_node_manifest( + trace_id, node_ordinal, phase_variant, orjson.dumps(envelope) + ) + + def _compute_build_stats(self) -> GraphStoreBuildStats: + """Snapshot the store's build-memory footprint from its live buffers. + + ``O(traces) + O(1)``: two ``len()`` reads, one ``len()`` over the trace + map, and one ``len()`` per trace's inner map. No pass over content and no + per-trace maxima -- the buffers already hold the running totals. + """ + return GraphStoreBuildStats( + segment_count=len(self._spans), + content_bytes=self._content_bytes_written, + node_manifest_count=sum(len(v) for v in self._node_offsets.values()), + manifest_bytes=len(self._nodes_buf), + trace_count=len(self._node_offsets), + peak_rss_mib=_peak_rss_mib(), + ) + + @property + def build_stats(self) -> GraphStoreBuildStats | None: + """The build-memory snapshot, or ``None`` until :meth:`finalize` runs. + + ``manifest_bytes`` and the counters reflect APPENDED (write-side) totals: + a duplicate ``(trace, ordinal, variant)`` write orphans the earlier blob + in ``_nodes_buf`` while ``node_manifest_count`` tracks live index entries, + so a future count/bytes divergence reads as that duplicate-write bug. + ``content_bytes`` now consumes the running ``_content_bytes_written`` + counter the incremental spill maintains (the seam this docstring used to + name as future work). ``manifest_bytes`` stays a ``len(self._nodes_buf)`` + read because the manifest region is not spilled. The counter is + monotonic, so the finalize-ENTRY snapshot still measures the full content + footprint even though the blob was streamed to disk during the write. + """ + return self._build_stats + + def abort(self) -> None: + """Best-effort teardown for a store that errored before finalize. + + The incremental spill streams ``content.blob`` as it goes, so an aborted + build would otherwise leave a partial blob for a later + :class:`GraphSegmentUnifiedClient` open to trip on. Close the write + handle and unlink the four store files. Idempotent and non-raising: safe + to call twice, and safe after a SUCCESSFUL finalize (the ``_finalized`` + guard skips the unlink so a complete store is never deleted). The + ``GraphStoreBuilder`` drain paths call this, then remove the store dir, + before re-raising the drain error. + """ + handle = self._content_blob_file + self._content_blob_file = None + if handle is not None: + with contextlib.suppress(OSError): + handle.close() + if self._finalized: + # A fully finalized store is complete on disk -- never unlink it. + return + for path in ( + self._content_blob_path, + self._content_idx_path, + self._nodes_blob_path, + self._nodes_idx_path, + ): + with contextlib.suppress(OSError): + path.unlink() + + async def finalize(self) -> None: + if self._finalized: + raise RuntimeError("finalize called twice") + # Snapshot BEFORE any remaining write so a mid-write failure still leaves + # the full-footprint measure here (the running counter is monotonic). + self._build_stats = self._compute_build_stats() + # content.blob was streamed to disk incrementally by put_segment; flush + # and close the handle instead of writing bytes(self._content_buf) -- the + # finalize transient double is gone. + if self._content_blob_file is not None: + self._content_blob_file.flush() + self._content_blob_file.close() + self._content_blob_file = None + async with aiofiles.open(self._content_idx_path, "wb") as f: + flat = _array.array(_IDX_TYPECODE) + for off, size in self._spans: + flat.append(off) + flat.append(size) + await f.write(flat.tobytes()) + async with aiofiles.open(self._nodes_blob_path, "wb") as f: + await f.write(bytes(self._nodes_buf)) + async with aiofiles.open(self._nodes_idx_path, "wb") as f: + await f.write(orjson.dumps(self._node_offsets)) + self._finalized = True + # Release the write-side accumulation state now that every file is + # flushed. The store object lives on through the structural merge / + # sidecar / prefix-cache tail with zero readers of these buffers + # (build_stats was snapshotted at entry above). Post-spill the + # load-bearing clears are _ids (dedup map) and _spans (content.idx + # source), which are what still pin RAM through the post-finalize tail; + # _content_buf is already empty (content spilled to disk at put time) so + # its clear is a formality. The put/add/handle methods guard on + # _finalized, so nothing touches them post-release. + self._ids = {} + self._spans = [] + self._content_buf = bytearray() + self._nodes_buf = bytearray() + self._node_offsets = {} + + +class GraphSegmentUnifiedClient: + """Worker-side reader for the unified store. Presents BOTH faces -- + per-node envelope addressing and interned content -- so worker_materialize + is handed this ONE object for both. A2-strict: only the interned packed + ``content.idx`` is accepted; a legacy JSON (A1) index fails loud.""" + + def __init__(self, base_path: Path | str, benchmark_id: str) -> None: + d = _unified_dir(Path(base_path), benchmark_id) + self._data_dir = d + self._content_blob_path = d / _CONTENT_BLOB + self._content_idx_path = d / _CONTENT_IDX + self._nodes_blob_path = d / _NODES_BLOB + self._nodes_idx_path = d / _NODES_IDX + self._content_mm: mmap.mmap | None = None + self._content_mv: memoryview | None = None + self._spans: list[list[int]] = [] + self._nodes_mm: mmap.mmap | None = None + self._nodes_mv: memoryview | None = None + self._node_offsets: dict[str, dict[str, list[int]]] = {} + self._opened = False + + @property + def data_dir(self) -> Path: + return self._data_dir + + def _read_or_raise(self, path: Path) -> bytes: + try: + return path.read_bytes() + except OSError as e: + raise MemoryMapFileOperationError(f"missing {path}") from e + + def _load_content_idx(self) -> None: + raw = self._read_or_raise(self._content_idx_path) + # A2 (packed) is raw 'Q' bytes; a legacy A1 (hex) index is a JSON object + # beginning with b'{' and is no longer readable. + if raw[:1] == b"{": + raise ValueError( + f"legacy non-interned unified store (pre-v3) at " + f"{self._content_idx_path}; re-parse required" + ) + flat = _array.array(_IDX_TYPECODE) + flat.frombytes(raw) + self._spans = [[flat[i], flat[i + 1]] for i in range(0, len(flat), 2)] + + def _load_idx(self, path: Path) -> dict: + try: + data = path.read_bytes() + except OSError as e: + raise MemoryMapFileOperationError(f"missing {path}") from e + try: + return orjson.loads(data) + except orjson.JSONDecodeError as e: + raise MemoryMapSerializationError(str(e)) from e + + def _map_blob(self, path: Path) -> tuple[mmap.mmap | None, memoryview | None]: + if not path.exists(): + raise MemoryMapFileOperationError(f"missing {path}") + if path.stat().st_size == 0: + return None, None + fh = path.open("rb") + mm = mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ) + fh.close() + return mm, memoryview(mm) + + def _validate_content_spans(self) -> None: + """Reject spans past the end of ``content.blob`` (mirrors the nodes-region + check in :meth:`get_node_envelope`) -- Python slice clamping would + otherwise return silently truncated bytes from a stale/partial store.""" + content_len = 0 if self._content_mv is None else len(self._content_mv) + for handle, (off, size) in enumerate(self._spans): + if off + size > content_len: + raise MemoryMapSerializationError( + f"content handle {handle}: offset {off}+{size} exceeds " + f"content.blob ({content_len})" + ) + + def open(self) -> GraphSegmentUnifiedClient: + try: + self._load_content_idx() + self._node_offsets = self._load_idx(self._nodes_idx_path) + self._content_mm, self._content_mv = self._map_blob(self._content_blob_path) + self._nodes_mm, self._nodes_mv = self._map_blob(self._nodes_blob_path) + self._validate_content_spans() + except BaseException: + self.close() + raise + self._opened = True + return self + + def _require_opened(self, operation: str) -> None: + """Reject reads on an unopened client with an actionable error.""" + if not self._opened: + raise RuntimeError( + f"GraphSegmentUnifiedClient.{operation}: unified store at " + f"{self._data_dir} is not opened; call open() (or use the " + "client as a context manager) before reading." + ) + + # --- addressing face (per-node envelope reads) --- + def get_node_envelope( + self, trace_id: str, node_ordinal: int, phase_variant: str = "profiling" + ) -> bytes | None: + self._require_opened("get_node_envelope") + trace_offsets = self._node_offsets.get(trace_id) + if trace_offsets is None: + return None + info = trace_offsets.get(_encode_inner_key(node_ordinal, phase_variant)) + if info is None: + return None + off, size = info + end = off + size + nodes_len = 0 if self._nodes_mv is None else len(self._nodes_mv) + if self._nodes_mv is None or end > nodes_len: + raise MemoryMapSerializationError( + f"{trace_id!r} node {node_ordinal} ({phase_variant!r}): " + f"offset {off}+{size} exceeds nodes.blob ({nodes_len})" + ) + return bytes(self._nodes_mv[off:end]) + + # --- interned (A2) int-handle face --- + def materialize_handles(self, handles: list[int]) -> list[dict[str, str]]: + self._require_opened("materialize_handles") + out: list[dict[str, str]] = [] + for h in handles: + if h < 0 or h >= len(self._spans): + raise MemoryMapSerializationError(f"unknown handle {h!r}") + assert self._content_mv is not None + off, size = self._spans[h] + out.append(orjson.loads(self._content_mv[off : off + size])) + return out + + def build_request_body_handles( + self, handles: list[int], overrides_inner: bytes + ) -> bytes: + self._require_opened("build_request_body_handles") + parts: list[bytes | memoryview] = [_HEAD] + first = True + for h in handles: + if h < 0 or h >= len(self._spans): + raise MemoryMapSerializationError(f"unknown handle {h!r}") + assert self._content_mv is not None + off, size = self._spans[h] + if not first: + parts.append(b",") + parts.append(self._content_mv[off : off + size]) + first = False + parts.append(b"]}" if not overrides_inner else b"]," + overrides_inner + b"}") + return b"".join(parts) + + def close(self) -> None: + for mv_attr, mm_attr in ( + ("_content_mv", "_content_mm"), + ("_nodes_mv", "_nodes_mm"), + ): + mv = getattr(self, mv_attr) + mm = getattr(self, mm_attr) + if mv is not None: + mv.release() + setattr(self, mv_attr, None) + if mm is not None: + mm.close() + setattr(self, mm_attr, None) + self._opened = False + + def __enter__(self) -> GraphSegmentUnifiedClient: + return self.open() if not self._opened else self + + def __exit__(self, *exc) -> None: + self.close() diff --git a/src/aiperf/dataset/loader/base_hf_dataset.py b/src/aiperf/dataset/loader/base_hf_dataset.py index 93afcfb916..62f1e7b422 100644 --- a/src/aiperf/dataset/loader/base_hf_dataset.py +++ b/src/aiperf/dataset/loader/base_hf_dataset.py @@ -178,7 +178,9 @@ def _extract_audio(self, row: dict[str, Any], audio_column: str) -> list[Audio]: return [Audio(name="", contents=[f"wav,{b64}"])] except (ImportError, OSError, ValueError, RuntimeError) as e: self.debug( - lambda exc=e: f"Failed to encode WAV from column '{audio_column}': {exc.__class__.__name__}: {exc}" + lambda exc=e: ( + f"Failed to encode WAV from column '{audio_column}': {exc.__class__.__name__}: {exc}" + ) ) return [] diff --git a/src/aiperf/dataset/loader/exgentic.py b/src/aiperf/dataset/loader/exgentic.py index 4316f7abdd..2670032ed7 100644 --- a/src/aiperf/dataset/loader/exgentic.py +++ b/src/aiperf/dataset/loader/exgentic.py @@ -360,7 +360,9 @@ def _parse_span( + _normalize_messages(attributes.get("gen_ai.input.messages")), raw_tools=_normalize_tools(attributes.get("gen_ai.tool.definitions")), extra_body=_request_extra_body(attributes), - extra_headers={"x-dynamo-session-id": session_id}, + # Session identity headers are no longer auto-stamped here: run + # with --session-routing dynamo_headers to stamp + # x-dynamo-session-id (now available for ANY dataset/endpoint). ) except (KeyError, TypeError, ValueError) as error: raise DatasetLoaderError( diff --git a/src/aiperf/dataset/loader/hf_asr.py b/src/aiperf/dataset/loader/hf_asr.py index 83f02b54c5..61de90e365 100644 --- a/src/aiperf/dataset/loader/hf_asr.py +++ b/src/aiperf/dataset/loader/hf_asr.py @@ -74,7 +74,9 @@ def _audio_from_bytes(self, audio_value: _HFAudioBytesRow) -> list[Audio]: return [Audio(name="", contents=[f"wav,{b64}"])] except (ImportError, OSError, ValueError, RuntimeError) as e: self.debug( - lambda exc=e: f"Failed to decode audio bytes: {exc.__class__.__name__}: {exc}" + lambda exc=e: ( + f"Failed to decode audio bytes: {exc.__class__.__name__}: {exc}" + ) ) return [] @@ -90,7 +92,9 @@ def _duration_seconds(self, audio_value: _HFAudioBytesRow) -> float | None: return info.duration except (ImportError, OSError, ValueError, RuntimeError) as e: self.debug( - lambda exc=e: f"Failed to estimate audio duration: {exc.__class__.__name__}: {exc}" + lambda exc=e: ( + f"Failed to estimate audio duration: {exc.__class__.__name__}: {exc}" + ) ) return None @@ -150,6 +154,8 @@ async def convert_to_conversations( f"Check that '{self.audio_column}' contains valid audio data." ) self.debug( - lambda: f"Converted {len(conversations)} rows (skipped {skipped} empty/long)" + lambda: ( + f"Converted {len(conversations)} rows (skipped {skipped} empty/long)" + ) ) return conversations diff --git a/src/aiperf/dataset/loader/sharegpt.py b/src/aiperf/dataset/loader/sharegpt.py index 255f56a11f..a5695c0343 100644 --- a/src/aiperf/dataset/loader/sharegpt.py +++ b/src/aiperf/dataset/loader/sharegpt.py @@ -148,7 +148,9 @@ async def convert_to_conversations( ) self.debug( - lambda: f"Filtered to {len(filtered_dataset)} dataset entries out of {len(dataset)} (skipped {skipped_entries})" + lambda: ( + f"Filtered to {len(filtered_dataset)} dataset entries out of {len(dataset)} (skipped {skipped_entries})" + ) ) return filtered_dataset diff --git a/src/aiperf/dataset/memory_map_utils.py b/src/aiperf/dataset/memory_map_utils.py index 806ab3f7a5..02f97b5067 100644 --- a/src/aiperf/dataset/memory_map_utils.py +++ b/src/aiperf/dataset/memory_map_utils.py @@ -326,12 +326,16 @@ async def _setup(self) -> None: """Open memory-mapped files (read-only).""" self._loop = asyncio.get_running_loop() self.debug( - lambda: f"Opening memory-mapped files: data={self._data_path}, index={self._index_path}" + lambda: ( + f"Opening memory-mapped files: data={self._data_path}, index={self._index_path}" + ) ) self._client = MemoryMapDatasetClient(self._data_path, self._index_path) self.debug( - lambda: f"Memory-mapped client store initialized with " - f"{len(self._client.index.conversation_ids)} conversations" + lambda: ( + f"Memory-mapped client store initialized with " + f"{len(self._client.index.conversation_ids)} conversations" + ) ) async def get_conversation(self, conversation_id: str) -> Conversation: @@ -373,7 +377,7 @@ class ConversationOffset(AIPerfBaseModel): class MemoryMapDatasetIndex(AIPerfBaseModel): """Index structure for the memory-mapped dataset. - All data is stored as uncompressed JSON bytes serialized with orjson. + All data is stored as uncompressed JSON bytes (Pydantic ``model_dump_json``). """ conversation_ids: list[str] = Field( @@ -465,7 +469,9 @@ def __init__(self, data_file_path: Path | str, index_file_path: Path | str) -> N ) _logger.debug( - lambda: f"MemoryMapDatasetClient initialized successfully: data_file={self.data_file_path}, index_file={self.index_file_path}, conversations={len(self.index.conversation_ids)}, size={self.index.total_size} bytes" + lambda: ( + f"MemoryMapDatasetClient initialized successfully: data_file={self.data_file_path}, index_file={self.index_file_path}, conversations={len(self.index.conversation_ids)}, size={self.index.total_size} bytes" + ) ) def __enter__(self) -> "MemoryMapDatasetClient": @@ -546,7 +552,9 @@ def get_conversation(self, conversation_id: str) -> Conversation: conv_bytes = self.data_mmap.read(offset_info.size) _logger.debug( - lambda: f"Loading conversation '{conversation_id}': offset={offset_info.offset}, size={offset_info.size} bytes" + lambda: ( + f"Loading conversation '{conversation_id}': offset={offset_info.offset}, size={offset_info.size} bytes" + ) ) return self._deserialize_conversation(conv_bytes) diff --git a/src/aiperf/exporters/aggregate/aggregate_base_exporter.py b/src/aiperf/exporters/aggregate/aggregate_base_exporter.py index 723ed735a7..021ad987a3 100644 --- a/src/aiperf/exporters/aggregate/aggregate_base_exporter.py +++ b/src/aiperf/exporters/aggregate/aggregate_base_exporter.py @@ -10,8 +10,57 @@ import aiofiles from aiperf.common.mixins import AIPerfLoggerMixin +from aiperf.common.scenario.submission_outcome import ( + CONTEXT_OVERFLOW_REASON, + RUN_CANCELLED_REASON, + compute_submission_outcome, +) from aiperf.orchestrator.aggregation.base import AggregateResult +__all__ = [ + "CONTEXT_OVERFLOW_REASON", + "RUN_CANCELLED_REASON", + "AggregateBaseExporter", + "AggregateExporterConfig", + "_build_run_metadata_dict", + "compute_submission_outcome", +] + + +def _build_run_metadata_dict( + *, + scenario_name: str | None, + submission_valid: bool | None, + submission_invalid_reasons: list[str] | None = None, +) -> dict: + """Build the scenario-submission sub-dict for the aggregate export. + + Mirrors the single-run ``RunInfo`` submission surface for the multi-run + (``--num-profile-runs > 1`` / sweep) aggregate export. Returns an empty + dict when ``scenario_name`` is ``None`` so non-scenario runs are not + polluted with submission-tracking fields. When ``scenario_name`` is set, + returns the ``scenario`` name plus a coerced ``submission_valid`` bool, and + includes ``submission_invalid_reasons`` only when that list is non-empty. + + Args: + scenario_name: Active scenario identifier, or ``None`` for a + non-scenario run. + submission_valid: Whether the run is a valid scenario submission. + Coerced to ``bool`` (``None`` becomes ``False``) when emitted. + submission_invalid_reasons: Machine-readable reason codes (e.g. + ``"unsafe_override"``, ``"context_overflow_rate_exceeded"``). + + Returns: + A dict suitable for merging into the top-level aggregate JSON output. + """ + md: dict = {} + if scenario_name is not None: + md["scenario"] = scenario_name + md["submission_valid"] = bool(submission_valid) + if submission_invalid_reasons: + md["submission_invalid_reasons"] = list(submission_invalid_reasons) + return md + @dataclass(slots=True) class AggregateExporterConfig: diff --git a/src/aiperf/exporters/aggregate/aggregate_confidence_csv_exporter.py b/src/aiperf/exporters/aggregate/aggregate_confidence_csv_exporter.py index cf10af3391..d7cc72070e 100644 --- a/src/aiperf/exporters/aggregate/aggregate_confidence_csv_exporter.py +++ b/src/aiperf/exporters/aggregate/aggregate_confidence_csv_exporter.py @@ -55,15 +55,28 @@ def _generate_content(self) -> str: def _write_metadata_section(self, writer: csv.writer) -> None: """Write metadata section to CSV. + The scenario-submission carrier keys (``_scenario_name`` etc.) are + internal writer->reader plumbing consumed by the JSON exporter's + submission fold; they are stripped here exactly as the JSON exporter + pops them, so they never leak into user-facing output. + Args: writer: CSV writer object """ + # Deferred import: cli_runner is the writer of the carrier-key + # contract; importing it lazily avoids pulling the CLI runner stack + # in at exporter module-import time. + from aiperf.cli_runner._aggregate import ( + strip_scenario_submission_carrier_keys, + ) + writer.writerow(["Aggregation Type", self._result.aggregation_type]) writer.writerow(["Total Runs", self._result.num_runs]) writer.writerow(["Successful Runs", self._result.num_successful_runs]) # Add custom metadata - for key, value in self._result.metadata.items(): + metadata = strip_scenario_submission_carrier_keys(self._result.metadata) + for key, value in metadata.items(): writer.writerow([key.replace("_", " ").title(), value]) def _write_metrics_section(self, writer: csv.writer) -> None: diff --git a/src/aiperf/exporters/aggregate/aggregate_confidence_json_exporter.py b/src/aiperf/exporters/aggregate/aggregate_confidence_json_exporter.py index 85c2797293..b9667b22eb 100644 --- a/src/aiperf/exporters/aggregate/aggregate_confidence_json_exporter.py +++ b/src/aiperf/exporters/aggregate/aggregate_confidence_json_exporter.py @@ -82,6 +82,11 @@ def _aggregate_to_export_data(self) -> "JsonExportData": aiperf_version=aiperf_version, ) + # ``result_metadata`` is consumed (carrier keys popped) by the submission + # fold below; the survivors flow into the public aggregate metadata. + result_metadata = dict(self._result.metadata) + run_metadata = self._build_submission_metadata(result_metadata) + # Add aggregate-specific metadata as extra field # (JsonExportData has extra="allow" to support this) aggregate_metadata = { @@ -89,7 +94,8 @@ def _aggregate_to_export_data(self) -> "JsonExportData": "num_profile_runs": self._result.num_runs, "num_successful_runs": self._result.num_successful_runs, "failed_runs": self._result.failed_runs, - **self._result.metadata, + **result_metadata, + **run_metadata, } export_data.metadata = aggregate_metadata @@ -118,3 +124,144 @@ def _aggregate_to_export_data(self) -> "JsonExportData": export_data.metrics = metrics_dict return export_data + + def _build_submission_metadata(self, result_metadata: dict) -> dict: + """Fold the aggregate (multi-run) scenario-submission verdict. + + Mirrors the single-run ``MetricsJsonExporter._fold_runtime_submission_outcome`` + for the ``--num-profile-runs > 1`` / sweep path: the scenario-lock + outcome (validator) is combined ACROSS RUNS with the cross-run + context-overflow rate and cancellation via ``compute_submission_outcome`` + (InferenceX AgentX RFC §7). Null-safe: a non-scenario run carries no + ``scenario_name`` and yields an empty dict (no submission fields). + + Args: + result_metadata: A mutable copy of ``self._result.metadata`` whose + underscore-prefixed carrier keys are popped in place so they do + not leak into the public aggregate metadata. + + Returns: + The ``scenario`` / ``submission_valid`` / ``submission_invalid_reasons`` + sub-dict to merge into the aggregate output, or ``{}`` for a + non-scenario run. + """ + from aiperf.exporters.aggregate.aggregate_base_exporter import ( + _build_run_metadata_dict, + compute_submission_outcome, + ) + + scenario_name = self._pop_scenario_name(result_metadata) + validator_submission_valid = result_metadata.pop( + "_validator_submission_valid", None + ) + validator_reasons = result_metadata.pop( + "_validator_submission_invalid_reasons", None + ) + total_responses, context_overflow_count = self._cross_run_response_counts( + result_metadata + ) + was_cancelled = bool(result_metadata.pop("_was_cancelled", False)) + + submission_valid, submission_invalid_reasons = compute_submission_outcome( + scenario_name=scenario_name, + validator_submission_valid=validator_submission_valid, + validator_reasons=validator_reasons, + total_responses=total_responses, + context_overflow_count=context_overflow_count, + was_cancelled=was_cancelled, + ) + return _build_run_metadata_dict( + scenario_name=scenario_name, + submission_valid=submission_valid, + submission_invalid_reasons=submission_invalid_reasons, + ) + + def _pop_scenario_name(self, result_metadata: dict) -> str | None: + """Resolve the active scenario name for the aggregate submission fold. + + Prefers the orchestrator-stamped ``_scenario_name`` carrier key (popped + so it does not leak into the public metadata), falling back to the + ``scenario`` key the aggregation strategy may already carry. Returns + ``None`` for a non-scenario run, which short-circuits the whole + submission fold to no output. + + Args: + result_metadata: A mutable copy of ``self._result.metadata`` whose + underscore-prefixed carrier keys are popped in place. + + Returns: + The scenario name, or ``None`` when no scenario is active. + """ + carrier = result_metadata.pop("_scenario_name", None) + if carrier is not None: + return str(carrier) + existing = result_metadata.get("scenario") + return str(existing) if existing is not None else None + + def _cross_run_response_counts(self, result_metadata: dict) -> tuple[int, int]: + """Resolve ``(total_responses, context_overflow_count)`` across runs. + + Prefers the orchestrator-stamped ``_total_responses`` / + ``_context_overflow_count`` carrier keys (summed across runs upstream), + popping them so they do not leak into the public metadata. When those + carriers are absent, derives the counts from the confidence aggregate's + own per-run-mean count metrics (``request_count_avg``, + ``error_request_count_avg``, ``context_overflow_count_avg``). + + The overflow RATE is rate-equivalent under both sourcing modes: the + cross-run mean and the cross-run sum share the same run count, so + ``mean(overflow) / mean(total) == sum(overflow) / sum(total)`` and + ``compute_submission_outcome`` derives the identical verdict. The + denominator is ``request_count + error_request_count + overflow`` per + the InferenceX AgentX spec §4.8 / §7 (all responses received). + + Args: + result_metadata: A mutable copy of ``self._result.metadata`` whose + underscore-prefixed carrier keys are popped in place. + + Returns: + A ``(total_responses, context_overflow_count)`` tuple of + non-negative ints. + """ + carrier_total = result_metadata.pop("_total_responses", None) + carrier_overflow = result_metadata.pop("_context_overflow_count", None) + if carrier_total is not None or carrier_overflow is not None: + return ( + int(carrier_total or 0), + int(carrier_overflow or 0), + ) + + overflow = self._aggregate_metric_mean("context_overflow_count") + total = ( + self._aggregate_metric_mean("request_count") + + self._aggregate_metric_mean("error_request_count") + + overflow + ) + return total, overflow + + def _aggregate_metric_mean(self, metric_tag: str) -> int: + """Return a count metric's cross-run mean from the aggregate, or 0. + + The confidence aggregation flattens each metric to ``{tag}_{stat}`` + keys; the per-run average lands on ``{tag}_avg``. A ``ConfidenceMetric`` + exposes that value as ``.mean``; a non-confidence fallback object may + expose ``.avg``. Missing or non-numeric -> 0. + + Args: + metric_tag: The base metric tag (e.g. ``"request_count"``). + + Returns: + The non-negative integer cross-run mean, or 0 when absent. + """ + metric = self._result.metrics.get(f"{metric_tag}_avg") + if metric is None: + return 0 + value = getattr(metric, "mean", None) + if value is None: + value = getattr(metric, "avg", None) + if value is None: + return 0 + try: + return max(0, int(value)) + except (TypeError, ValueError): + return 0 diff --git a/src/aiperf/exporters/metrics_csv_exporter.py b/src/aiperf/exporters/metrics_csv_exporter.py index e1e32deabe..130d045157 100644 --- a/src/aiperf/exporters/metrics_csv_exporter.py +++ b/src/aiperf/exporters/metrics_csv_exporter.py @@ -28,7 +28,9 @@ def __init__(self, exporter_config: ExporterConfig, **kwargs) -> None: self._percentile_keys = _percentile_keys_from(STAT_KEYS) self.trace_or_debug( lambda: f"Initializing MetricsCsvExporter with config: {exporter_config}", - lambda: f"Initializing MetricsCsvExporter with file path: {self._file_path}", + lambda: ( + f"Initializing MetricsCsvExporter with file path: {self._file_path}" + ), ) def get_export_info(self) -> FileExportInfo: diff --git a/src/aiperf/exporters/metrics_json_exporter.py b/src/aiperf/exporters/metrics_json_exporter.py index a88ab2b304..2da1649360 100644 --- a/src/aiperf/exporters/metrics_json_exporter.py +++ b/src/aiperf/exporters/metrics_json_exporter.py @@ -34,7 +34,9 @@ def __init__(self, exporter_config: ExporterConfig, **kwargs) -> None: self._file_path = exporter_config.cfg.artifacts.profile_export_json_file self.trace_or_debug( lambda: f"Initializing MetricsJsonExporter with config: {exporter_config}", - lambda: f"Initializing MetricsJsonExporter with file path: {self._file_path}", + lambda: ( + f"Initializing MetricsJsonExporter with file path: {self._file_path}" + ), ) def get_export_info(self) -> FileExportInfo: @@ -89,6 +91,15 @@ def _generate_content(self) -> str: for metric_tag, json_result in prepared_json_metrics.items(): setattr(export_data, metric_tag, json_result) + # Fold the runtime context-overflow-rate contribution (InferenceX + # AgentX RFC §7) into the lock-only ``submission_valid`` surfaced by + # ``RunInfo.from_run``. ``apply_scenario`` stamps lock violations + + # ``unsafe_override``; ``compute_submission_outcome`` additionally flips + # the verdict when ``context_overflow_count / total_responses`` exceeds + # the configured threshold (or the run was cancelled). Null-safe: only + # runs out of a scenario carry a non-None ``scenario_name``. + self._fold_runtime_submission_outcome(export_data, prepared_json_metrics) + # Splice DAG branch orchestration counters when present. Non-DAG # runs leave ``branch_stats`` unset on ProfileResults so the # section is omitted entirely (model_dump_json with @@ -113,6 +124,57 @@ def _generate_content(self) -> str: scrub_non_finite(payload), option=orjson.OPT_INDENT_2 ).decode("utf-8") + def _fold_runtime_submission_outcome( + self, + export_data: JsonExportData, + prepared_json_metrics: dict[str, JsonMetricResult], + ) -> None: + """Fold the runtime overflow-rate contribution into ``run_info``. + + ``RunInfo.from_run`` surfaces the lock-only scenario outcome + (invariant violations + ``--unsafe-override``). This re-derives the + final ``submission_valid`` / ``submission_invalid_reasons`` via + ``compute_submission_outcome``, which additionally flips the verdict + when the runtime context-overflow rate exceeds + ``Environment.AGENTX.CONTEXT_OVERFLOW_RATE_LIMIT`` (RFC §7) or the run + was cancelled. No-op for non-scenario runs (``scenario_name`` None). + + Args: + export_data: The export envelope whose ``run_info`` is mutated + in place. + prepared_json_metrics: Tag -> JsonMetricResult map; aggregate + counters expose their value via ``.avg``. + """ + from aiperf.common.scenario import compute_submission_outcome + + run_info = export_data.run_info + if run_info is None or run_info.scenario_name is None: + return + + def _metric_avg(tag: str) -> int: + m = prepared_json_metrics.get(tag) + if m is None or m.avg is None: + return 0 + return int(m.avg) + + context_overflow_count = _metric_avg("context_overflow_count") + total_responses = ( + _metric_avg("request_count") + + _metric_avg("error_request_count") + + context_overflow_count + ) + + submission_valid, submission_invalid_reasons = compute_submission_outcome( + scenario_name=run_info.scenario_name, + validator_submission_valid=run_info.submission_valid, + validator_reasons=run_info.submission_invalid_reasons, + total_responses=total_responses, + context_overflow_count=context_overflow_count, + was_cancelled=bool(self._results.was_cancelled), + ) + run_info.submission_valid = submission_valid + run_info.submission_invalid_reasons = submission_invalid_reasons or None + def _prepare_metrics_for_json( self, metric_results: Iterable[MetricResult] ) -> dict[str, JsonMetricResult]: diff --git a/src/aiperf/exporters/timeslice_metrics_csv_exporter.py b/src/aiperf/exporters/timeslice_metrics_csv_exporter.py index 161aca6107..9da5f52d49 100644 --- a/src/aiperf/exporters/timeslice_metrics_csv_exporter.py +++ b/src/aiperf/exporters/timeslice_metrics_csv_exporter.py @@ -39,8 +39,12 @@ def __init__(self, exporter_config: ExporterConfig, **kwargs) -> None: exporter_config.cfg.artifacts.profile_export_timeslices_csv_file ) self.trace_or_debug( - lambda: f"Initializing TimesliceMetricsCsvExporter with config: {exporter_config}", - lambda: f"Initializing TimesliceMetricsCsvExporter with file path: {self._file_path}", + lambda: ( + f"Initializing TimesliceMetricsCsvExporter with config: {exporter_config}" + ), + lambda: ( + f"Initializing TimesliceMetricsCsvExporter with file path: {self._file_path}" + ), ) def get_export_info(self) -> FileExportInfo: diff --git a/src/aiperf/exporters/timeslice_metrics_json_exporter.py b/src/aiperf/exporters/timeslice_metrics_json_exporter.py index e3fe3040d8..b352986dde 100644 --- a/src/aiperf/exporters/timeslice_metrics_json_exporter.py +++ b/src/aiperf/exporters/timeslice_metrics_json_exporter.py @@ -44,8 +44,12 @@ def __init__(self, exporter_config: ExporterConfig, **kwargs) -> None: exporter_config.cfg.artifacts.profile_export_timeslices_json_file ) self.trace_or_debug( - lambda: f"Initializing TimesliceMetricsJsonExporter with config: {exporter_config}", - lambda: f"Initializing TimesliceMetricsJsonExporter with file path: {self._file_path}", + lambda: ( + f"Initializing TimesliceMetricsJsonExporter with config: {exporter_config}" + ), + lambda: ( + f"Initializing TimesliceMetricsJsonExporter with file path: {self._file_path}" + ), ) def get_export_info(self) -> FileExportInfo: diff --git a/src/aiperf/graph/__init__.py b/src/aiperf/graph/__init__.py new file mode 100644 index 0000000000..6b71d5e947 --- /dev/null +++ b/src/aiperf/graph/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Graph runtime package: the async-dataflow executor and its dispatch spine. + +No package-level re-exports: consumers import submodules directly +(`aiperf.graph.executor`, `aiperf.graph.channel_store`, +`aiperf.graph.credit_dispatch_adapter`, `aiperf.graph.worker_materialize`, ...). +""" diff --git a/src/aiperf/graph/analysis/__init__.py b/src/aiperf/graph/analysis/__init__.py new file mode 100644 index 0000000000..88767d4ad8 --- /dev/null +++ b/src/aiperf/graph/analysis/__init__.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Native static trace analysis over the async-dataflow graph engine. + +A pure, side-effect-free elaboration of a parsed graph + one trace into an +ordered firing timeline. Reuses the dataflow ``Scheduler`` (adjacency) and the +channel model's ``producers_per_channel`` fan-in primitive rather than a +separate analysis scheduler. The only ordering index is a +parallel-readiness ``cohort`` frontier counter, which is a derived view, not +an execution barrier. +""" + +from __future__ import annotations + +from aiperf.graph.analysis.snapshot import ( + compute_snapshot, + trace_duration_us, +) + +__all__ = [ + "compute_snapshot", + "trace_duration_us", +] diff --git a/src/aiperf/graph/analysis/snapshot.py b/src/aiperf/graph/analysis/snapshot.py new file mode 100644 index 0000000000..82e4b6de7b --- /dev/null +++ b/src/aiperf/graph/analysis/snapshot.py @@ -0,0 +1,117 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Snapshot-at-t* partition of a trace timeline (native, over ``TraceTimeline``). + +Reproduces the agentx "snapshot at t*" liveness partition (see +``aiperf/timing/trajectory_source.py::_snapshot_for``) natively over the firing +timeline produced by :func:`elaborate_trace`. + +A sampling instant ``t*`` splits a trace's firings into a warmup prefix +(``arrival_offset_us < t*`` -- cache-priming history dispatched immediately) and +a profiled set (``arrival_offset_us >= t*`` -- dispatched at its offset from +``t*``). Firings carrying no offset are profiled at dispatch 0 (agentx's +no-timestamp fallback). + +The flat IR emits a single-stream timeline (every firing on the trace's own +stream), so the partition is a plain ``arrival_offset_us`` vs ``t*`` test per +firing. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from aiperf.dataset.graph.models import ParsedGraph, TraceRecord +from aiperf.graph.analysis.timeline import ( + Firing, + TraceTimeline, + elaborate_trace, +) + + +@dataclass(slots=True, frozen=True) +class SnapshotFiring: + """One timeline firing located in the snapshot, with its dispatch timing. + + Attributes: + firing: The underlying :class:`Firing` from the trace timeline. + dispatch_offset_us: ``0`` for warmup firings; ``max(0, arrival - t*)`` + for profiled firings (a firing with no ``arrival_offset_us`` is + profiled at ``0``). + """ + + firing: Firing + dispatch_offset_us: int + + +@dataclass(slots=True, frozen=True) +class Snapshot: + """A trace timeline partitioned at a sampling instant ``t*``. + + Attributes: + t_star_us: The sampling instant in microseconds (pinned to ``0`` for a + ``full_replay`` snapshot). + warmup: Firings whose ``arrival_offset_us < t*`` (cache-priming history; + dispatch offset ``0``), in timeline order. + profiled: Firings at/after ``t*`` (dispatched during PROFILING at their + offset from ``t*``), sorted by ``(dispatch_offset_us, cohort)``. + """ + + t_star_us: int + warmup: tuple[SnapshotFiring, ...] + profiled: tuple[SnapshotFiring, ...] + + +def trace_duration_us(parsed: ParsedGraph, trace: TraceRecord) -> int: + """Return the trace's intrinsic wall-clock span in microseconds. + + The largest ``arrival_offset_us`` across the trace's firings. A trace where + no node carries timing returns ``0``. Snapshot-at-``t*`` lane construction + uses this to choose a sampling instant as a fraction of the trace's duration. + """ + return elaborate_trace(parsed, trace).duration_us() + + +def compute_snapshot( + parsed: ParsedGraph, + trace: TraceRecord, + *, + t_star_us: int, + full_replay: bool = False, +) -> Snapshot: + """Partition ``trace``'s firing timeline at the sampling instant ``t_star_us``. + + ``arrival_offset_us < t*`` firings go to warmup (dispatch offset ``0``); the + rest go to profiled at ``max(0, arrival - t*)`` (a firing with no offset is + profiled at ``t*`` -> dispatch ``0``). + + ``full_replay`` (recycle path) pins ``t*`` to ``0`` and places every firing + in profiled at its own ``arrival_offset_us`` -- byte parity with agentx's + turn-0 recycle. + """ + if full_replay: + t_star_us = 0 + + timeline: TraceTimeline = elaborate_trace(parsed, trace) + + warmup: list[SnapshotFiring] = [] + profiled: list[SnapshotFiring] = [] + for firing in timeline.firings: + off = firing.arrival_offset_us + if not full_replay and off is not None and off < t_star_us: + warmup.append(SnapshotFiring(firing=firing, dispatch_offset_us=0)) + else: + absolute = off if off is not None else t_star_us + profiled.append( + SnapshotFiring( + firing=firing, + dispatch_offset_us=max(0, absolute - t_star_us), + ) + ) + + profiled.sort(key=lambda s: (s.dispatch_offset_us, s.firing.cohort)) + return Snapshot( + t_star_us=t_star_us, + warmup=tuple(warmup), + profiled=tuple(profiled), + ) diff --git a/src/aiperf/graph/analysis/timeline.py b/src/aiperf/graph/analysis/timeline.py new file mode 100644 index 0000000000..b279aedeab --- /dev/null +++ b/src/aiperf/graph/analysis/timeline.py @@ -0,0 +1,193 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Native static elaboration of a trace into an ordered firing timeline. + +``elaborate_trace`` performs a native dataflow dry-run: it seeds the readiness +frontier from the dataflow ``Scheduler.entry_nodes()``, advances a lightweight +static arrival tracker on each fired node's ``write_channels``, and computes +the next frontier from every scheduled successor -- ``successors_after`` +(completion edges) AND ``start_anchored_successors`` (dispatch-anchored edges, +which the runtime schedules at the predecessor's dispatch) -- with fan-in +satisfaction against ``channels.producers_per_channel``. + +There are no lockstep rounds and no channel reduction: the dry-run only needs +arrival counts (how many producers have written each channel) to resolve +AND-fan-in joins, never the channel values themselves. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from aiperf.dataset.graph.models import ( + GraphRecord, + NodeKind, + ParsedGraph, + TraceRecord, + resolve_trace_graph, +) +from aiperf.graph.channels import producers_per_channel +from aiperf.graph.scheduler import Scheduler + + +class GraphCycleError(RuntimeError): + """Raised when ``elaborate_trace(..., depth_cap=N)`` exceeds the firing cap. + + Native replacement for the firing-plan cycle guard: a trace whose readiness + walk emits more than the cap likely traverses a cycle. Runtime callers pass + ``depth_cap=None`` (unbounded) on already-validated graphs; validators pass + a finite cap so a cyclic graph fails loudly instead of looping forever. + """ + + +@dataclass(slots=True, frozen=True) +class Firing: + """One node firing in a trace's static elaboration.""" + + node_id: str + kind: NodeKind + arrival_offset_us: int | None + cohort: int + + +@dataclass(slots=True, frozen=True) +class TraceTimeline: + """A trace's full firing sequence in parallel-readiness order. + + ``cohort`` indexes the parallel-readiness frontier; it is a derived view + for parallel-TTFT cohorting and visualization, NOT an execution barrier. + """ + + firings: tuple[Firing, ...] + + def duration_us(self) -> int: + """Return the max ``arrival_offset_us`` over the timeline (0 default). + + Firings carrying no offset contribute nothing; a trace where no node + carries timing returns 0. + """ + return max( + ( + f.arrival_offset_us + for f in self.firings + if f.arrival_offset_us is not None + ), + default=0, + ) + + +def _inputs_satisfied( + node: object, + arrivals: dict[str, int], + all_counts: dict[str, int], +) -> bool: + """Return True iff every ``ChannelRequirement`` on ``node`` is satisfied. + + Pure static analogue of the dataflow channel store's ``"all"`` fan-in gate: + ``count == "all"`` resolves to the channel's static producer count + (``producers_per_channel``); a finite ``count`` requires that many arrivals. + A node with no ``inputs`` always passes (OR-fan-in successor-walk default). + """ + inputs = getattr(node, "inputs", None) or [] + for req in inputs: + required = all_counts.get(req.channel, 0) if req.count == "all" else req.count + if arrivals.get(req.channel, 0) < required: + return False + return True + + +def _elaborate_graph( + graph: GraphRecord, + scheduler: Scheduler, + *, + depth_cap: int | None, + emitted: list[Firing], +) -> None: + """Walk the graph's readiness frontier, appending a Firing per fired node. + + Static mirror of the executor's scheduling: seed from ``entry_nodes()``, + emit firings per frontier in lexical node order, advance the arrival + tracker on each fired node's ``write_channels``, then schedule successors. + Both anchor kinds are followed -- ``successors_after`` (completion edges) + AND ``start_anchored_successors`` (the runtime schedules those at the + predecessor's DISPATCH, which collapses onto the same firing step here) -- + so start-anchored subtrees elaborate exactly as they fire live. A scheduled + node whose fan-in is not yet satisfied stays ``pending`` and re-enters the + frontier once later arrivals satisfy it, mirroring the runtime's parked + ``await_inputs`` (a runtime node is scheduled once, then blocks). ``cohort`` + is the monotone frontier index. + """ + all_counts = producers_per_channel(graph) + arrivals: dict[str, int] = {} + scheduled: set[str] = set(scheduler.entry_nodes()) + pending: set[str] = set(scheduled) + + def _satisfied() -> list[str]: + return [ + nid + for nid in pending + if _inputs_satisfied(graph.nodes[nid], arrivals, all_counts) + ] + + frontier = _satisfied() + cohort = 0 + fired_count = 0 + while frontier: + pending.difference_update(frontier) + for node_id in sorted(frontier): + node = graph.nodes[node_id] + emitted.append( + Firing( + node_id=node_id, + kind=node.node_type, + arrival_offset_us=getattr(node, "arrival_offset_us", None), + cohort=cohort, + ) + ) + fired_count += 1 + if depth_cap is not None and fired_count > depth_cap: + raise GraphCycleError( + f"trace elaboration exceeded {depth_cap} firings; " + "trace likely traverses a cycle" + ) + for ch in node.write_channels: + arrivals[ch] = arrivals.get(ch, 0) + 1 + for node_id in frontier: + for succ in ( + *scheduler.successors_after(node_id), + *scheduler.start_anchored_successors(node_id), + ): + if succ in scheduled: + continue + scheduled.add(succ) + pending.add(succ) + frontier = _satisfied() + cohort += 1 + + +def elaborate_trace( + parsed: ParsedGraph, + trace: TraceRecord, + *, + depth_cap: int | None = None, +) -> TraceTimeline: + """Native dataflow dry-run of one trace into a full firing timeline. + + Walks the trace's top-level graph (``resolve_trace_graph``) frontier. + Fan-in reuses the dataflow scheduler and channel model; no lockstep-round + state, no channel reduction. + + ``depth_cap`` bounds the total firings emitted (raises ``GraphCycleError`` + on overflow) for validators that may be called on cyclic graphs; ``None`` + (default) is unbounded for runtime callers on validated graphs. + """ + graph = resolve_trace_graph(parsed, trace) + scheduler = Scheduler(graph) + emitted: list[Firing] = [] + _elaborate_graph( + graph, + scheduler, + depth_cap=depth_cap, + emitted=emitted, + ) + return TraceTimeline(firings=tuple(emitted)) diff --git a/src/aiperf/graph/channel_store.py b/src/aiperf/graph/channel_store.py new file mode 100644 index 0000000000..32b12d70a8 --- /dev/null +++ b/src/aiperf/graph/channel_store.py @@ -0,0 +1,431 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""VersionedChannelStore - per-trace channel state for the async-dataflow executor. + +A versioned log per channel. Writes are linearized via a single monotonic +`_next_seq` counter; readers capture the per-channel versions at firing time and +reducers consume them in (write_seq, writer_node_id) order on read. No deepcopy +anywhere - values are stored share-by-ref. +""" + +from __future__ import annotations + +import asyncio +import contextlib +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +from aiperf.dataset.graph.models import ChannelRequirement, ChannelSpec, ReducerName +from aiperf.graph.reducers import ( + UNSET, + OverwriteConflictError, + Reducer, + get_reducer, +) + +__all__ = [ + "ChannelOrphanedError", + "ReducerRegistry", + "UnknownChannelError", + "VersionedChannelStore", +] + + +class UnknownChannelError(KeyError): + """Raised when writing or reading a channel not declared in `channel_specs`.""" + + +class ChannelOrphanedError(RuntimeError): + """Raised when an `await_inputs` requirement can no longer be satisfied. + + Surfaces when every producer of a required channel has reported done + (with `success=False` or `wrote=False`) before the waiter's count target + is reached. + """ + + +# Type alias for the reducer dispatch dependency. The store does not depend +# on a concrete registry class - any callable that maps reducer-name to a +# `(current, [(writer_id, value), ...]) -> reduced` reducer works. Defaults +# to `aiperf.graph.reducers.get_reducer`. +ReducerRegistry = Callable[[str], Reducer] + + +@dataclass(slots=True, frozen=True) +class _LogEntry: + """One committed write to a value channel.""" + + write_seq: int + """Monotonic global sequence number assigned at commit time.""" + writer_node_id: str + """Identifier of the node that produced this write; secondary sort key.""" + value: Any + """The raw write value; stored by reference, never deep-copied.""" + + +@dataclass(slots=True, frozen=True) +class _VersionCapture: + """Immutable snapshot of which channel versions count for one firing.""" + + channel: str + """Channel this capture refers to.""" + captured_seqs: tuple[int, ...] + """`write_seq` values included in this capture, in commit order.""" + + +@dataclass(slots=True) +class _Waiter: + """One pending `await_inputs` requirement on a single channel.""" + + channel: str + required_count: int + event: asyncio.Event = field(default_factory=asyncio.Event) + orphaned_reason: str | None = None + + +class VersionedChannelStore: + """Per-trace channel store backed by per-channel append-only logs. + + Writers call `write(channels, value, writer_node_id=...)`; readers call + `await await_inputs(reqs)` to capture versions, then `read(reqs, capture)` + to materialize reduced values. Cancellation is communicated by the + executor via `mark_producer_done(channel, success=..., wrote=...)`. + """ + + __slots__ = ( + "_specs", + "_logs", + "_arrival_count", + "_producers_remaining", + "_producers_declared", + "_initial", + "_next_seq", + "_waiters", + "_reducers", + "_orphaned", + "_overwrite_writer", + ) + + def __init__( + self, + initial: dict[str, Any], + channel_specs: dict[str, ChannelSpec], + producers_per_channel: dict[str, int], + ) -> None: + self._specs: dict[str, ChannelSpec] = dict(channel_specs) + self._logs: dict[str, list[_LogEntry]] = {ch: [] for ch in channel_specs} + self._arrival_count: dict[str, int] = dict.fromkeys(channel_specs, 0) + self._producers_remaining: dict[str, int] = { + ch: int(producers_per_channel.get(ch, 0)) for ch in channel_specs + } + self._producers_declared: dict[str, int] = { + ch: int(producers_per_channel.get(ch, 0)) for ch in channel_specs + } + self._initial: dict[str, Any] = dict(initial) + # `_next_seq` is incremented before every commit. Reserve 0 for + # synthetic initial-state writes; node writes start at 1. + self._next_seq: int = 0 + self._waiters: dict[str, list[_Waiter]] = {ch: [] for ch in channel_specs} + self._reducers: ReducerRegistry = get_reducer + self._orphaned: dict[str, str] = {} + # Tracks the writer id of the first non-init write to each overwrite + # channel so a different writer can be rejected with + # OverwriteConflictError across the trace's lifetime (cross-spec + # decision 1: strict over the whole trace, not per-dispatch). + self._overwrite_writer: dict[str, str] = {} + + for ch, value in self._initial.items(): + if ch not in self._specs: + raise UnknownChannelError( + f"initial state names unknown channel: {ch!r}" + ) + self._logs[ch].append( + _LogEntry(write_seq=0, writer_node_id="__init__", value=value) + ) + # Note: init seed does NOT bump `_arrival_count`. arrival_count + # counts node writes only, so `count=N` requirements measure + # producer arrivals. + + # ------------------------------------------------------------------ + # write + # ------------------------------------------------------------------ + def write( + self, + channel_names: list[str], + value: Any, + *, + writer_node_id: str, + ) -> None: + """Commit `value` to every channel in `channel_names`.""" + if not channel_names: + return + for ch in channel_names: + self._validate_write_channel(ch, value, writer_node_id=writer_node_id) + + for ch in channel_names: + self._commit_write_channel(ch, value, writer_node_id=writer_node_id) + self._wake_waiters(ch) + + def _validate_write_channel( + self, channel: str, value: Any, *, writer_node_id: str + ) -> None: + if channel not in self._specs: + raise UnknownChannelError(f"unknown channel: {channel!r}") + if self._specs[channel].reducer is not ReducerName.OVERWRITE: + return + prior = self._overwrite_writer.get(channel) + if prior is not None: + raise OverwriteConflictError( + f"overwrite channel {channel!r} already written by " + f"{prior!r}; rejecting second writer {writer_node_id!r}" + ) + + def _commit_write_channel( + self, channel: str, value: Any, *, writer_node_id: str + ) -> None: + self._next_seq += 1 + entry = _LogEntry( + write_seq=self._next_seq, + writer_node_id=writer_node_id, + value=value, + ) + self._logs[channel].append(entry) + self._arrival_count[channel] += 1 + if self._specs[channel].reducer is ReducerName.OVERWRITE: + self._overwrite_writer.setdefault(channel, writer_node_id) + + # ------------------------------------------------------------------ + # await_inputs + # ------------------------------------------------------------------ + async def await_inputs( + self, requirements: list[ChannelRequirement] + ) -> dict[str, _VersionCapture]: + """Block until every requirement is satisfied; return frozen captures. + + For each requirement: `count=N` waits for the N-th arrival; `count="all"` + resolves to the static topology count at call time and then behaves + identically. Raises `ChannelOrphanedError` if any required channel can + no longer fulfil its count. + """ + captures: dict[str, _VersionCapture] = {} + for req in requirements: + ch = req.channel + if ch not in self._specs: + raise UnknownChannelError(f"unknown channel: {ch!r}") + if ch in self._orphaned: + raise ChannelOrphanedError( + f"channel {ch!r} orphaned: {self._orphaned[ch]}" + ) + target = self._resolve_count(req) + await self._await_count(ch, target) + captures[ch] = self._capture(ch, target) + return captures + + def _resolve_count(self, req: ChannelRequirement) -> int: + if req.count == "all": + # "all" resolves to the static topology count at call time. Late + # producer cancellation may render this unreachable; orphan check + # inside _await_count handles that case. + return self._producers_declared[req.channel] + return int(req.count) + + async def _await_count(self, channel: str, target: int) -> None: + if target <= 0: + return + while self._arrival_count[channel] < target: + if channel in self._orphaned: + raise ChannelOrphanedError( + f"channel {channel!r} orphaned: {self._orphaned[channel]}" + ) + # Check reachability: arrivals so far + still-live producers must + # reach target, else orphan immediately. + reachable = ( + self._arrival_count[channel] + self._producers_remaining[channel] + ) + if reachable < target: + self._orphaned[channel] = "insufficient_producers_remaining" + raise ChannelOrphanedError( + f"channel {channel!r} cannot reach count={target}: " + f"only {reachable} arrivals possible " + f"({self._arrival_count[channel]} so far, " + f"{self._producers_remaining[channel]} producers remain)" + ) + waiter = _Waiter(channel=channel, required_count=target) + self._waiters[channel].append(waiter) + try: + await waiter.event.wait() + finally: + with contextlib.suppress(ValueError): + self._waiters[channel].remove(waiter) + if waiter.orphaned_reason is not None: + raise ChannelOrphanedError( + f"channel {channel!r} orphaned: {waiter.orphaned_reason}" + ) + + def _capture(self, channel: str, target: int) -> _VersionCapture: + entries = self._logs[channel] + if target <= 0: + return _VersionCapture(channel=channel, captured_seqs=()) + # Init seed (write_seq=0) is the reducer seed, not an "arrival"; the + # count-N requirement measures node writes only. The init entry is + # always implicitly part of read (handled in _reduce_value_channel). + non_init = [e for e in entries if e.write_seq != 0] + sorted_entries = sorted(non_init, key=lambda e: (e.write_seq, e.writer_node_id)) + chosen = sorted_entries[:target] + return _VersionCapture( + channel=channel, + captured_seqs=tuple(e.write_seq for e in chosen), + ) + + def _wake_waiters(self, channel: str) -> None: + if not self._waiters[channel]: + return + current = self._arrival_count[channel] + for waiter in list(self._waiters[channel]): + if current >= waiter.required_count: + waiter.event.set() + + # ------------------------------------------------------------------ + # read + # ------------------------------------------------------------------ + def read( + self, + requirements: list[ChannelRequirement], + capture: dict[str, _VersionCapture], + ) -> dict[str, Any]: + """Return the reduced value per channel at the captured versions.""" + out: dict[str, Any] = {} + for req in requirements: + ch = req.channel + if ch not in self._specs: + raise UnknownChannelError(f"unknown channel: {ch!r}") + cap = capture.get(ch) + if cap is None: + raise KeyError(f"no capture for channel {ch!r}") + out[ch] = self._reduce_value_channel(ch, cap.captured_seqs) + return out + + def _reduce_value_channel( + self, channel: str, captured_seqs: tuple[int, ...] + ) -> Any: + spec = self._specs[channel] + reducer = self._reducers(spec.reducer) + entries = self._logs[channel] + # Init seed (write_seq=0) is the reducer's starting value; it is not + # itself a "captured" write. Capture sets only hold node writes. + init_entries = [e for e in entries if e.write_seq == 0] + current: Any = init_entries[0].value if init_entries else UNSET + if not captured_seqs: + return current + seq_set = set(captured_seqs) + chosen = [e for e in entries if e.write_seq in seq_set] + chosen.sort(key=lambda e: (e.write_seq, e.writer_node_id)) + tuples = [(e.writer_node_id, e.value) for e in chosen] + return reducer(current, tuples) + + # ------------------------------------------------------------------ + # snapshot + # ------------------------------------------------------------------ + def snapshot(self) -> dict[str, Any]: + """Return the final user-visible view of every channel.""" + out: dict[str, Any] = {} + for ch in self._specs: + entries = self._logs[ch] + if not entries: + out[ch] = UNSET + continue + sorted_entries = sorted( + entries, key=lambda e: (e.write_seq, e.writer_node_id) + ) + init_entries = [e for e in sorted_entries if e.write_seq == 0] + non_init = [e for e in sorted_entries if e.write_seq != 0] + current: Any = init_entries[0].value if init_entries else UNSET + tuples = [(e.writer_node_id, e.value) for e in non_init] + spec = self._specs[ch] + reducer = self._reducers(spec.reducer) + out[ch] = reducer(current, tuples) if tuples else current + return out + + @property + def current_seq(self) -> int: + """Latest committed write sequence number; 0 when only init seeds exist. + + Reading this immediately after `await_inputs` returns gives a causal + cap for `snapshot_at_seq` that includes every write that satisfied + the input gate and excludes any concurrent write committed afterward. + """ + return self._next_seq + + def snapshot_at_seq(self, max_seq: int) -> dict[str, Any]: + """Reduce every channel considering only entries with `write_seq <= max_seq`. + + A causally-consistent global cut: an LlmNode firing's `inputs` view + (built in the executor's `_prepare_node_inputs`) covers every channel + — including ones the node did not declare in `inputs` — as of the + captured sequence. + """ + out: dict[str, Any] = {} + for ch in self._specs: + entries = [e for e in self._logs[ch] if e.write_seq <= max_seq] + if not entries: + # Unset channels are absent from the global snapshot dict so + # consumers can treat `snapshot.get(ch) is None` as unwritten; + # returning the `UNSET` sentinel here would smuggle a non-None + # placeholder into their value path. + continue + sorted_entries = sorted( + entries, key=lambda e: (e.write_seq, e.writer_node_id) + ) + init_entries = [e for e in sorted_entries if e.write_seq == 0] + non_init = [e for e in sorted_entries if e.write_seq != 0] + current: Any = init_entries[0].value if init_entries else UNSET + tuples = [(e.writer_node_id, e.value) for e in non_init] + spec = self._specs[ch] + reducer = self._reducers(spec.reducer) + out[ch] = reducer(current, tuples) if tuples else current + return out + + # ------------------------------------------------------------------ + # producer accounting / orphan propagation + # ------------------------------------------------------------------ + def mark_producer_done( + self, + channel: str, + *, + success: bool, + wrote: bool, + ) -> None: + """Inform the store that one producer of `channel` has terminated. + + `success=True, wrote=True`: a real write already landed; nothing to do + beyond decrementing the remaining-producer count (the write itself + bumped `_arrival_count`). + `success=True, wrote=False`: producer ran to completion without writing + (e.g. an expected no-write exit). Decrement only. + `success=False`: producer cancelled or failed. Decrement and, if this + was the last potential producer AND outstanding waiters cannot reach + their count, orphan the channel. + """ + if channel not in self._specs: + raise UnknownChannelError(f"unknown channel: {channel!r}") + if self._producers_remaining[channel] > 0: + self._producers_remaining[channel] -= 1 + reachable = self._arrival_count[channel] + self._producers_remaining[channel] + orphan_reason = ( + "insufficient_producers_remaining" if success else "all_producers_cancelled" + ) + for waiter in list(self._waiters[channel]): + if reachable < waiter.required_count: + waiter.orphaned_reason = orphan_reason + waiter.event.set() + has_init_seed = channel in self._logs and any( + entry.write_seq == 0 for entry in self._logs[channel] + ) + if ( + not success + and self._producers_remaining[channel] == 0 + and self._arrival_count[channel] == 0 + and not has_init_seed + ): + self._orphaned[channel] = orphan_reason diff --git a/src/aiperf/graph/channels.py b/src/aiperf/graph/channels.py new file mode 100644 index 0000000000..41fe2c0cae --- /dev/null +++ b/src/aiperf/graph/channels.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Native channel primitives for the async-dataflow graph engine. + +Single source of truth for static channel topology derivations shared by the +executor (``"all"`` fan-in resolution in the channel store) and the static +trace analyzer. The only derivation here today is ``producers_per_channel``; +the dataflow channel store and the analyzer both consume it instead of +re-deriving per-channel producer counts at their own callsites. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from aiperf.dataset.graph.models import GraphRecord + +__all__ = ["producers_per_channel"] + + +def producers_per_channel(graph: GraphRecord) -> dict[str, int]: + """Count the nodes that statically write each channel. + + A channel's producer count is the number of nodes whose ``write_channels`` + (the IR property on each node) include that channel. In the flat IR the only + writer is ``LlmNode.output``; ``write_channels`` covers it uniformly, so no + per-kind special-casing is needed here. + + Every channel declared in ``graph.state`` is seeded to ``0`` so the result + carries an entry for declared-but-unwritten channels (e.g. pure initial + state). Counts for channels a node writes that are not in ``graph.state`` + (such as error markers) are still included. + + Args: + graph: The static graph IR whose nodes' ``write_channels`` are counted. + + Returns: + Mapping of channel name to the number of nodes that write it. + """ + counts: dict[str, int] = {ch: 0 for ch in graph.state} + for node in graph.nodes.values(): + for ch in node.write_channels: + counts[ch] = counts.get(ch, 0) + 1 + return counts diff --git a/src/aiperf/graph/context.py b/src/aiperf/graph/context.py new file mode 100644 index 0000000000..47e7d60bbf --- /dev/null +++ b/src/aiperf/graph/context.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-trace state container, sentinel exception, and node-result shape for the +async-dataflow executor. + +Three small types co-located to avoid circular imports between the executor, +the channel store, and the node-dispatch table. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from aiperf.dataset.graph.models import TraceRecord + from aiperf.graph.channel_store import VersionedChannelStore + + +class _NodeExpectedExit(BaseException): + """Signals a _fire task exiting cleanly without writing successors. + + Base class for clean-exit signals a node can raise to terminate without + propagating to successors. Inherits BaseException (not Exception) so the + asyncio.TaskGroup cascade does not treat it as a programming error. + + The sole concrete raiser today is the ``_NodeOverflowTerminate`` subclass + (defined in ``aiperf.graph.credit_dispatch_adapter``), which the credit + path raises on a context-overflow response so the executor terminates the + whole trajectory cleanly. + """ + + +@dataclass(slots=True, frozen=True) +class Write: + """A single channel write produced by a node's _execute return. + + Frozen so accidental mutation between produce and consume is impossible. + """ + + channel: str + value: object + + +@dataclass(slots=True) +class NodeExecutionResult: + """The return value of every node-kind's _execute. + + `writes` carries the channel writes the node produced (possibly empty). + """ + + writes: list[Write] = field(default_factory=list) + + +@dataclass(slots=True) +class _TraceContext: + """Per-trace mutable state passed into every _fire call. + + Owned by `TraceExecutor.run` for the duration of a single trace; never + shared across traces. Mutable fields default to empty collections. + """ + + trace: TraceRecord + store: VersionedChannelStore + tg: asyncio.TaskGroup | None = None + scheduled_node_ids: set[str] = field(default_factory=set) + tasks_by_node_id: dict[str, asyncio.Task] = field(default_factory=dict) + # Set True when any LlmNode in this trace returned a context-overflow error + # (early termination). The overflowed node exits cleanly via + # ``_NodeOverflowTerminate``; this flag lets the executor treat the resulting + # downstream ``ChannelOrphanedError`` cascade (successors awaiting the + # never-written output channel) as a CLEAN trajectory stop rather than a trace + # error, so the rest of the trace's turns do not dispatch and the instance is + # not counted as ``errored_traces``. + overflow_terminated: bool = False + # Wall-clock-us at which each node's `_fire` reached its `finally` block, + # i.e. the moment it actually finished executing (success, expected exit, + # or race-cancel). Read by successor `_apply_firing_delay` to anchor + # incoming-edge `delay_after_predecessor_us` gates on the predecessor's + # ACTUAL finish time. + # No lock needed: writes happen in the predecessor's `_fire` `finally` + # before successor scheduling (`_schedule(succ_id, ...)` runs AFTER + # `mark_producer_done`, which runs in the same `finally`), so the + # successor's `await_inputs` is guaranteed to see the write. + node_finish_wall_us: dict[str, float] = field(default_factory=dict) + # Wall-clock-us at which each node's firing gate cleared and it proceeded + # to execute (its dispatch instant). Written in `_prepare_node_inputs` + # immediately after `_apply_firing_delay` returns and BEFORE start-anchored + # successors are scheduled on the same loop iteration, so a successor's + # `_compute_firing_gate_us` read is guaranteed to see the write (same + # single-loop happens-before argument as `node_finish_wall_us`). + node_dispatch_wall_us: dict[str, float] = field(default_factory=dict) + # Wall-clock-us at which each node's first token was observed (FirstToken + # event routed through the dispatch adapter's stamp closure). Written on + # the same single loop that reads it in successor gate computation -- + # same happens-before argument as node_dispatch_wall_us. A node absent + # from this map after its first-token event is set terminated without + # streaming a first token (fallback gate applies). + node_first_token_wall_us: dict[str, float] = field(default_factory=dict) + # Per-node first-token latches. SET by the stamp closure (first token + # observed) or by _finalize_node (terminal without one). Successors with + # first-token-anchored incoming edges await these before computing their + # firing gate. + node_first_token_events: dict[str, asyncio.Event] = field(default_factory=dict) + + def first_token_event(self, node_id: str) -> asyncio.Event: + """Lazy per-node first-token latch; single-loop access needs no lock.""" + return self.node_first_token_events.setdefault(node_id, asyncio.Event()) diff --git a/src/aiperf/graph/credit_dispatch_adapter.py b/src/aiperf/graph/credit_dispatch_adapter.py new file mode 100644 index 0000000000..4c35518be5 --- /dev/null +++ b/src/aiperf/graph/credit_dispatch_adapter.py @@ -0,0 +1,545 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CreditDispatchAdapter — the Future bridge from the dataflow executor to v1 credits. + +The async-dataflow ``TraceExecutor`` calls exactly one method on its injected +``credit_issuer``:: + + async def dispatch(self, node, request, ctx, **kwargs) -> Any + +(see ``aiperf.graph.dispatch.llm``). The return value feeds ``node.output`` +(the dispatch module substitutes a type-correct empty list for messages-typed +channels). No downstream node consumes an LLM output's value (prompts are +materialized worker-side from recorded envelopes; output channels gate fan-in +only), so a placeholder ``str`` is contractually correct. + +This adapter turns that fire-and-forget executor call into an awaitable backed by +the v1 credit pipeline: + +1. Resolve the fired node's ``node_ordinal`` from the build-time catalog. +2. Mint a collision-free correlation key (see ``_mint``). +3. Park an ``asyncio.Future`` under that key. +4. Issue a real ``TurnToSend`` (via ``CreditIssuer.issue_graph_credit``, which + BYPASSES the linear session-slot lifecycle -- the strategy owns concurrency). +5. ``await`` the Future under a timeout guard; ``resolve`` (driven by the + unconditional graph-return hook on ``CreditCallbackHandler``) sets the result + on a normal return or rejects it on cancel / error. + +Why the adapter owns its own context (NOT ``ctx``) +-------------------------------------------------- +The executor passes a ``PlacementContext`` carrying ONLY ``parent_trace_id`` / +``parent_node_id`` -- a frozen ``slots=True`` dataclass. +It has no ``trace``, ``phase_variant``, ``agent_depth``, or correlation fields, +and cannot be monkey-patched. The adapter is therefore constructed PER-RUN with +the per-trace identity it needs and derives node identity from +``request.node_id`` + ``ctx.parent_trace_id`` alone. +""" + +from __future__ import annotations + +import asyncio +import uuid +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +from aiperf.common.aiperf_logger import AIPerfLogger +from aiperf.common.environment import Environment +from aiperf.common.scenario import is_context_overflow_response +from aiperf.credit.structs import TurnToSend +from aiperf.dataset.graph.graph_path_catalog import ( + CatalogContext, + node_ordinal_for, +) +from aiperf.graph.context import _NodeExpectedExit + +if TYPE_CHECKING: + from aiperf.graph.placement import DispatchRequest, PlacementContext + +_logger = AIPerfLogger(__name__) + +__all__ = ["CreditDispatchAdapter", "GraphDispatchError", "_NodeOverflowTerminate"] + +# Placeholder resolved to the dispatch caller and (for non-messages channels) +# written to ``node.output``. Validated: no downstream node consumes an LLM +# output channel's value (fan-in gating only), so the concrete value is +# irrelevant. +_PLACEHOLDER = "" + + +class GraphDispatchError(RuntimeError): + """A graph dispatch failed (worker error OR worker-reported cancellation). + + A plain ``Exception`` subclass (NOT ``asyncio.CancelledError``) so the + executor's per-trace TaskGroup unwinds it as a normal task failure rather + than mistaking it for cooperative cancellation of the awaiting coroutine. + """ + + +class GraphStickinessError(GraphDispatchError): + """A dynamic slot's pool value was missing on the routed worker (§6.3). + + Broken stickiness (worker death re-route) or dynamic-pool backstop + eviction. Non-containable: the executor re-raises it as a trace error -- + the workload's content dependency is unsatisfiable, so continuing with + omission would silently corrupt the trace shape. + """ + + +_POOL_MISSING_PREFIX = "aiperf.graph.pool_missing:" + + +class CreditIssueRefusedError(GraphDispatchError): + """The issuer refused to place the credit on the wire (stop gate / caps). + + Distinct from a post-issue dispatch failure: the executor's + mid-conversation containment must NOT sentinel-continue past a refusal -- + at duration end every remaining node would otherwise churn through + sentinel writes without ever dispatching. Refusal is a trace-stop. + """ + + +class _NodeOverflowTerminate(_NodeExpectedExit): + """A node dispatch returned a context-overflow error (early termination). + + Subclasses ``_NodeExpectedExit`` so ``_run_node`` catches it as a clean exit + (no successors, no trace error); the dedicated subclass lets tests + distinguish overflow. + """ + + +def _split_shaped_id(node_id: str) -> tuple[str, int] | None: + """``(scope, turn)`` of a ``{scope}:{turn}`` node id; None otherwise. + + Recorded adapters and dag_jsonl lowering always emit the shape; native + graphs keep author-chosen ids verbatim (``"plan"``, ``"phase:review"``), + which return None so callers fall back to root-trajectory identity. + """ + scope, sep, turn = node_id.rpartition(":") + if sep and scope and turn.isdigit(): + return scope, int(turn) + return None + + +class CreditDispatchAdapter: + """Bridges one trace's executor dispatches onto the v1 credit pipeline. + + Constructed per-trace by ``GraphIRReplayStrategy`` with that trace's + identity. Thread-confined to the TimingManager event loop: ``dispatch`` is + awaited from the executor's per-trace TaskGroup and ``resolve`` is invoked + synchronously from the credit-return callback on the SAME loop, so the + waiter dict needs no locking. + """ + + def __init__( + self, + *, + credit_issuer: Any, + catalog_context: CatalogContext, + trace_id: str, + instance_id: str | None = None, + phase_variant: str = "profiling", + parent_correlation_id: str | None = None, + dispatch_timeout_s: float | None = None, + on_drained: Callable[[CreditDispatchAdapter], None] | None = None, + first_token_sources: frozenset[str] = frozenset(), + node_identity: dict[str, tuple[int, str | None]] | None = None, + ) -> None: + """Initialize the per-trace dispatch adapter. + + Args: + credit_issuer: The real ``CreditIssuer``; the adapter calls its + ``issue_graph_credit`` graph path (session-slot-bypassing). + catalog_context: Build-time ``{trace_id: {node_key: ordinal}}`` + + namespace maps; resolves a fired node to its ``node_ordinal``. + trace_id: This trace's BASE/template root id (e.g. ``"t-1"``). Keys + BOTH the build-time catalog (node-ordinal resolution) AND the + graph store mmap (the worker strips any ``#`` recycle suffix + back to this base). ``ctx.parent_trace_id`` carries this base as + the bare instance id. + instance_id: The per-recycle INSTANCE id stamped on ``credit.trace_id`` + (e.g. ``"t-1#0.0"`` then ``"t-1#0.1"`` after one recycle). Distinct + per (lane, recycle pass) so the worker's cache-bust marker ROTATES + AND the strategy's return observer routes by it without two + concurrent instances of one template colliding. Defaults to + ``trace_id`` for non-recycling callers (unit harness). + phase_variant: Graph phase-variant label stamped on every credit + (``"profiling"`` / ``"warmup"``); also part of the correlation + key so the two variants of one node never collide. + parent_correlation_id: Parent session correlation id for children. + dispatch_timeout_s: Per-dispatch deadlock guard; defaults to + ``Environment.GRAPH.DISPATCH_TIMEOUT``. + on_drained: Optional callback invoked (with ``self``) whenever a return + empties the in-flight waiter set (``inflight_count`` -> 0); the + strategy uses it to defer adapter teardown until it is truly idle + (C1 guard). + first_token_sources: Node ids whose issued turn must carry + ``first_token_event=True`` -- the sources of this trace's + first-token-anchored edges (post-TTFT anchoring). A dispatch of + such a node emits a ``FirstToken`` so a successor gated on the + node's observed first token can be released. Empty (default) => + no node emits a first-token event (pre-anchoring byte parity). + node_identity: ``node_id -> (agent_depth, parent_node_id)`` legacy + record identity map (dag_jsonl lowering's ``metadata["dag"]`` + stamp). A dispatch whose node hits the map carries the mapped + depth and, when ``parent_node_id`` is set, the parent NODE's + derived correlation id. ``None`` (weka/dynamo default) keeps + behavior byte-identical: depth 0 and the constructor's + ``parent_correlation_id`` on every credit. + """ + self._issuer = credit_issuer + self._catalog = catalog_context + self._trace_id = trace_id + # Instance id stamped on every credit's ``trace_id`` (marker rotation + + # strategy return de-mux); falls back to the base id when no recycle suffix. + self._instance_id = instance_id if instance_id is not None else trace_id + self._phase_variant = phase_variant + # scope -> trajectory-instance x_correlation_id, minted lazily per + # scope: fresh per adapter (= per instance/recycle), stable across all + # of that trajectory's turns within this instance. + self._scope_corrs: dict[str, str] = {} + self._parent_correlation_id = parent_correlation_id + # scope -> recorded turn count, from the build-time catalog: the max + # {scope}:{turn} coordinate per scope (+1), with native bare ids + # counting as root-trajectory turns by ordinal. Gives every credit a + # REAL num_turns so ``is_final_turn`` carries the trajectory's + # recorded session-final fact (session-routing modes key bind/close + # and session-final semantics on it). + scope_last_turn: dict[str, int] = {} + for node_key, ordinal in catalog_context.catalog.get(trace_id, {}).items(): + shaped = _split_shaped_id(node_key) + scope, turn = shaped if shaped is not None else (trace_id, ordinal) + if turn > scope_last_turn.get(scope, -1): + scope_last_turn[scope] = turn + self._scope_num_turns = { + scope: last + 1 for scope, last in scope_last_turn.items() + } + self._timeout_s = ( + dispatch_timeout_s + if dispatch_timeout_s is not None + else Environment.GRAPH.DISPATCH_TIMEOUT + ) + # Correlation key -> parked Future. The key embeds runtime trace id + # (instance identity), node id, phase variant, and a monotonic counter + # so two in-flight dispatches of the same node NEVER share a waiter. + self._waiters: dict[tuple[str, int], asyncio.Future[str]] = {} + # Node ids whose dispatch stamps ``first_token_event=True`` (sources of + # this trace's first-token-anchored edges); the successor gates on the + # emitted ``FirstToken`` (post-TTFT anchoring). + self._first_token_sources = first_token_sources + # node_id -> (agent_depth, parent_node_id) legacy identity map (dag); + # None => every dispatch is a root-chain firing (weka/dynamo). + self._node_identity = node_identity + # Correlation key -> zero-arg callback fired ONCE when a FirstToken for + # that key arrives (``on_first_token``). Parked per-dispatch beside the + # waiter and popped in the SAME resolve/finally cleanup so nothing leaks. + self._first_token_cbs: dict[tuple[str, int], Callable[[], None]] = {} + # Invoked whenever the in-flight waiter set drains to empty so the owner + # can pop the adapter only once it is truly idle. + self._on_drained = on_drained + + @property + def inflight_count(self) -> int: + """Number of dispatches currently awaiting a return (test/diagnostic).""" + return len(self._waiters) + + @property + def phase_variant(self) -> str: + """Graph-IR phase variant label this adapter stamps on its credits.""" + return self._phase_variant + + @property + def instance_id(self) -> str: + """The per-recycle instance id this adapter stamps on ``credit.trace_id``. + + The strategy's return observer routes credits back to the owning adapter + by this id (NOT the base template id), so two concurrent recycle + instances of one template never collide on the de-mux registry. + """ + return self._instance_id + + async def dispatch( + self, + node: Any, + request: DispatchRequest, + ctx: PlacementContext, + first_token_cb: Callable[[], None] | None = None, + **kwargs: Any, + ) -> str: + """Issue a graph credit for ``node`` and await its correlated return. + + Tolerates and ignores extra keyword arguments from the LLM dispatch + path. Returns a placeholder ``str`` on a normal return; raises on + cancel / error / timeout so the executor coroutine unwinds rather than + hangs. + + ``first_token_cb`` (optional): a zero-arg callable parked under this + dispatch's waiter key and invoked AT MOST ONCE by :meth:`on_first_token` + when the emitting credit's ``FirstToken`` arrives -- the release hook a + first-token-anchored successor registers (post-TTFT anchoring). It is + popped in the SAME resolve/timeout/finally paths as the waiter, so a + dispatch that resolves without a TTFT never leaks or late-fires it. + """ + runtime_trace_id = ctx.parent_trace_id or self._trace_id + node_id = request.node_id + node_ordinal = self._resolve_ordinal(runtime_trace_id, node_id) + + x_corr, turn_index, num_turns = self._mint(node_id, node_ordinal) + key = (x_corr, turn_index) + if key in self._waiters: + raise RuntimeError( + f"duplicate in-flight dispatch for {key!r} (node {node_id!r}, " + f"instance {self._instance_id!r}): the executor fires each " + "node at most once per instance run" + ) + conversation_id = self._conversation_identity(node_id) + agent_depth, parent_correlation_id = self._dag_identity( + runtime_trace_id, node_id + ) + + loop = asyncio.get_running_loop() + fut: asyncio.Future[str] = loop.create_future() + self._waiters[key] = fut + if first_token_cb is not None: + self._first_token_cbs[key] = first_token_cb + + turn = TurnToSend( + conversation_id=conversation_id, + x_correlation_id=x_corr, + turn_index=turn_index, + num_turns=num_turns, + agent_depth=agent_depth, + parent_correlation_id=parent_correlation_id, + # The instance's ROOT trajectory corr: session-tree identity for + # session-routing plugins (dynamo parent/tree grouping). Minted + # lazily, so it is stable even before the root scope fires. + root_correlation_id=self._corr_of(self._trace_id), + trace_id=self._instance_id, + node_ordinal=node_ordinal, + phase_variant=self._phase_variant, + first_token_event=node_id in self._first_token_sources, + ) + try: + issued = await self._issuer.issue_graph_credit(turn) + if not issued: + pending = self._waiters.pop(key, None) + # No credit reached the wire => no FirstToken will arrive; drop + # any parked first-token cb so it cannot late-fire on reuse. + self._first_token_cbs.pop(key, None) + if pending is not None and not pending.done(): + pending.set_exception( + CreditIssueRefusedError( + "graph credit refused by issuer (stop/duration/" + "request-count cap reached or run cancelled); " + f"trace={self._trace_id!r} node_ordinal={node_ordinal!r}" + " -- no return will arrive, stopping this trace" + ) + ) + return await asyncio.wait_for(fut, timeout=self._timeout_s) + except (TimeoutError, asyncio.CancelledError): + # Drop the orphaned waiter so a late return can't resolve a dead + # Future, then re-raise so the executor coroutine unwinds. + self._waiters.pop(key, None) + self._first_token_cbs.pop(key, None) + raise + finally: + # Defensive: a resolved/rejected Future may already be popped by + # ``resolve``; ensure no waiter (or parked first-token cb) lingers. + self._waiters.pop(key, None) + self._first_token_cbs.pop(key, None) + + def resolve(self, credit: Any, error: str | None, cancelled: bool) -> None: + """Resolve (or reject) the Future parked for ``credit``'s correlation key. + + Driven by ``CreditCallbackHandler``'s unconditional graph-return hook. + Unknown keys (already resolved, or never parked) are a graceful no-op. + Fires ``on_drained`` whenever this return empties the in-flight waiter set, + so the owner can defer de-mux teardown until the adapter is idle. + """ + key = (credit.x_correlation_id, credit.turn_index) + fut = self._waiters.pop(key, None) + # The dispatch is settling; a first-token cb not already consumed by + # ``on_first_token`` (error/cancel before TTFT, or non-streaming node) + # is now moot -- drop it so it cannot late-fire. + self._first_token_cbs.pop(key, None) + try: + if fut is None: + _logger.debug( + lambda: ( + f"graph return for unknown waiter key {key} " + f"(trace={getattr(credit, 'trace_id', None)}); dropped" + ) + ) + return + if fut.done(): + return + if error is not None and is_context_overflow_response(body=error): + # An overflowed node TERMINATES the trajectory early -- later + # turns carry even more context and would only overflow too. The + # clean exit suppresses downstream dispatch; the overflow record + # still flows to the RecordProcessor metrics-skip path. + _logger.info( + lambda: ( + f"Terminating trajectory {self._instance_id!r} early at " + f"node_ordinal={getattr(credit, 'node_ordinal', None)!r}: " + "context-overflow error from server" + ) + ) + fut.set_exception( + _NodeOverflowTerminate(f"context-overflow early-term: {error}") + ) + elif error is not None and error.startswith(_POOL_MISSING_PREFIX): + fut.set_exception(GraphStickinessError(error)) + elif error is not None: + fut.set_exception( + GraphDispatchError(f"graph dispatch errored: {error}") + ) + elif cancelled: + fut.set_exception( + GraphDispatchError("graph dispatch cancelled by worker return") + ) + else: + fut.set_result(_PLACEHOLDER) + finally: + if self._on_drained is not None and not self._waiters: + self._on_drained(self) + + def on_first_token( + self, x_correlation_id: str | None, turn_index: int | None + ) -> None: + """Fire the first-token callback parked for one dispatch (AT MOST ONCE). + + Driven by ``GraphIRReplayStrategy._on_graph_first_token`` when the + emitting graph credit's TTFT arrives (post-TTFT anchoring). Pops the + callback registered under ``(x_correlation_id, turn_index)`` and invokes + it, so a successor gated on this node's observed first token is released. + Popping guarantees the single-fire contract: a second call for the same + key (or a return that already dropped the cb) finds nothing. An unknown + key -- a node that carried no first-token successor, or a ``None`` field + on a non-graph fast-path token -- is a graceful no-op. + """ + cb = self._first_token_cbs.pop((x_correlation_id, turn_index), None) + if cb is not None: + # The observer path runs as a fire-and-forget task, so an exception + # here would surface only as an unretrieved-task warning. Log it and + # move on: the cb is already popped, so the successor simply waits for + # its finalize latch instead of the (lost) first-token release. + try: + cb() + except Exception: + _logger.exception( + lambda: ( + "first-token callback raised for key " + f"{(x_correlation_id, turn_index)!r}" + ) + ) + + # ------------------------------------------------------------------ helpers + + def _resolve_ordinal(self, runtime_trace_id: str, node_id: str) -> int | None: + """Map a fired executor node to its build-time ``node_ordinal``. + + Every live producer lowers to a flat LlmNode graph, so the fired node's + bare ``node_id`` IS its catalog key (the executor never descends into + child scopes; ``runtime_trace_id`` is always the bare instance id). + """ + ordinal = node_ordinal_for(self._catalog, self._trace_id, node_id) + if ordinal is None: + _logger.warning( + lambda: ( + f"no node_ordinal for trace={self._trace_id!r} " + f"key={node_id!r} (runtime_trace={runtime_trace_id!r}); " + f"credit will carry node_ordinal=None" + ) + ) + return ordinal + + def _conversation_identity(self, node_id: str) -> str: + """Return the ``conversation_id`` stamped on one fired node's credit. + + The trajectory TEMPLATE id (legacy semantics: stable across recycles, + deliberately duplicated in exports) -- the trace id for root-scope + nodes, ``{trace_id}::{scope}`` for child trajectories. Instance + identity rides ``x_correlation_id``; depth/parent identity is per-NODE + and lives in :meth:`_dag_identity`. Native-authored ids (no + ``{scope}:{turn}`` shape) belong to the root trajectory, mirroring + :meth:`_mint`. + """ + shaped = _split_shaped_id(node_id) + return self._conversation_of(shaped[0] if shaped else self._trace_id) + + def _dag_identity( + self, runtime_trace_id: str, node_id: str + ) -> tuple[int, str | None]: + """Return ``(agent_depth, parent_correlation_id)`` for one fired node. + + A ``node_identity`` hit (dag_jsonl lowering's ``metadata["dag"]`` stamp) + yields the instance's depth; when the stamp names a triggering parent + node, the parent's correlation id is derived from the SAME parts + :meth:`_mint` folds (stable per node + variant -- the counter rides + ``turn_index``, not the string), so the child's + ``parent_correlation_id`` equals the parent node's minted + ``x_correlation_id`` exactly. No map / miss / parent-less stamp falls + back to the legacy root-chain identity: depth from the stamp (or 0) and + the constructor's ``parent_correlation_id`` (weka/dynamo byte-identical + when the map is absent). + """ + identity = ( + None if self._node_identity is None else self._node_identity.get(node_id) + ) + if identity is None: + return 0, self._parent_correlation_id + agent_depth, parent_node_id = identity + if parent_node_id is None: + return agent_depth, self._parent_correlation_id + shaped = _split_shaped_id(parent_node_id) + return agent_depth, self._corr_of(shaped[0] if shaped else self._trace_id) + + def _mint(self, node_id: str, node_ordinal: int | None) -> tuple[str, int, int]: + """Mint ``(x_correlation_id, turn_index, num_turns)`` for one fire -- + legacy semantics. + + ``x_correlation_id`` is the node's TRAJECTORY-INSTANCE id: one fresh + ``{conversation_id}::{uuid4().hex}`` minted lazily per scope per + adapter instance (all turns of one trajectory share it, exactly like a + linear session; recycles get a fresh adapter and therefore fresh + corrs). ``turn_index`` is the node's own 0-based turn within its + trajectory -- the ``{scope}:{turn}`` node id IS the legacy + ``(conversation, turn_index)`` coordinate. Uniqueness of the waiter + key follows from the executor firing each node at most once per + instance run; a re-fire would collide loudly in ``dispatch`` rather + than silently share a waiter. + + A node id WITHOUT the ``{scope}:{turn}`` shape (native-authored graphs + keep author-chosen ids like ``"plan"`` verbatim) maps to the ROOT + trajectory with the catalog ``node_ordinal`` as its turn: the whole + graph rides one sticky session (worker-local dynamic-slot content + stays reachable) and ordinals keep waiter keys unique per node. A + pathological mix of a bare id and a root-scoped ``{trace}:{ordinal}`` + id could collide -- caught loudly by the duplicate-waiter guard. + """ + shaped = _split_shaped_id(node_id) + if shaped is not None: + scope, turn = shaped + else: + scope, turn = self._trace_id, node_ordinal or 0 + # Recorded trajectory turn count from the catalog; a runtime scope the + # catalog does not know (defensive) reads as NON-final -- a wrong + # session-close on the wire is worse than a missing one. + num_turns = self._scope_num_turns.get(scope, turn + 2) + return self._corr_of(scope), turn, max(num_turns, turn + 1) + + def _corr_of(self, scope: str) -> str: + """The trajectory-instance correlation id for ``scope`` (lazy mint).""" + corr = self._scope_corrs.get(scope) + if corr is None: + corr = f"{self._conversation_of(scope)}::{uuid.uuid4().hex}" + self._scope_corrs[scope] = corr + return corr + + def _conversation_of(self, scope: str) -> str: + """Trajectory TEMPLATE id: the trace id for the root scope, else + ``{trace_id}::{scope}`` (corpus-unique; recorded child scopes like + weka ``agent_001`` recur across traces).""" + if scope == self._trace_id: + return scope + return f"{self._trace_id}::{scope}" diff --git a/src/aiperf/graph/dispatch/__init__.py b/src/aiperf/graph/dispatch/__init__.py new file mode 100644 index 0000000000..1cb829ae0d --- /dev/null +++ b/src/aiperf/graph/dispatch/__init__.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-node-kind _execute implementations for the async-dataflow TraceExecutor. + +Each module in this package registers a `_execute_` function onto the +TraceExecutor.__dict__["_execute"] singledispatchmethod via a side-effect at +import time. The TraceExecutor calls `_import_dispatch_modules()` once at +__init__ to trigger registration. + +Per-kind ownership: +- llm.py -> LlmNode (the sole dispatched kind; every live producer lowers + to LlmNode + StaticEdge) +""" + +from __future__ import annotations + + +def _import_dispatch_modules() -> None: + """Import every dispatch module so its registration side-effect runs. + + Called once at module import time (see bottom of this file) and also + invoked from `TraceExecutor.__init__` defensively. Safe to call multiple + times; singledispatch.register is idempotent for the same (cls, func) pair. + """ + # Imports are deferred to function-call time to avoid circular import with + # executor.py during module load. Each module registers its + # _execute_ on import. + # noqa: F401 — imports are for side effect. + from aiperf.graph.dispatch import ( # noqa: F401 + llm, + ) + + +# Auto-trigger registration on first import of this package. Any consumer +# that imports the `dispatch` subpackage thereby gets every node-kind handler +# registered without having to remember to call _import_dispatch_modules() or +# to side-effect-import individual dispatch modules. (`aiperf.graph` alone +# does not import this package; `TraceExecutor.__init__` does, defensively.) +_import_dispatch_modules() diff --git a/src/aiperf/graph/dispatch/llm.py b/src/aiperf/graph/dispatch/llm.py new file mode 100644 index 0000000000..9aa6d965c7 --- /dev/null +++ b/src/aiperf/graph/dispatch/llm.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""LlmNode `_execute` registration for the async-dataflow TraceExecutor. + +Build a `DispatchRequest` (node id), await `credit_issuer.dispatch(node, +request, placement_ctx, first_token_cb=...)`, and write the returned value onto +`node.output` (a type-correct empty list for messages-typed channels, the +placeholder string otherwise). The `first_token_cb` is the per-dispatch stamp +closure (`_make_first_token_stamp`) the adapter fires when this node's +`FirstToken` arrives, releasing any first-token-anchored successor (post-TTFT +anchoring). + +Errors propagate to the executor: `GraphDispatchError` is contained +mid-conversation by `_handle_node_exception` (segment-store parses write a +type-correct sentinel to the node's channels so the conversation continues; +issuer refusals and stickiness errors re-raise as trace-stops). Anything else +(`asyncio.TimeoutError`, transport errors) unwinds the trace: `node.output` +stays UNSET and the `mark_producer_done` in `_fire`'s `finally` orphans +waiters whose count target is unreachable. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +from aiperf.dataset.graph.models import ( + ChannelType, + LlmNode, +) +from aiperf.graph.context import ( + NodeExecutionResult, + Write, + _TraceContext, +) +from aiperf.graph.executor import TraceExecutor +from aiperf.graph.placement import ( + DispatchRequest, + PlacementContext, +) + +__all__ = ["_execute_llm"] + + +def _make_first_token_stamp( + ctx: _TraceContext, node_id: str, loop_wall_us: Callable[[], float] +) -> Callable[[], None]: + """Build the zero-arg release hook a first-token-anchored successor registers. + + Passed as the per-dispatch ``first_token_cb`` kwarg; the dispatch adapter + parks it and invokes it AT MOST ONCE when this node's ``FirstToken`` arrives. + The stamp records the observing loop's wall time and sets the node's + first-token latch so a successor gated on this node's observed first token is + released. A late / duplicate invocation is a guarded no-op: the wall is never + overwritten and the clock is not re-read (the guard returns before the read). + """ + + def _stamp() -> None: + if node_id in ctx.node_first_token_wall_us: + return # late / duplicate first token: no-op + ctx.node_first_token_wall_us[node_id] = loop_wall_us() + ctx.first_token_event(node_id).set() + + return _stamp + + +async def _execute_llm( + self: TraceExecutor, + node: LlmNode, + inputs: dict[str, Any], + ctx: _TraceContext, +) -> NodeExecutionResult: + """Dispatch one LlmNode through the injected credit issuer.""" + spec = self._parsed.graph.state.get(node.output) + channel_type = spec.type if spec is not None else ChannelType.TEXT + + producer_id = next(nid for nid, n in self._parsed.graph.nodes.items() if n is node) + + request = DispatchRequest( + node_id=producer_id, + ) + + placement_ctx = PlacementContext( + parent_trace_id=ctx.trace.id, + parent_node_id=producer_id, + ) + + response = await self._credit_issuer.dispatch( + node, + request, + placement_ctx, + first_token_cb=_make_first_token_stamp(ctx, producer_id, self._loop_wall_us), + ) + + # The credit path resolves a content-free placeholder string (content stays + # worker-side). A messages-typed output channel needs a type-correct empty + # instead: the global `snapshot_at_seq` reduces EVERY channel and + # `add_messages` rejects non-list values. + value: Any = [] if channel_type is ChannelType.MESSAGES else response + writes: list[Write] = [Write(channel=node.output, value=value)] + return NodeExecutionResult(writes=writes) + + +TraceExecutor.__dict__["_execute"].register(LlmNode, _execute_llm) diff --git a/src/aiperf/graph/dynamic_pool.py b/src/aiperf/graph/dynamic_pool.py new file mode 100644 index 0000000000..6fe1dea270 --- /dev/null +++ b/src/aiperf/graph/dynamic_pool.py @@ -0,0 +1,256 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Worker-local dynamic-content pool for graph traces. + +Captured assistant responses keyed ``(trace_id, phase_variant) -> ordinal -> +value`` so a later credit of the SAME trace on the SAME worker (guaranteed by +per-trace sticky routing) can splice a predecessor's actual response into its +prompt. Content never crosses back to the timing plane; this pool is +the graph twin of the worker-cached linear ``UserSession`` state. + +Values are :class:`GraphCapturedReply` (the reply's joined replayable text +plus, for tool_calls / structured replies, the verbatim orjson-serialized +assistant message) or a :class:`GraphPoolSentinel`: ``FAILED`` (dispatch error +/ timeout / cancellation) and ``EMPTY`` (successful response with no +replayable content at all) both splice as omission downstream but stay +distinct for observability. A tool-calls-only reply is NOT ``EMPTY``: it is a +structured capture whose ``message_json`` splices the recorded assistant +message -- ``tool_calls`` included -- into successor prompts verbatim. + +Lifecycle: the router forwards ``GraphTraceEnd`` when the strategy reaps a +trace's dispatch adapter; eviction is DEFERRED while the worker still has +in-flight credits for the trace (their capture writes may land after the end +message on cancelled paths). The byte cap is a load-bearing LRU backstop for +lost end messages (worker death re-routes, version skew): evicting a live +trace's entry surfaces as a loud pool-missing error at its consumer, never a +silent content truncation. +""" + +from __future__ import annotations + +from collections import OrderedDict +from dataclasses import dataclass, field +from enum import Enum +from itertools import count + +__all__ = [ + "GraphCapturedReply", + "GraphDynamicPool", + "GraphPoolMissingError", + "GraphPoolSentinel", +] + + +class GraphPoolMissingError(RuntimeError): + """A dynamic slot referenced a pool entry this worker does not hold. + + Broken stickiness (worker death re-route) or backstop eviction; the worker + converts it to a ``credit_context.error`` with the + ``aiperf.graph.pool_missing:`` prefix, which the dispatch adapter raises + as a non-containable trace error — never a silent omission. + """ + + def __init__(self, src_ordinal: int) -> None: + super().__init__(f"missing dynamic pool value for producer {src_ordinal}") + self.src_ordinal = src_ordinal + + +class GraphPoolSentinel(str, Enum): + FAILED = "failed" + EMPTY = "empty" + + +@dataclass(slots=True, frozen=True) +class GraphCapturedReply: + """One captured assistant reply, structured so splices can reproduce it.""" + + text: str + """Joined replayable text of the reply (``''`` when the reply is + tool_calls-only). Consumed by ``{"sv"}`` block parts (string + concatenation) and as the fallback ``{"s"}`` splice content.""" + + message_json: str | None = None + """Verbatim orjson-serialized assistant message when the endpoint's + ``build_assistant_turn`` returned ``raw_messages`` (tool_calls / + structured replies); ``None`` for plain-text replies. ``{"s"}`` splices + ``orjson.loads(message_json)`` verbatim, byte-matching the legacy + child-seed rendering.""" + + +PoolValue = GraphCapturedReply | GraphPoolSentinel + +_SENTINEL_COST_BYTES = 8 + +# Default bound on the recently-ended-key tombstone FIFO (see +# `GraphDynamicPool.__init__`). Keys are two short strings (~100 B each), so +# the worst-case footprint is a few hundred KiB per worker. +_TOMBSTONE_CAPACITY = 4096 + + +def _value_bytes(value: PoolValue) -> int: + if isinstance(value, GraphPoolSentinel): + return _SENTINEL_COST_BYTES + return len(value.text.encode()) + ( + len(value.message_json.encode()) if value.message_json else 0 + ) + + +@dataclass(slots=True) +class _TraceEntry: + """One trace execution's captured values plus its lifecycle state.""" + + values: dict[int, PoolValue] = field(default_factory=dict) + """Captured value per node ordinal.""" + total_bytes: int = 0 + """Byte cost of ``values`` (UTF-8 text + message JSON length; flat cost + per sentinel).""" + inflight: int = 0 + """Credits of this trace currently processing on this worker.""" + end_pending: bool = False + """A ``GraphTraceEnd`` arrived while credits were still in flight.""" + last_touch: int = 0 + """Monotonic LRU stamp (pool-level counter, not wall time).""" + + +class GraphDynamicPool: + """Byte-capped, LRU-backstopped store of captured graph responses.""" + + def __init__( + self, max_bytes: int, tombstone_capacity: int = _TOMBSTONE_CAPACITY + ) -> None: + self._max_bytes = max_bytes + self._entries: dict[tuple[str, str], _TraceEntry] = {} + self._total_bytes = 0 + self._clock = count() + # FIFO tombstones of recently-ENDED keys: `credit_started` / `put` for + # a tombstoned key is a no-op, so a straggler event racing behind its + # `GraphTraceEnd` (cancelled dispatch's start/capture landing after the + # end) cannot resurrect a zero-byte immortal entry. Bounded so recycle- + # heavy runs don't trade the entry leak for a tombstone leak; a key + # aged out of the FIFO could in principle resurrect, but by then its + # straggler window is long past. + self._tombstone_capacity = tombstone_capacity + self._tombstones: OrderedDict[tuple[str, str], None] = OrderedDict() + + @property + def total_bytes(self) -> int: + return self._total_bytes + + @property + def entry_count(self) -> int: + """Number of live trace entries (test/diagnostic).""" + return len(self._entries) + + def get( + self, trace_id: str, phase_variant: str, node_ordinal: int + ) -> PoolValue | None: + """The captured value for one node, or ``None`` when absent. + + ``None`` means missing (never captured, or evicted) -- the Phase 3 + consumer converts it to a loud pool-missing error, never a default. + """ + entry = self._entries.get((trace_id, phase_variant)) + if entry is None: + return None + entry.last_touch = next(self._clock) + return entry.values.get(node_ordinal) + + def put( + self, + trace_id: str, + phase_variant: str, + node_ordinal: int, + value: PoolValue, + ) -> None: + key = (trace_id, phase_variant) + entry = self._entries.get(key) + if entry is None: + if key in self._tombstones: + # Straggler capture for an already-ENDED trace: nothing will + # ever read it (the end already reaped the consumers), so + # storing it would recreate an immortal entry. + return + entry = _TraceEntry() + self._entries[key] = entry + previous = entry.values.get(node_ordinal) + if previous is not None: + delta = _value_bytes(previous) + entry.total_bytes -= delta + self._total_bytes -= delta + entry.values[node_ordinal] = value + added = _value_bytes(value) + entry.total_bytes += added + self._total_bytes += added + entry.last_touch = next(self._clock) + self._enforce_cap(protect=key) + + def credit_started(self, trace_id: str, phase_variant: str) -> None: + key = (trace_id, phase_variant) + entry = self._entries.get(key) + if entry is None: + if key in self._tombstones: + # The trace already ENDED; a late credit_started (cancelled + # dispatch racing behind GraphTraceEnd) must not resurrect a + # zero-byte entry with no end message left to evict it. + return + entry = _TraceEntry() + self._entries[key] = entry + entry.inflight += 1 + entry.last_touch = next(self._clock) + + def credit_finished(self, trace_id: str, phase_variant: str) -> None: + key = (trace_id, phase_variant) + entry = self._entries.get(key) + if entry is None: + return + entry.inflight = max(0, entry.inflight - 1) + if entry.end_pending and entry.inflight == 0: + self._evict(key) + self._mark_ended(key) + + def trace_end(self, trace_id: str, phase_variant: str) -> None: + """Handle ``GraphTraceEnd``: evict now, or defer until in-flight drains. + + Deferral covers the cancelled paths where a credit's capture write can + land AFTER the end message (the strategy reaps on adapter drain, but a + locally-cancelled dispatch leaves the worker task still running). + Idempotent; unknown traces still tombstone the key so straggler events + arriving after the end cannot resurrect an entry. + """ + key = (trace_id, phase_variant) + entry = self._entries.get(key) + if entry is not None and entry.inflight > 0: + entry.end_pending = True + return + self._evict(key) + self._mark_ended(key) + + def _mark_ended(self, key: tuple[str, str]) -> None: + """Tombstone ``key`` as recently ended (bounded FIFO, oldest dropped).""" + self._tombstones[key] = None + self._tombstones.move_to_end(key) + while len(self._tombstones) > self._tombstone_capacity: + self._tombstones.popitem(last=False) + + def _evict(self, key: tuple[str, str]) -> None: + entry = self._entries.pop(key, None) + if entry is not None: + self._total_bytes -= entry.total_bytes + + def _enforce_cap(self, protect: tuple[str, str]) -> None: + """LRU-evict whole trace entries until under the byte cap. + + The entry just written is evicted only as the last resort (it alone + exceeds the cap); its consumers then fail loudly at splice time. + """ + if self._total_bytes <= self._max_bytes: + return + by_lru = sorted(self._entries.items(), key=lambda kv: kv[1].last_touch) + for key, _entry in by_lru: + if self._total_bytes <= self._max_bytes: + return + if key == protect: + continue + self._evict(key) + if self._total_bytes > self._max_bytes: + self._evict(protect) diff --git a/src/aiperf/graph/executor.py b/src/aiperf/graph/executor.py new file mode 100644 index 0000000000..760e6c0206 --- /dev/null +++ b/src/aiperf/graph/executor.py @@ -0,0 +1,603 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TraceExecutor - async-dataflow per-trace executor for the graph package. + +Nodes fire as soon as their input channels are ready. Per-trace +state lives on `_TraceContext`; the executor itself is shared across traces in +a phase (all graph-derived state is immutable post-`__init__`). + +Per-kind ``_execute`` implementations live in +``aiperf.graph.dispatch``; see that package's +``__init__.py`` for the ownership map. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from functools import singledispatchmethod +from typing import TYPE_CHECKING, Any + +from aiperf.common.aiperf_logger import AIPerfLogger +from aiperf.common.clock import AIPerfClock, WallClock +from aiperf.common.environment import Environment +from aiperf.dataset.graph.models import ( + ChannelType, + LlmNode, + NodeUnion, + ParsedGraph, +) +from aiperf.graph.channel_store import ( + VersionedChannelStore, +) +from aiperf.graph.channels import ( + producers_per_channel as compute_producers_per_channel, +) +from aiperf.graph.context import ( + NodeExecutionResult, + _NodeExpectedExit, + _TraceContext, +) +from aiperf.graph.scheduler import Scheduler + +if TYPE_CHECKING: + from aiperf.dataset.graph.models import ( + ChannelSpec, + TraceRecord, + ) + +__all__ = ["TraceExecutor", "TraceResult"] + + +def _default_store_factory( + channel_specs: dict[str, ChannelSpec], + initial: dict[str, Any], + producers_per_channel: dict[str, int], +) -> VersionedChannelStore: + return VersionedChannelStore( + initial=initial, + channel_specs=channel_specs, + producers_per_channel=producers_per_channel, + ) + + +@dataclass(slots=True) +class TraceResult: + """The return value of `TraceExecutor.run`. + + Carries the trace id and the post-run channel snapshot. + """ + + trace_id: str + channels: dict[str, Any] + + +_logger = AIPerfLogger(__name__) + + +class TraceExecutor: + """Async-dataflow trace executor. Shared across traces in a phase. + + All mutable per-trace state lives on `_TraceContext`; constructor-injected + collaborators are graph-derived and immutable. Node-kind ``_execute`` + bodies are registered on the dispatch table by sibling modules under + ``dispatch/`` without touching this class (each + ``_execute.register()`` lives in its own dispatch module). + """ + + def __init__( + self, + parsed: ParsedGraph, + *, + credit_issuer: Any = None, + scheduler: Scheduler | None = None, + producers_per_channel: dict[str, int] | None = None, + compress_edge_delays: bool = False, + absolute_start_offsets: bool = False, + clock: AIPerfClock | None = None, + ) -> None: + from aiperf.graph.dispatch import _import_dispatch_modules + + _import_dispatch_modules() + self._parsed = parsed + # Time source for the firing loop. Defaults to WallClock (reads + # ``time.perf_counter_ns`` + ``asyncio.sleep`` -- behavior identical to + # the prior direct ``time.*`` calls). A VirtualClock can be injected so + # a driver pump fast-forwards sim time, letting a multi-hour recorded + # trace replay in milliseconds while reproducing the exact firing + # timeline. + self._clock: AIPerfClock = clock if clock is not None else WallClock() + # Anchor node-level ``min_start_delay_us`` to the INSTANCE run-start (the + # t* / phase-start anchor) instead of each node's input-ready instant. + # The PROFILING snapshot rewrite (``timing/snapshot_chop.py``, + # adapter-agnostic) stamps every surviving frontier + # turn's ``min_start_delay_us`` as its ABSOLUTE ``dispatch_offset_us`` from + # t*. + # Measuring it from input-ready double-counts the lead for any stream + # whose first turn's inputs arrive late (a co-scoped subagent/worker + # gated on its spawn), so those streams drift out of recorded-time order + # and occupy lanes longer than recorded. Anchoring to the shared instance + # run-start makes every stream's frontier turn fire at its absolute + # recorded offset (clamped to input-readiness). Off by + # default == relative semantics. + self._absolute_start_offsets = absolute_start_offsets + self._anchor_wall_us: float | None = None + # Per-instance counterpart of the global ``AIPERF_GRAPH_IGNORE_EDGE_DELAYS`` + # env: when True this single executor collapses ALL incoming firing-edge + # delays (zero-idle / burst pacing), independent of the env default. The + # accelerated cache-pressure warmup sets this True ONLY on + # the WARMUP-phase executors (knob-gated), so the live trajectories replay + # with zero idle delay to drive the server's KV cache to pressure. Default + # False == honor every captured edge delay exactly (byte-for-byte). + self._compress_edge_delays = compress_edge_delays + # `scheduler` and `producers_per_channel` are pure functions of the + # immutable `parsed.graph`. When a + # caller already memoized them per template_id, inject them to skip the + # two O(graph) rescans on this per-bind instance. Both are read-only + # after construction, so sharing them across instances is safe; only the + # mutable `_credit_issuer` (monkeypatched per-run) stays per-instance. + self._scheduler = scheduler or Scheduler(parsed.graph) + self._credit_issuer = credit_issuer + self._producers_per_channel = ( + producers_per_channel + if producers_per_channel is not None + else compute_producers_per_channel(parsed.graph) + ) + + # ------------------------------------------------------------------ + # run + # ------------------------------------------------------------------ + async def run( + self, + trace: TraceRecord, + ) -> TraceResult: + """Drive one trace through the dataflow firing loop. + + Opens a per-trace `asyncio.TaskGroup`, schedules every entry node + returned by `Scheduler.entry_nodes()`, and blocks until every `_fire` + task finishes. The channel snapshot is captured once the frontier + drive completes; a cancellation or unhandled trace error propagates + out of `run` before a `TraceResult` is built. + + When `Environment.GRAPH.EXECUTOR_WATCHDOG_TIMEOUT` is set, the + frontier drive is bounded by that wall-clock deadline: a firing loop + wedged on an unsatisfiable channel input (a producer-accounting bug) + raises `asyncio.TimeoutError` instead of hanging. Default `None` + preserves the faithful, unbounded idle-gap replay bare/count/session + runs rely on (see `GraphIRReplayStrategy` AgentX count-mode parity). + """ + # Pin the instance t* origin once, on the first run, so every firing in + # this executor shares one absolute-offset anchor (AgentX dispatches all + # streams from t*). + if self._absolute_start_offsets and self._anchor_wall_us is None: + self._anchor_wall_us = self._loop_wall_us() + + store = _default_store_factory( + self._parsed.graph.state, + dict(trace.initial_state), + self._producers_per_channel, + ) + + ctx = _TraceContext(trace=trace, store=store) + + watchdog_s = Environment.GRAPH.EXECUTOR_WATCHDOG_TIMEOUT + if watchdog_s is not None: + await asyncio.wait_for(self._drive_frontier(ctx), timeout=watchdog_s) + else: + await self._drive_frontier(ctx) + + return TraceResult( + trace_id=trace.id, + channels=store.snapshot(), + ) + + async def _drive_frontier(self, ctx: _TraceContext) -> None: + """Open the per-trace TaskGroup and drive the readiness frontier. + + Schedules every entry node and blocks until every `_fire` finishes. + """ + async with asyncio.TaskGroup() as tg: + ctx.tg = tg + for entry_id in self._scheduler.entry_nodes(): + self._schedule(entry_id, ctx) + + def _schedule(self, node_id: str, ctx: _TraceContext) -> None: + """Schedule a node's `_fire` task on the per-trace TaskGroup. + + Dedup semantics combine §5.1 (AND-fan-in: silent skip when the same + successor is reached via multiple predecessors) with §5.2 (cycle + guard: re-entering an already-COMPLETED node is a `RuntimeError`). + The disambiguator is `tasks_by_node_id[node_id].done()`: + + - In-flight task or no task entry yet -> legitimate concurrent + scheduling from a fan-in predecessor; silently skip. + - Completed task -> cycle; the loader's cycle check should have + rejected this graph, so raise `RuntimeError`. + """ + if node_id in ctx.scheduled_node_ids: + existing = ctx.tasks_by_node_id.get(node_id) + if existing is not None and existing.done(): + raise RuntimeError( + f"cycle detected: node {node_id!r} re-scheduled after " + "completing. The loader's cycle check should have " + "rejected this graph; a mixed-anchor fan-in (one " + "start-anchored plus one completion in-edge on the same " + "target, rejected at Scheduler construction) is another " + "likely cause." + ) + # In-flight: AND-fan-in dedup, silent no-op. + return + + ctx.scheduled_node_ids.add(node_id) + assert ctx.tg is not None # set inside run()'s TaskGroup + task = ctx.tg.create_task( + self._fire(node_id, ctx), name=f"fire:{ctx.trace.id}:{node_id}" + ) + ctx.tasks_by_node_id[node_id] = task + + async def _prepare_node_inputs( + self, node_id: str, node: NodeUnion, ctx: _TraceContext + ) -> dict[str, Any]: + requirements = node.inputs + capture = await ctx.store.await_inputs(requirements) + gate_seq = ctx.store.current_seq + node_firable_wall_us = self._loop_wall_us() + await self._apply_firing_delay(node_id, ctx, node_firable_wall_us) + ctx.node_dispatch_wall_us[node_id] = self._loop_wall_us() + for succ_id in self._scheduler.start_anchored_successors(node_id): + self._schedule(succ_id, ctx) + if isinstance(node, LlmNode): + return ctx.store.snapshot_at_seq(gate_seq) + return ctx.store.read(requirements, capture) + + # ------------------------------------------------------------------ + # _fire + # ------------------------------------------------------------------ + async def _fire(self, node_id: str, ctx: _TraceContext) -> None: + """Drive one node through the dataflow firing path.""" + node = self._parsed.graph.nodes[node_id] + success = False + wrote_channels: set[str] = set() + try: + result, wrote_channels = await self._run_node(node_id, node, ctx) + success = result is not None + finally: + self._finalize_node( + node_id, + node, + ctx, + success=success, + wrote_channels=wrote_channels, + ) + if result is None: + return + self._schedule_successors(node_id, result, ctx) + + async def _run_node( + self, + node_id: str, + node: NodeUnion, + ctx: _TraceContext, + ) -> tuple[NodeExecutionResult | None, set[str]]: + from aiperf.graph.channel_store import ChannelOrphanedError + from aiperf.graph.credit_dispatch_adapter import _NodeOverflowTerminate + + try: + inputs = await self._prepare_node_inputs(node_id, node, ctx) + result = await self._execute(node, inputs, ctx) + wrote_channels = self._publish_writes(node_id, result, ctx) + return result, wrote_channels + except _NodeOverflowTerminate: + # This node's response was a context-overflow. Flag the trace so + # the downstream orphan cascade (successors awaiting this node's + # never-written output) is also treated as a clean stop, and exit + # cleanly WITHOUT scheduling successors -- the trajectory terminates + # here, since later turns carry even more context and would only + # overflow too. + ctx.overflow_terminated = True + return None, set() + except _NodeExpectedExit: + return None, set() + except ChannelOrphanedError: + # A successor of an overflow-terminated node will orphan on its + # never-written input channel. That is the EXPECTED downstream + # consequence of the early-termination, not a trace error: swallow it as + # a clean exit so the whole trajectory stops without inflating + # ``errored_traces``. Outside an overflow termination, re-raise so a + # genuine producer-failure orphan still surfaces. + if ctx.overflow_terminated: + return None, set() + raise + except Exception as exc: + return self._handle_node_exception(node_id, node, ctx=ctx, exc=exc) + + def _publish_writes( + self, node_id: str, result: NodeExecutionResult, ctx: _TraceContext + ) -> set[str]: + wrote_channels: set[str] = set() + for write in result.writes: + ctx.store.write([write.channel], write.value, writer_node_id=node_id) + wrote_channels.add(write.channel) + return wrote_channels + + def _handle_node_exception( + self, + node_id: str, + node: NodeUnion, + *, + ctx: _TraceContext, + exc: Exception, + ) -> tuple[NodeExecutionResult | None, set[str]]: + from aiperf.graph.credit_dispatch_adapter import ( + CreditIssueRefusedError, + GraphDispatchError, + GraphStickinessError, + ) + + # Issuer refusal (duration / request-count caps, cancellation) and + # broken stickiness (missing dynamic-pool value) are TRACE-STOPS, + # never contained: sentinel-continuing past a refusal would churn + # every remaining node without dispatching, and past a pool miss + # would silently corrupt the workload's content dependencies. + # Checked before the containment branch (both subclass + # GraphDispatchError so the untyped check below would swallow them). + if isinstance(exc, (CreditIssueRefusedError, GraphStickinessError)): + raise exc + + # Mid-conversation resilience: a dispatch/request FAILURE (connection + # reset, HTTP error, worker-cancel -> GraphDispatchError) + # CONTAINS to this node instead of unwinding the whole per-trace + # TaskGroup (which would lose every remaining turn + subagent of the + # conversation). The failed request is ALREADY recorded as an error by + # the RecordProcessor; the conversation continues to its next turn. + # Returning a result (not raising) lets _finalize_node mark this node's + # output channels producer-done so downstream await_inputs resolve, and + # _schedule_successors fires the next turn. Keeping the instance alive + # also preserves its adapter, so any still-in-flight returns still + # resolve (no "no live adapter" drops). Faithful: downstream weka turns + # splice the RECORDED @messages delta, not this node's live response, so + # a failed turn does not corrupt downstream content. Structural / + # programming errors (not GraphDispatchError) still raise. + if isinstance(exc, GraphDispatchError): + _logger.warning( + lambda: ( + f"node {node_id!r} dispatch failed; continuing " + f"conversation past the failed turn: {exc!r}" + ) + ) + # Segment-store IRs ONLY: still PRODUCE this node's output channels + # (a content-neutral sentinel) so downstream AND-fan-in readers + # unblock instead of orphaning. These IRs wire each turn to read its + # predecessors' ``{nid}_out`` channels purely as a completion GATE + # (prompts come from the segment pool, not the channel), so an empty + # sentinel lets the conversation continue past the failed turn + # rather than cascade-erroring the whole trace. The sentinel is + # TYPE-CORRECT per channel: ``[]`` for messages-typed channels (the + # global ``snapshot_at_seq`` reduces EVERY channel and + # ``add_messages`` rejects non-list values), ``None`` for the rest + # (never read). Non-segment-store parses + # keep the prior behavior (return empty -> downstream orphans -> + # trace errors). + if getattr(self._parsed, "segment_pool", None) is not None: + wrote_channels: set[str] = set() + for channel in node.write_channels: + spec = self._parsed.graph.state.get(channel) + sentinel = ( + [] + if spec is not None and spec.type is ChannelType.MESSAGES + else None + ) + ctx.store.write([channel], sentinel, writer_node_id=node_id) + wrote_channels.add(channel) + return NodeExecutionResult(), wrote_channels + return NodeExecutionResult(), set() + raise exc + + def _finalize_node( + self, + node_id: str, + node: NodeUnion, + ctx: _TraceContext, + *, + success: bool, + wrote_channels: set[str], + ) -> None: + ctx.node_finish_wall_us[node_id] = self._loop_wall_us() + # Latch the first-token event even when no first token was observed + # (error/cancel/non-streaming terminal). Idempotent: the stamp closure + # may have already set it. A successor with a first-token-anchored edge + # awaits this latch; the absence of a `node_first_token_wall_us` entry + # tells its gate computation to fall back to the dispatch anchor. + ctx.first_token_event(node_id).set() + for channel in node.write_channels: + ctx.store.mark_producer_done( + channel, + success=success, + wrote=channel in wrote_channels, + ) + + def _schedule_successors( + self, node_id: str, result: NodeExecutionResult, ctx: _TraceContext + ) -> None: + for succ_id in self._scheduler.successors_after(node_id): + self._schedule(succ_id, ctx) + + # ------------------------------------------------------------------ + # edge delay + # ------------------------------------------------------------------ + async def _apply_firing_delay( + self, + node_id: str, + ctx: _TraceContext, + node_firable_wall_us: float, + ) -> None: + """Sleep until ``node_id``'s incoming-edge gate clears. + + Edge-gate semantics, anchored to the dataflow loop: + + - ``edge.delay_after_predecessor_us``: gate >= predecessor finish + wall + delay. The predecessor's finish wall is recorded in its + ``_fire`` ``finally`` block on + ``ctx.node_finish_wall_us[edge.source]``. Because the successor + is scheduled after the predecessor's ``mark_producer_done`` + (same ``finally``), the write happens-before the successor's + ``_apply_firing_delay`` read on the single asyncio loop — no + lock needed. + - ``edge.min_start_delay_us``: gate >= ``node_firable_wall_us`` + (the moment ``await_inputs`` returned) + delay, measured from the + moment the node became input-ready. + - ``edge.delay_after_predecessor_start_us``: gate >= predecessor + DISPATCH wall (``ctx.node_dispatch_wall_us[edge.source]``, stamped + when its firing gate cleared) + delay; the successor was scheduled at + that dispatch, so the stamp happens-before this read. + - ``edge.delay_after_predecessor_first_token_us``: gate >= predecessor + OBSERVED FIRST-TOKEN wall (``ctx.node_first_token_wall_us[edge.source]``, + stamped by the dispatch adapter's first-token cb) + delay. Refines a + start-anchored edge: this method first awaits the source's first-token + latch, so the wall (or its absence -- a terminal without a first token) + is settled before the gate is computed. An observed first token + SUPERSEDES the dispatch anchor for that edge; a source that terminated + without one has no wall entry, so the dispatch-anchor fallback applies. + - Node-level ``min_start_delay_us``: gate >= ``node_firable_wall_us`` + + node delay. + + AND-fan-in: every incoming static edge contributes a gate; the runtime + takes the ``max``. + + ``AIPERF_GRAPH_IGNORE_EDGE_DELAYS=1`` short-circuits ALL of the + above uniformly — useful for "how fast on infinite hardware" + A/B comparison. The per-instance ``compress_edge_delays`` flag + (set by the accelerated cache-pressure warmup on + WARMUP-phase executors only) does the same for a single executor. + """ + if Environment.GRAPH.IGNORE_EDGE_DELAYS or self._compress_edge_delays: + return + # First-token-anchored edges gate on the source's OBSERVED first token, + # so wait for the source's first-token latch (set by its stamp closure + # OR by ``_finalize_node`` when it terminates without one) before + # computing the gate. Placed AFTER the ignore/compress early return so + # compressed replays skip the wait entirely -- their dispatch-time + # scheduling already happened via ``start_anchored_successors``. + for edge in self._scheduler.incoming_static_edges(node_id): + if edge.delay_after_predecessor_first_token_us is not None: + await ctx.first_token_event(edge.source).wait() + gate_us = self._compute_firing_gate_us(node_id, ctx, node_firable_wall_us) + if gate_us <= 0.0: + return + wait_us = gate_us - self._loop_wall_us() + if wait_us <= 0: + return + # Replay the recorded firing gate FAITHFULLY (verbatim). The edge delays + # are already on the build-plane per-gap-warped clock (the shared + # ``ActiveIdleWarp`` in ``segment_ir/trie_content.py``, byte-faithful to + # agentx's ``_IdleGapTimeWarp`` per-gap cap); there is + # NO runtime clamp here. A faithfully-replayed trace can therefore span + # its full recorded wall time -- the GraphIRReplayStrategy bounds phase + # completion by the AgentX stop-condition model (duration / count), so a + # long idle-gap replay finalizes on the stop condition rather than being + # truncated mid-gap by an unfaithful per-node sleep cap. + # + # Sleeps on the injected clock (``WallClock`` -> ``asyncio.sleep``; + # ``VirtualClock`` -> park until the driver pump advances sim time past + # this gate). ``wait_us`` is microseconds; the clock takes ns. + await self._clock.sleep_ns(int(wait_us * 1_000.0)) + + def _compute_firing_gate_us( + self, + node_id: str, + ctx: _TraceContext, + node_firable_wall_us: float, + ) -> float: + """Return the max incoming-edge gate wall (us) for ``node_id``. + + AND-fan-in: every incoming static edge plus the node-level + ``min_start_delay_us`` contributes a gate; the runtime takes the ``max``. + See :meth:`_apply_firing_delay` for the edge-gate semantics. + """ + gate_us = 0.0 + finishes = ctx.node_finish_wall_us + dispatches = ctx.node_dispatch_wall_us + first_tokens = ctx.node_first_token_wall_us + for edge in self._scheduler.incoming_static_edges(node_id): + if edge.delay_after_predecessor_first_token_us is not None: + ft_us = first_tokens.get(edge.source) + if ft_us is not None: + gate_us = max( + gate_us, + ft_us + edge.delay_after_predecessor_first_token_us, + ) + # Observed first token supersedes the dispatch fallback for + # this edge. A source that terminated WITHOUT a first token + # has no wall entry, so fall through to the dispatch-anchor + # block below. + continue + if edge.delay_after_predecessor_us is not None: + finish_us = finishes.get(edge.source) + if finish_us is not None: + gate_us = max(gate_us, finish_us + edge.delay_after_predecessor_us) + if edge.min_start_delay_us is not None: + gate_us = max(gate_us, node_firable_wall_us + edge.min_start_delay_us) + if edge.delay_after_predecessor_start_us is not None: + dispatch_us = dispatches.get(edge.source) + if dispatch_us is not None: + gate_us = max( + gate_us, + dispatch_us + edge.delay_after_predecessor_start_us, + ) + node = self._parsed.graph.nodes.get(node_id) + node_min_start = node.min_start_delay_us if node is not None else None + if node_min_start is not None: + # In absolute mode the snapshot rewrite stamped this as the firing's + # ABSOLUTE dispatch_offset_us from t*; anchor it to the shared + # instance run-start (t* origin) so a late-input co-scoped stream + # fires at its recorded offset, not input-ready + offset (which + # double-counts the lead and drifts it out of recorded-time order). + # The wait is still clamped to input-readiness upstream (await_inputs + # precedes this gate), so the node never fires before its data. + anchor = ( + self._anchor_wall_us + if self._absolute_start_offsets and self._anchor_wall_us is not None + else node_firable_wall_us + ) + gate_us = max(gate_us, anchor + node_min_start) + return gate_us + + def _loop_wall_us(self) -> float: + """Return the firing-loop clock reading in microseconds. + + Reads the injected ``AIPerfClock`` (``now_ns``) so the gate + computation and any caller-side telemetry share ONE time source. With + the default ``WallClock`` this is ``time.perf_counter_ns() / 1_000`` + (monotonic wall, behavior unchanged); with a ``VirtualClock`` it is + sim time advanced by the driver pump. + """ + return self._clock.now_ns() / 1_000.0 + + # ------------------------------------------------------------------ + # dispatch table + # ------------------------------------------------------------------ + # LlmNode's `_execute` body lives in `dispatch/llm.py`. The stub below is + # a fallback raiser that fires only if the dispatch module failed to load; + # any other node kind hits the singledispatch default (every live producer + # lowers to LlmNode-only graphs). + + @singledispatchmethod + async def _execute( + self, + node: Any, + inputs: dict[str, Any], + ctx: _TraceContext, + ) -> NodeExecutionResult: + raise NotImplementedError( + f"no _execute registered for node kind {type(node).__name__!r}" + ) + + @_execute.register + async def _execute_llm( + self, node: LlmNode, inputs: dict[str, Any], ctx: _TraceContext + ) -> NodeExecutionResult: + raise NotImplementedError( + "LlmNode dispatch was not loaded; ensure 'import " + "aiperf.graph.dispatch' ran before " + "TraceExecutor instantiation" + ) diff --git a/src/aiperf/graph/placement.py b/src/aiperf/graph/placement.py new file mode 100644 index 0000000000..3b64de8f58 --- /dev/null +++ b/src/aiperf/graph/placement.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dispatch payload/context carriers for the graph runtime. + +`DispatchRequest` / `PlacementContext` carry the per-call payload and runtime +context that the live dispatch path (`graph.dispatch.llm`) threads into the +credit issuer. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +__all__ = [ + "DispatchRequest", + "PlacementContext", +] + + +@dataclass(slots=True) +class DispatchRequest: + """One dispatch: the fired node id. + + The credit path resolves a node's actual request worker-side from the + recorded envelope keyed by ``node_id``. + """ + + node_id: str + + +@dataclass(slots=True) +class PlacementContext: + """Per-dispatch routing context the executor stamps for the credit path. + + Populated by the executor; not authorable. The credit adapter reads only + ``parent_trace_id`` (the per-trace instance id); ``parent_node_id`` rides + along as the fired node's identity. + """ + + parent_trace_id: str | None = None + parent_node_id: str | None = None diff --git a/src/aiperf/graph/reducers.py b/src/aiperf/graph/reducers.py new file mode 100644 index 0000000000..4002b3d5e8 --- /dev/null +++ b/src/aiperf/graph/reducers.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Graph reducers for overwrite and add_messages channels.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +# Sentinel marking a channel that has never been written. Compare with +# ``is UNSET``; msgspec's singleton keeps identity across pickle/deepcopy. +from msgspec import UNSET + +Reducer = Callable[[Any, list[tuple[str, Any]]], Any] + + +class OverwriteConflictError(ValueError): + """Raised when two or more nodes write an overwrite-typed channel concurrently.""" + + +class ReducerNameError(ValueError): + """Raised when a reducer name is unknown or reserved for a future spec version.""" + + +def overwrite_reducer(current: Any, writes: list[tuple[str, Any]]) -> Any: + """Return the single writer's value; reject concurrent multi-writer.""" + if not writes: + return current + if len(writes) > 1: + ids = ", ".join(w[0] for w in writes) + raise OverwriteConflictError( + f"overwrite-typed channel written by multiple nodes concurrently: {ids}" + ) + return writes[0][1] + + +class MessageList(list): + """`list` subclass for add_messages accumulators with an `id -> index` map. + + The map (`_id_index`) is internal bookkeeping that lets + `add_messages_reducer` replace by-id in O(1) instead of O(n). Channel + readers consume the value as a plain list — `isinstance(x, list)` + passes and all list ops behave normally. + + The index is rebuilt on `__init__`/`__deepcopy__`/`__reduce__` so + copy/pickle round-trips stay correct even when the resulting object + is a `MessageList`. Mutating the list directly (append, __setitem__, + etc.) bypasses the index, so all production mutation must go through + the reducer; the reducer is the only writer. + """ + + __slots__ = ("_id_index",) + + def __init__(self, iterable: Any = ()) -> None: + super().__init__(iterable) + self._id_index: dict[Any, int] = {} + for i, msg in enumerate(self): + if isinstance(msg, dict): + msg_id = msg.get("id") + if msg_id is not None: + self._id_index[msg_id] = i + + def __deepcopy__(self, memo: dict[int, Any]) -> MessageList: + import copy as _copy + + return MessageList(_copy.deepcopy(list(self), memo)) + + def __reduce__(self) -> tuple[Any, tuple[list[Any]]]: + return (MessageList, (list(self),)) + + +def add_messages_reducer( + current: Any, writes: list[tuple[str, Any]] +) -> list[dict[str, Any]]: + """Append writer values; replace prior messages whose `id` matches a new message. + + Maintains an `id -> index` map on the accumulator (`MessageList`) so + id-based replacement is O(1). Messages without `id` take the plain + fast-append path with no dict overhead beyond a sentinel check. + """ + if isinstance(current, MessageList): + accumulator = current + elif current is UNSET: + accumulator = MessageList() + else: + accumulator = MessageList(current) + id_index = accumulator._id_index + for writer_id, value in writes: + if not isinstance(value, list): + raise TypeError( + f"add_messages reducer expected a list for writer '{writer_id}', got {type(value).__name__}" + ) + for msg in value: + msg_id = msg.get("id") if isinstance(msg, dict) and "id" in msg else None + if msg_id is not None: + idx = id_index.get(msg_id) + if idx is not None: + accumulator[idx] = msg + continue + id_index[msg_id] = len(accumulator) + accumulator.append(msg) + return accumulator + + +_REDUCERS: dict[str, Reducer] = { + "overwrite": overwrite_reducer, + "add_messages": add_messages_reducer, +} + + +def get_reducer(name: str) -> Reducer: + """Look up a reducer function by name; raise ReducerNameError if unknown.""" + if name not in _REDUCERS: + raise ReducerNameError(f"unknown reducer '{name}'") + return _REDUCERS[name] diff --git a/src/aiperf/graph/scheduler.py b/src/aiperf/graph/scheduler.py new file mode 100644 index 0000000000..9ecefa38b3 --- /dev/null +++ b/src/aiperf/graph/scheduler.py @@ -0,0 +1,194 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Graph-derived adjacency helper for the async-dataflow executor. + +A pure adjacency lookup over the parsed graph's static edges; holds no mutable +per-trace state. Per-trace firing dedup lives on +`_TraceContext.scheduled_node_ids`. +""" + +from __future__ import annotations + +from collections import defaultdict + +import msgspec + +from aiperf.dataset.graph.models import ( + END_NODE_ID, + START_NODE_ID, + GraphRecord, + StaticEdge, +) + + +class Scheduler: + """Graph-derived adjacency helper for the async-dataflow executor. + + All state is derived from the parsed graph at construction time and is + immutable across traces (no mutable per-trace fields). Per-trace state + lives on `_TraceContext.scheduled_node_ids`; this class is shared across + every trace in a phase. + """ + + __slots__ = ( + "_static_succ", + "_start_anchored_succ", + "_static_pred_edges", + "_entry", + ) + + def __init__(self, graph: GraphRecord) -> None: + static_succ: dict[str, list[str]] = defaultdict(list) + start_anchored_succ: dict[str, list[str]] = defaultdict(list) + static_pred_edges: dict[str, list[StaticEdge]] = defaultdict(list) + + for edge in graph.edges: + if isinstance(edge, StaticEdge): + if edge.delay_after_predecessor_start_us is not None: + start_anchored_succ[edge.source].append(edge.target) + else: + static_succ[edge.source].append(edge.target) + static_pred_edges[edge.target].append(edge) + + _reject_mixed_anchor_fan_in(static_pred_edges) + + entry: list[str] = [] + seen: set[str] = set() + for target in static_succ.get(START_NODE_ID, []): + if target == END_NODE_ID or target in seen: + continue + seen.add(target) + entry.append(target) + + self._static_succ: dict[str, list[str]] = dict(static_succ) + self._start_anchored_succ: dict[str, list[str]] = dict(start_anchored_succ) + self._static_pred_edges: dict[str, list[StaticEdge]] = dict(static_pred_edges) + self._entry: list[str] = entry + + def entry_nodes(self) -> list[str]: + """Return node ids that should fire at trace start (successors of START). + + END is suppressed. + """ + return list(self._entry) + + def successors_after(self, node_id: str) -> list[str]: + """Return successor node ids that should fire after `node_id` completes. + + Static successors via StaticEdge. Start-anchored StaticEdges + (`delay_after_predecessor_start_us`) are EXCLUDED here; they are + scheduled at `node_id`'s DISPATCH instead of its completion (see + `start_anchored_successors`). END is suppressed. + """ + return [t for t in self._static_succ.get(node_id, []) if t != END_NODE_ID] + + def incoming_static_edges(self, node_id: str) -> list[StaticEdge]: + """Return StaticEdge objects targeting `node_id`. + + Used by edge-gate computation for `min_start_delay_us` / + `delay_after_predecessor_us` / `delay_after_predecessor_start_us`. + """ + return list(self._static_pred_edges.get(node_id, [])) + + def start_anchored_successors(self, node_id: str) -> list[str]: + """Successors wired via start-anchored edges. The executor schedules + these at `node_id`'s DISPATCH (firing-gate clear), not its completion; + they are deliberately absent from `successors_after` so a child that + finishes before its still-running parent is not re-scheduled into the + cycle guard.""" + return [ + t for t in self._start_anchored_succ.get(node_id, []) if t != END_NODE_ID + ] + + +def _reject_mixed_anchor_fan_in( + static_pred_edges: dict[str, list[StaticEdge]], +) -> None: + """Reject a start-anchored in-edge that is not its target's ONLY in-edge. + + The runtime half-supports ANY fan-in involving a start anchor: the start + anchor schedules the target at its anchor parent's DISPATCH, so the target + fires without waiting for its other predecessors (a completion + predecessor's recorded ordering is silently ignored; a second start-anchor + parent that has not yet dispatched is silently dropped by the firing gate) + and, when the other predecessor later finishes/dispatches, it re-schedules + the DONE target into the cycle guard (spurious "cycle detected"). No + shipped lowering emits either shape (`apply_start_anchors` replaces a + start-anchored node's WHOLE in-edge set with exactly one edge), so it is + gated loudly at graph load instead. + """ + for target, edges in static_pred_edges.items(): + start_anchored = [ + e for e in edges if e.delay_after_predecessor_start_us is not None + ] + if not start_anchored or len(edges) == 1: + continue + completion = [e for e in edges if e.delay_after_predecessor_start_us is None] + if completion: + shape = "mixed-anchor fan-in" + detail = ( + f"start-anchored edge {start_anchored[0].source!r} -> {target!r} " + f"(delay_after_predecessor_start_us) and completion edge " + f"{completion[0].source!r} -> {target!r} arrive at the same node" + ) + else: + shape = "multi-start-anchored fan-in" + detail = ( + f"start-anchored edges {start_anchored[0].source!r} -> {target!r} " + f"and {start_anchored[1].source!r} -> {target!r} " + f"(delay_after_predecessor_start_us) arrive at the same node" + ) + raise NotImplementedError( + f"node {target!r}: {shape} is unsupported: {detail}. " + "The runtime schedules a start-anchored target at its anchor " + "parent's DISPATCH, so the node would fire without waiting for " + "its other predecessor and then be re-scheduled into the cycle " + "guard when that predecessor finishes or dispatches. A " + "start-anchored in-edge must be its target's ONLY in-edge." + ) + + +def collapse_leading_start_offsets(graph: GraphRecord) -> GraphRecord: + """Zero the leading phase-start offsets (`--burst-phase-starts` collapse). + + The anchor carrier for a firing's t*-relative leading offset is the + `min_start_delay_us` on its START in-edge: the trie build roots gap-started + chains at START with the warped arrival offset + (`interval_order.build_interval_edges`) and the t* snapshot chop re-roots + every surviving frontier node at START with its `arrival - t*` offset + (`snapshot_chop._chop_edges`). The executor gates on that edge field + (`_compute_firing_gate_us`). Burst collapses ONLY those leading offsets -- + every node fires the instant its inputs are ready -- while non-START edges + keep their recorded inter-turn pacing untouched. Node-level + `min_start_delay_us` is collapsed too, but ONLY on a node with no non-START + static in-edge (same leading-anchor semantics, hand-authored graphs): a + node-level delay on a node that HAS a real predecessor is mid-graph pacing, + not a leading offset -- it carries the AND-fan-in residual fold from + `snapshot_chop.chop_trie_at_frontier` and must survive burst, else the fold + is silently reverted on `--burst-phase-starts` runs. + + Pure function: returns a rebuilt `GraphRecord` (msgspec replace), never + mutates. Identity-preserving for untouched nodes/edges. + """ + has_real_pred = { + edge.target + for edge in graph.edges + if isinstance(edge, StaticEdge) and edge.source != START_NODE_ID + } + new_edges = [ + msgspec.structs.replace(edge, min_start_delay_us=0.0) + if isinstance(edge, StaticEdge) + and edge.source == START_NODE_ID + and edge.min_start_delay_us + else edge + for edge in graph.edges + ] + new_nodes = { + nid: ( + msgspec.structs.replace(node, min_start_delay_us=0.0) + if getattr(node, "min_start_delay_us", None) and nid not in has_real_pred + else node + ) + for nid, node in graph.nodes.items() + } + return msgspec.structs.replace(graph, nodes=new_nodes, edges=new_edges) diff --git a/src/aiperf/graph/worker_materialize.py b/src/aiperf/graph/worker_materialize.py new file mode 100644 index 0000000000..a2675876be --- /dev/null +++ b/src/aiperf/graph/worker_materialize.py @@ -0,0 +1,533 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Worker-side materialization of a graph-IR node's request. + +On a graph credit the worker rebuilds the node's request from its OWN per-node +envelope in the shared unified segment store (``GraphSegmentUnifiedClient``) -- +NO per-worker session, NO co-location, any worker serves any node. Each envelope +self-describes the node's full message list as interned int ``handles`` +(materialized via ``materialize_handles``); a slot-carrying node also carries an +``items`` assembly program that splices predecessor responses from the +worker-local dynamic pool (see ``_assemble_items``) to reconstruct the full +user/assistant alternation. The build interns the whole conversation prefix per +node, so there is no ancestor concatenation at dispatch. Warmup reuses the +profiling envelope and applies the warmup 1-token output cap at materialization +time. After the messages are built the node's own ``dispatch_overrides``, +run-level model fallback, warmup cap, and ``stream`` mode are layered on. No +tokenizer, no synthesis, no catalog at dispatch. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +import orjson + +from aiperf.common.enums import CacheBustTarget +from aiperf.common.environment import Environment +from aiperf.timing.strategies.cache_bust import ( + build_trace_instance_marker, + inject_marker_at_first_user_message, +) + +if TYPE_CHECKING: + from aiperf.common.models import EndpointInfo + from aiperf.dataset.graph_segment_unified_store import GraphSegmentUnifiedClient + +# The node's recorded output cap is carried under this dispatch-override key; the +# worker maps it to the wire token field per the endpoint convention (see +# ``endpoints/openai_chat.py::format_payload``). +_MAX_OUTPUT_TOKENS_KEY = "max_output_tokens" +_LEGACY_TOKEN_FIELD = "max_tokens" +_MODERN_TOKEN_FIELD = "max_completion_tokens" +_PROFILING_VARIANT = "profiling" +_WARMUP_VARIANT = "warmup" + +# Dynamo session-identity headers stamped verbatim at build time +# (``trie_lowering.py::_node_meta``). These carry the RECORDED session ids and +# must be uniquified per replay instance; ``x-dynamo-session-final`` is a +# per-turn flag, not an identity, and is forwarded untouched. +_DYNAMO_SESSION_ID_HEADERS = ("x-dynamo-session-id", "x-dynamo-parent-session-id") + + +def encode_overrides_inner(overrides: dict | None) -> bytes: + """Outer body fields (max_tokens/model/stream/extra/...) as the inner JSON to + splice after the messages array. orjson handles all value typing/escaping.""" + if not overrides: + return b"" + return orjson.dumps(overrides)[1:-1] # strip the { } + + +def read_node_envelope( + client: GraphSegmentUnifiedClient, + trace_id: str, + node_ordinal: int, + phase_variant: str = "profiling", +) -> dict[str, Any] | None: + """Fetch + decode one node's manifest envelope (variant-aliased), or None. + + The single pre-read seam: the worker calls this ONCE per graph credit -- + to route bytes-vs-dict and read the dynamic-content envelope flags + (``capture``, ``items``) -- and passes the decoded envelope into the + materialize functions so the manifest is never fetched or decoded twice. + """ + lookup_variant = ( + _PROFILING_VARIANT if phase_variant == _WARMUP_VARIANT else phase_variant + ) + raw = client.get_node_envelope(trace_id, node_ordinal, lookup_variant) + if raw is None: + return None + return orjson.loads(raw) + + +def _assemble_items( + client: GraphSegmentUnifiedClient, + items: list[dict[str, Any]], + slot_resolver: Callable[[int], Any], +) -> list[dict[str, Any]]: + """Assemble a slot-carrying node's messages from its ``items`` program. + + ``{"h": handle}`` appends the interned static message; ``{"s": {"src"}}`` + appends the producer's pooled reply as an assistant message — the verbatim + recorded assistant message (``orjson.loads(message_json)``, ``tool_calls`` + preserved) when the capture is structured, else ``{"role": "assistant", + "content": text}`` — or nothing on FAILED/EMPTY (deliberate omission); + ``{"m": {"role", "parts"}}`` emits one composed message whose content + concatenates static text parts and pooled slot texts (FAILED/EMPTY + substitute the empty string, so the role and static instruction survive a + failed producer). A MISSING pool value is :class:`GraphPoolMissingError` — + broken stickiness or backstop eviction is a loud trace error, never a + default. + """ + from aiperf.graph.dynamic_pool import ( + GraphCapturedReply, + GraphPoolMissingError, + GraphPoolSentinel, + ) + + def _slot_value(src: int) -> GraphCapturedReply | None: + value = slot_resolver(src) + if value is None: + raise GraphPoolMissingError(src) + if isinstance(value, GraphPoolSentinel): + return None + return value + + messages: list[dict[str, Any]] = [] + for token in items: + if "h" in token: + messages.extend(client.materialize_handles([token["h"]])) + elif "s" in token: + reply = _slot_value(token["s"]["src"]) + if reply is None: + pass # FAILED/EMPTY producer: deliberate omission + elif reply.message_json: + messages.append(orjson.loads(reply.message_json)) + else: + messages.append({"role": "assistant", "content": reply.text}) + elif "m" in token: + parts: list[str] = [] + for part in token["m"]["parts"]: + if "t" in part: + parts.append(part["t"]) + else: + reply = _slot_value(part["sv"]) + parts.append(reply.text if reply is not None else "") + messages.append({"role": token["m"]["role"], "content": "".join(parts)}) + else: + raise ValueError(f"unknown assembly item {token!r}") + if Environment.GRAPH.MERGE_CONSECUTIVE_USER: + messages = _merge_consecutive_user(messages) + return messages + + +def _merge_consecutive_user( + messages: list[dict[str, Any]], +) -> list[dict[str, Any]]: + """Collapse consecutive user-role messages into one (contents newline-joined). + + Consecutive user turns arise in a reconstructed conversation two ways: + omitting a FAILED/EMPTY producer's assistant turn leaves its user turn + adjacent to the next one, and a user-role init/delta boundary can place two + authored user turns back to back. Some chat APIs reject consecutive + same-role messages. Gated by ``AIPERF_GRAPH_MERGE_CONSECUTIVE_USER``. Only + user turns merge (adjacent assistants are intended); non-string content is + left untouched. + """ + merged: list[dict[str, Any]] = [] + for message in messages: + if ( + merged + and message.get("role") == "user" + and merged[-1].get("role") == "user" + and isinstance(message.get("content"), str) + and isinstance(merged[-1].get("content"), str) + ): + merged[-1] = { + "role": "user", + "content": f"{merged[-1]['content']}\n{message['content']}", + } + else: + merged.append(message) + return merged + + +def materialize_graph_request_unified( + client: GraphSegmentUnifiedClient, + trace_id: str, + node_ordinal: int, + phase_variant: str = "profiling", + *, + use_legacy_max_tokens: bool = False, + envelope: dict[str, Any] | None = None, + slot_resolver: Callable[[int], Any] | None = None, + default_model: str | None = None, +) -> dict[str, Any] | None: + """Rebuild an INTERNED (A2) node's request payload from int-handle manifests. + + Reads the node's envelope from the unified store and ONLY proceeds when the + envelope carries ``handles`` (the int-handle path), materializing ``messages`` + via :meth:`GraphSegmentUnifiedClient.materialize_handles` (or via + :func:`_assemble_items` when the envelope carries a slot-splice ``items`` + program). ``dispatch_overrides`` / warmup / stream handling layer on after, + stamping ``stream`` only when the envelope recorded one. + + Returns: + ``{"messages": [...], **dispatch_overrides}`` -- plus ``"stream": bool`` + only when the envelope recorded one -- the materialized request payload + pieces -- or ``None`` when the node's envelope is absent OR carries no + ``handles`` (every interned-store envelope carries ``handles``, so a + ``None`` is a genuine miss). + """ + if envelope is None: + envelope = read_node_envelope(client, trace_id, node_ordinal, phase_variant) + if envelope is None: + return None + handles = envelope.get("handles") + if handles is None: + return None + + items = envelope.get("items") + if items is not None and slot_resolver is not None: + messages = _assemble_items(client, items, slot_resolver) + else: + messages = client.materialize_handles(handles) + request: dict[str, Any] = {"messages": messages} + _apply_dispatch_overrides( + request, + envelope.get("dispatch_overrides") or {}, + use_legacy_max_tokens=use_legacy_max_tokens, + ) + # Run-level model fallback: a node with no per-node ``dispatch_overrides`` + # model (native graphs -- weka/dynamo stamp their recorded model) needs the + # run ``--model`` in the wire body, mirroring the linear path's + # ``turn.model or primary_model_name``. The endpoint's ``format_payload`` is + # bypassed for graph credits, so nothing else adds it. + if default_model is not None: + request.setdefault("model", default_model) + if phase_variant == _WARMUP_VARIANT: + request.pop(_MODERN_TOKEN_FIELD, None) + request[_LEGACY_TOKEN_FIELD] = Environment.GRAPH.WARMUP_MAX_OUTPUT_TOKENS + env_stream = envelope.get("stream") + if env_stream is not None: + request["stream"] = bool(env_stream) + return request + + +def materialize_graph_request_unified_bytes( + client: GraphSegmentUnifiedClient, + trace_id: str, + node_ordinal: int, + phase_variant: str = "profiling", + *, + use_legacy_max_tokens: bool = False, + endpoint: EndpointInfo, + envelope: dict[str, Any] | None = None, + default_model: str | None = None, +) -> tuple[bytes, str | None, bool] | None: + """Build a pre-serialized body from INTERNED (A2) int-handle manifests. + + The contiguous-bytes sibling of :func:`materialize_graph_request_unified` + (no per-segment ``orjson.loads``, no ``orjson.dumps`` of the messages + array): reads the SAME envelope, ONLY proceeds when it carries ``handles``, + and folds every outer field into the overrides tail at build time (the node's + ``dispatch_overrides`` with the mapped token cap, the warmup cap, and the + run-level options :func:`apply_run_level_payload_options` adds), then builds + the body via :meth:`GraphSegmentUnifiedClient.build_request_body_handles`. + + Because the body is built once and cannot be mutated after, the cache-bust + path is NOT handled here: it mutates message content (which bytes cannot + do), so the caller takes this path only when + ``endpoint.cache_bust == CacheBustTarget.NONE``. The overrides tail is built + with the SAME helpers + order as the dict path; parity is parsed-JSON + equality (``orjson.loads(body) == dict_path_payload``), not raw-byte + key-order parity. + + Returns: + ``(body, model, effective_stream)`` -- the contiguous + ``{"messages":[...],}`` body bytes; the node's per-node model + (``dispatch_overrides["model"]`` as folded, or ``None``) surfaced for + parity checks against the dict path's ``payload.get("model")`` -- the + worker deliberately does NOT stamp ``Turn.model`` with it (the recorded + model rides only the body bytes; ``record.model_name`` falls back to the + run ``--model`` for tokenizer selection); and the FINAL stamped wire + ``stream`` mode, so the caller can carry the recorded per-node override + onto ``RequestInfo`` for the transport (``effective_streaming``). Or + ``None`` when the node's envelope is absent OR carries no ``handles``. + """ + if envelope is None: + envelope = read_node_envelope(client, trace_id, node_ordinal, phase_variant) + if envelope is None: + return None + handles = envelope.get("handles") + if handles is None: + return None + + overrides: dict[str, Any] = {} + _apply_dispatch_overrides( + overrides, + envelope.get("dispatch_overrides") or {}, + use_legacy_max_tokens=use_legacy_max_tokens, + ) + # Run-level model fallback (see the dict path): fold the run ``--model`` into + # the body when the node carries none, so native graph requests are not + # rejected for a missing ``model`` field. + if default_model is not None: + overrides.setdefault("model", default_model) + if phase_variant == _WARMUP_VARIANT: + overrides.pop(_MODERN_TOKEN_FIELD, None) + overrides[_LEGACY_TOKEN_FIELD] = Environment.GRAPH.WARMUP_MAX_OUTPUT_TOKENS + env_stream = envelope.get("stream") + stream_override = bool(env_stream) if env_stream is not None else None + apply_run_level_payload_options( + overrides, + endpoint, + stream_override=stream_override, + skip_endpoint_extra=bool(envelope.get("endpoint_extra_applied")), + ) + + body = client.build_request_body_handles(handles, encode_overrides_inner(overrides)) + return body, overrides.get("model"), bool(overrides["stream"]) + + +def _apply_dispatch_overrides( + request: dict[str, Any], + dispatch_overrides: dict[str, Any], + *, + use_legacy_max_tokens: bool, +) -> None: + """Merge ``dispatch_overrides`` into ``request``, mapping the token cap. + + ``max_output_tokens`` becomes ``max_tokens`` (legacy) or + ``max_completion_tokens`` (modern); all other keys (``model`` and any + provider-specific tunables) pass through verbatim. + """ + for key, value in dispatch_overrides.items(): + if key == _MAX_OUTPUT_TOKENS_KEY: + token_field = ( + _LEGACY_TOKEN_FIELD if use_legacy_max_tokens else _MODERN_TOKEN_FIELD + ) + request[token_field] = value + else: + request[key] = value + + +def stamp_cache_bust_marker( + payload: dict[str, Any], + *, + benchmark_id: str, + trace_instance_id: str, + target: CacheBustTarget, +) -> None: + """Stamp the FIRST_TURN_PREFIX cache-bust marker on the wire payload. + + The marker is PER-TRACE-INSTANCE: every dispatch of one trace instance + (``credit.trace_id``, e.g. ``t-1#0`` -- including its nested/subagent + dispatches, which share the same ``trace_id``) carries the SAME marker, so + the instance's own conversation prefix stays consistent and prefix-caches + WITHIN the instance. Distinct instances get distinct markers (cross-instance + bust), and a recycled template (a fresh instance id like ``t-1#1``) mints a + fresh marker. Because the marker is deterministic from ``(benchmark_id, + trace_instance_id)``, the stateless worker needs no shared ledger to agree + across dispatches. + + The marker is prepended to the first ``role == "user"`` message of the + materialized ``payload["messages"]`` -- the FIRST user turn only, idempotent, + in the ``[rid:<12hex>]\\n\\n`` prefix format of agentx's cache-bust path + (``inject_marker_at_first_user_message`` mirrors agentx + ``worker.py::_inject_marker_into_first_user_turn``). + + No-op (payload byte-identical) when ``target`` is ``NONE`` -- the default -- + so the verbatim-replay path is unchanged unless the run passes ``--cache-bust``. + + Args: + payload: The materialized request payload (mutated in place); its + ``messages`` list is the wire prefix. + benchmark_id: Run-scoped salt so two runs mint different markers. + trace_instance_id: The trace INSTANCE id (``credit.trace_id``); shared by + every dispatch of the instance, distinct per instance, fresh on recycle. + target: Cache-bust mode; ``NONE`` stamps nothing. + """ + if target == CacheBustTarget.NONE: + return + marker = build_trace_instance_marker(benchmark_id, trace_instance_id, target=target) + messages = payload.get("messages") + if marker is None or not isinstance(messages, list): + return + inject_marker_at_first_user_message(messages, marker) + + +def uniquify_dynamo_session_headers( + extra_headers: dict[str, str] | None, + *, + trace_instance_id: str | None, + phase_variant: str, +) -> dict[str, str] | None: + """Uniquify recorded dynamo session ids per trace REPLAY INSTANCE. + + Build-time lowering stamps the RECORDED ``x-dynamo-session-id`` / + ``x-dynamo-parent-session-id`` verbatim into the node envelope, but the + graph replay strategy runs multiple concurrent instances of the SAME trace + (lanes/recycles wrap the corpus). Forwarded verbatim, every instance would + share one server-side session: affinity collapses to one bucket and the + first instance's ``x-dynamo-session-final`` evicts KV under its + still-running siblings. Mirror of :func:`stamp_cache_bust_marker`'s + placement -- the worker is the only party that knows the instance id. + + Both identity headers get the SAME deterministic ``#.`` + suffix derived from ``trace_instance_id`` (``t-1#0.0`` -> ``0.0``) and + ``phase_variant``, so parent-child linkage WITHIN an instance is preserved + while distinct instances -- and a warmup instance vs the profiling + instance of the same ``(lane, pass)`` slot -- never collide. + ``x-dynamo-session-final`` is forwarded untouched: each instance dispatches + its own copy of the session's last turn, closing only its own session. + + No-op (input returned as-is) when there are no headers, no dynamo identity + header, or no ``#`` instance suffix on ``trace_instance_id`` -- plain + (non-instanced) replay is unaffected. + + Args: + extra_headers: The node envelope's per-node HTTP headers, or ``None``. + trace_instance_id: The credit's trace INSTANCE id + (``{template}#{lane}.{pass}``); ``None`` / suffix-less ids disable + the transform. + phase_variant: The credit's graph phase variant (``"profiling"`` / + ``"warmup"``); folded into the suffix so warmup and profiling + instances of one slot open distinct sessions. + + Returns: + A NEW dict with the identity headers suffixed, or the input unchanged. + """ + if not extra_headers or not trace_instance_id: + return extra_headers + _, sep, instance_nonce = trace_instance_id.partition("::") + if not sep: + return extra_headers + if not any(header in extra_headers for header in _DYNAMO_SESSION_ID_HEADERS): + return extra_headers + tag = f"::{phase_variant}-{instance_nonce}" + transformed = dict(extra_headers) + for header in _DYNAMO_SESSION_ID_HEADERS: + value = transformed.get(header) + if value is not None: + transformed[header] = f"{value}{tag}" + return transformed + + +def strip_dynamo_session_headers( + extra_headers: dict[str, str] | None, +) -> dict[str, str] | None: + """Drop RECORDED dynamo session-identity headers from a node envelope. + + When an active ``--session-routing`` plugin owns session identity, the + build-time recorded ``x-dynamo-session-id`` / ``-parent-session-id`` / + ``-session-final`` headers are stale replay artifacts: forwarding them + alongside the plugin's live headers would put two (case-distinct, + conflicting) identities on the wire. Non-identity recorded headers are + preserved. Returns the input unchanged when nothing needs stripping. + """ + if not extra_headers: + return extra_headers + identity = (*_DYNAMO_SESSION_ID_HEADERS, "x-dynamo-session-final") + if not any(header in extra_headers for header in identity): + return extra_headers + stripped = {k: v for k, v in extra_headers.items() if k not in identity} + return stripped or None + + +def apply_run_level_payload_options( + payload: dict[str, Any], + endpoint: EndpointInfo, + stream_override: bool | None = None, + *, + skip_endpoint_extra: bool = False, +) -> None: + """Layer run-level endpoint concerns onto a materialized graph payload. + + The materialized payload carries the node's own per-node + ``dispatch_overrides`` (``model``, the mapped token cap). This layers the + run-level concerns the verbatim ``raw_payload`` path would otherwise drop + (the chat endpoint's ``format_payload`` is bypassed for graph credits): + + - ``stream`` is stamped from the RECORDED per-node ``stream_override`` when + the caller supplies one (weka ``"n"``/``"s"``, dynamo ``ttft_ms``); a + ``None`` override falls back to the GLOBAL ``endpoint.streaming`` setting. + The recorded per-node mode wins for graph credits so a recorded ``n``-type + turn stays non-streaming inside an otherwise-streaming run; the transport + picks the matching wire mode per-request (``effective_streaming``). + - ``endpoint.extra`` (the user's ``--extra-inputs`` vendor tunables) is + merged with the USER winning over any per-node key. + SKIPPED when ``skip_endpoint_extra`` is set: an adapter that already folded + the run's ``--extra-inputs`` into the node's ``dispatch_overrides`` at parse + owns those keys, so re-merging here would clobber the adapter-owned values. + - ``stream_options.include_usage = True`` is forced when the FINAL stamped + ``stream`` is on AND ``endpoint.use_server_token_count`` is set, so + server-side token-count metrics still get usage on graph credits (it keys + on ``payload.get("stream")`` AFTER the stamp above). Any author-supplied + ``stream_options`` keys are preserved. + + Args: + payload: The materialized request payload (mutated in place). + endpoint: The run's endpoint info carrying the run-level options. + stream_override: Recorded per-node wire mode for graph credits; ``None`` + follows the global ``endpoint.streaming``. + skip_endpoint_extra: When set, the ``endpoint.extra`` merge is skipped + (adapter-owned extras precedence); the ``stream`` stamp and + ``include_usage`` forcing are unaffected. + """ + payload["stream"] = ( + stream_override if stream_override is not None else bool(endpoint.streaming) + ) + + if not skip_endpoint_extra: + for key, value in endpoint.extra or []: + payload[key] = value + + if payload.get("stream") and endpoint.use_server_token_count: + _ensure_include_usage(payload) + + +def _ensure_include_usage(payload: dict[str, Any]) -> None: + """Force ``stream_options.include_usage = True``, preserving author keys. + + Mirrors ``ChatEndpoint._ensure_include_usage`` so graph credits get the + same usage opt-in the linear path applies in ``format_payload``. + """ + stream_options = payload.get("stream_options") + if not isinstance(stream_options, dict): + payload["stream_options"] = {"include_usage": True} + return + if "include_usage" not in stream_options: + stream_options["include_usage"] = True + + +__all__ = [ + "encode_overrides_inner", + "materialize_graph_request_unified", + "materialize_graph_request_unified_bytes", + "apply_run_level_payload_options", + "stamp_cache_bust_marker", + "strip_dynamo_session_headers", + "uniquify_dynamo_session_headers", +] diff --git a/src/aiperf/metrics/types/context_overflow_count_metric.py b/src/aiperf/metrics/types/context_overflow_count_metric.py new file mode 100644 index 0000000000..7d511dcaca --- /dev/null +++ b/src/aiperf/metrics/types/context_overflow_count_metric.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Aggregate counter for runtime context-overflow detections. + +Companion to ``aiperf.common.scenario.is_context_overflow_response`` and the +``RequestRecord.context_overflow`` flag set by the inference result parser. +Increments by 1 per record whose request was tagged as a context-overflow +error; otherwise contributes 0. Used by the InferenceX AgentX scenario +(RFC §7) to flip ``submission_valid=false`` when the overflow rate exceeds the +configured threshold. +""" + +from aiperf.common.enums import GenericMetricUnit, MetricFlags +from aiperf.common.models import ParsedResponseRecord +from aiperf.metrics.base_aggregate_counter_metric import BaseAggregateCounterMetric +from aiperf.metrics.metric_dicts import MetricRecordDict + + +class ContextOverflowCountMetric(BaseAggregateCounterMetric[int]): + """Counts records flagged as context-overflow by the runtime classifier. + + Formula: + ``` + Context Overflow Count = Sum(1 if request.context_overflow else 0) + ``` + + Marked ``ERROR_ONLY`` because context-overflow records are by definition + error responses, and ``NO_INDIVIDUAL_RECORDS`` because the count is an + aggregate-only signal that doesn't make sense on a per-record export. + """ + + tag = "context_overflow_count" + header = "Context Overflow Count" + short_header = "Ctx Overflow" + short_header_hide_unit = True + unit = GenericMetricUnit.REQUESTS + flags = MetricFlags.ERROR_ONLY | MetricFlags.NO_INDIVIDUAL_RECORDS + required_metrics = None + + def _parse_record( + self, record: ParsedResponseRecord, record_metrics: MetricRecordDict + ) -> int: + """Return 1 iff the underlying RequestRecord was flagged as overflow.""" + return 1 if getattr(record.request, "context_overflow", False) else 0 diff --git a/src/aiperf/metrics/types/inter_chunk_latency_metric.py b/src/aiperf/metrics/types/inter_chunk_latency_metric.py index 04e9834682..f78441e5de 100644 --- a/src/aiperf/metrics/types/inter_chunk_latency_metric.py +++ b/src/aiperf/metrics/types/inter_chunk_latency_metric.py @@ -1,5 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from aiperf.common.constants import STREAMED_REQUEST_TAG from aiperf.common.enums import MetricConsoleGroup, MetricFlags, MetricTimeUnit from aiperf.common.exceptions import NoMetricValue from aiperf.common.models import ParsedResponseRecord @@ -35,7 +36,7 @@ class InterChunkLatencyMetric(BaseRecordMetric[list[int]]): display_unit = MetricTimeUnit.MILLISECONDS flags = MetricFlags.STREAMING_TOKENS_ONLY console_group = MetricConsoleGroup.NONE - required_metrics = None + required_metrics = {STREAMED_REQUEST_TAG} def _parse_record( self, @@ -50,9 +51,12 @@ def _parse_record( Note: Only includes content responses (with actual data), excluding usage-only or [DONE] chunks. Raises: - NoMetricValue: If the record does not have at least two content responses - ValueError: If any of the inter chunk latencies are not positive. + NoMetricValue: If the record did not stream, or does not have at least two content responses + ValueError: If any of the inter chunk latencies are negative. """ + if STREAMED_REQUEST_TAG not in record_metrics: + raise NoMetricValue("record did not stream; streaming metrics skipped") + content_responses = record.content_responses if len(content_responses) < 2: diff --git a/src/aiperf/metrics/types/stream_latency_metrics.py b/src/aiperf/metrics/types/stream_latency_metrics.py index 8707146f85..5fff6074be 100644 --- a/src/aiperf/metrics/types/stream_latency_metrics.py +++ b/src/aiperf/metrics/types/stream_latency_metrics.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from aiperf.common.constants import STREAMED_REQUEST_TAG from aiperf.common.enums import MetricFlags, MetricTimeUnit from aiperf.common.exceptions import NoMetricValue from aiperf.common.models import ParsedResponseRecord @@ -29,7 +30,7 @@ class StreamSetupLatencyMetric(BaseRecordMetric[int]): unit = MetricTimeUnit.NANOSECONDS display_unit = MetricTimeUnit.MILLISECONDS flags = MetricFlags.STREAMING_ONLY | MetricFlags.EXPERIMENTAL - required_metrics = None + required_metrics = {STREAMED_REQUEST_TAG} def _parse_record( self, @@ -37,6 +38,8 @@ def _parse_record( record_metrics: MetricRecordDict, ) -> int: """This method extracts the request and receive start timestamps, and calculates the stream setup time.""" + if STREAMED_REQUEST_TAG not in record_metrics: + raise NoMetricValue("record did not stream; streaming metrics skipped") if not record.request.recv_start_perf_ns or not record.start_perf_ns: raise NoMetricValue( diff --git a/src/aiperf/metrics/types/streamed_request_count_metric.py b/src/aiperf/metrics/types/streamed_request_count_metric.py new file mode 100644 index 0000000000..74ad0436a3 --- /dev/null +++ b/src/aiperf/metrics/types/streamed_request_count_metric.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Visible aggregate counter of requests that streamed on the wire. + +This is the user-facing streaming denominator, displayed beside Request Count. It +mirrors ``RequestCountMetric`` but counts only records whose underlying request +streamed (``RequestRecord.streamed``); non-streamed records raise +``NoMetricValue`` and contribute nothing. The hidden per-record gate that streaming +metrics depend on is the companion ``StreamedRequestMetric`` predicate. +""" + +from aiperf.common.constants import STREAMED_REQUEST_COUNT_TAG +from aiperf.common.enums import GenericMetricUnit, MetricFlags +from aiperf.common.exceptions import NoMetricValue +from aiperf.common.models import ParsedResponseRecord +from aiperf.metrics.base_aggregate_counter_metric import BaseAggregateCounterMetric +from aiperf.metrics.metric_dicts import MetricRecordDict + + +class StreamedRequestCountMetric(BaseAggregateCounterMetric[int]): + """Counts requests that actually streamed on the wire. + + Formula: + ``` + Streamed Request Count = Sum(1 if request.streamed else skip) + ``` + """ + + tag = STREAMED_REQUEST_COUNT_TAG + header = "Streamed Request Count" + short_header = "Streamed Requests" + short_header_hide_unit = True + unit = GenericMetricUnit.REQUESTS + display_order = 1101 + flags = MetricFlags.STREAMING_TOKENS_ONLY | MetricFlags.NO_INDIVIDUAL_RECORDS + required_metrics = None + + def _parse_record( + self, + record: ParsedResponseRecord, + record_metrics: MetricRecordDict, + ) -> int: + """Return 1 iff the underlying request streamed; else raise to skip it. + + Raises: + NoMetricValue: If the request did not stream on the wire. + """ + if not record.request.streamed: + raise NoMetricValue("request did not stream") + return 1 diff --git a/src/aiperf/metrics/types/streamed_request_metric.py b/src/aiperf/metrics/types/streamed_request_metric.py new file mode 100644 index 0000000000..a5a86bb0b1 --- /dev/null +++ b/src/aiperf/metrics/types/streamed_request_metric.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-record streaming predicate that gates every per-record streaming metric. + +This metric's PRESENCE in a record's ``MetricRecordDict`` is the predicate: it +parses to 1 for a request that actually streamed on the wire +(``RequestRecord.streamed``), and raises ``NoMetricValue`` otherwise so it is +absent from that record's dict. + +Streaming metrics whose ``_parse_record`` reads response timing directly declare +this tag in ``required_metrics`` (so the topological order computes it first) and +check membership; when the record did not stream the tag is absent and the +streaming metric skips itself instead of, e.g., reporting full request latency as +TTFT for a non-streamed record whose single response timestamp is the completion. + +It is a ``BaseRecordMetric`` (not an aggregate counter) so RECORD-typed streaming +metrics may depend on it: the dependency validator only permits a RECORD metric to +depend on other RECORD metrics. It carries ``INTERNAL`` (excluded from the +``summarize()`` JSON/CSV export, where a constant-1 stat row is meaningless), +``NO_INDIVIDUAL_RECORDS`` (excluded from per-record exports), and +``console_group=NONE`` (hidden from the console); INTERNAL metrics stay computable +as dependencies since ``get_filters`` never disallows INTERNAL. The user-facing +streamed-request count is the companion ``StreamedRequestCountMetric`` aggregate. +""" + +from aiperf.common.constants import STREAMED_REQUEST_TAG +from aiperf.common.enums import GenericMetricUnit, MetricConsoleGroup, MetricFlags +from aiperf.common.exceptions import NoMetricValue +from aiperf.common.models import ParsedResponseRecord +from aiperf.metrics import BaseRecordMetric +from aiperf.metrics.metric_dicts import MetricRecordDict + + +class StreamedRequestMetric(BaseRecordMetric[int]): + """Per-record streaming predicate that gates streaming metrics. + + Formula: + ``` + Streamed Request = 1 if request.streamed else skip + ``` + """ + + tag = STREAMED_REQUEST_TAG + header = "Streamed Request" + short_header = "Streamed Request" + short_header_hide_unit = True + unit = GenericMetricUnit.REQUESTS + display_order = 1102 + flags = ( + MetricFlags.STREAMING_TOKENS_ONLY + | MetricFlags.NO_INDIVIDUAL_RECORDS + | MetricFlags.INTERNAL + ) + console_group = MetricConsoleGroup.NONE + required_metrics = None + + def _parse_record( + self, + record: ParsedResponseRecord, + record_metrics: MetricRecordDict, + ) -> int: + """Return 1 iff the underlying request streamed; else raise to stay absent. + + Raises: + NoMetricValue: If the request did not stream on the wire. + """ + if not record.request.streamed: + raise NoMetricValue("request did not stream") + return 1 diff --git a/src/aiperf/metrics/types/theoretical_prefix_cache_metric.py b/src/aiperf/metrics/types/theoretical_prefix_cache_metric.py new file mode 100644 index 0000000000..dd3d739018 --- /dev/null +++ b/src/aiperf/metrics/types/theoretical_prefix_cache_metric.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import NoReturn + +from aiperf.common.enums import ( + GenericMetricUnit, + MetricConsoleGroup, + MetricFlags, +) +from aiperf.common.exceptions import NoMetricValue +from aiperf.metrics import BaseDerivedMetric +from aiperf.metrics.metric_dicts import MetricResultsDict + + +class TheoreticalPrefixCacheHitMetric(BaseDerivedMetric[float]): + """Cumulative infinite-cache prefix hit rate over a trace replay, in percent. + + Invariant: externally injected by + `TheoreticalPrefixCacheAccumulator.summarize` from loader-stamped per-node / + per-turn prefix-cache block counts. `_derive_value` is intentionally + non-functional; `MetricsAccumulator._resolve_derived_metrics` is expected + to catch NoMetricValue and skip the tag during its derivation walk. The + class exists so display consumers (realtime dashboard table, console + exporter) can resolve the tag's display metadata from the MetricRegistry. + """ + + tag = "theoretical_prefix_cache_hit" + header = "Theoretical Prefix Cache Hit" + unit = GenericMetricUnit.PERCENT + display_order = 1018 + flags = MetricFlags.NONE + console_group = MetricConsoleGroup.CACHE + + def _derive_value(self, metric_results: MetricResultsDict) -> NoReturn: + raise NoMetricValue( + "Cannot derive 'theoretical_prefix_cache_hit' from MetricResultsDict: " + "this metric is externally injected by " + "TheoreticalPrefixCacheAccumulator.summarize. If this exception " + "surfaces, the derivation walk is missing its NoMetricValue handler " + "(see MetricsAccumulator._resolve_derived_metrics)." + ) diff --git a/src/aiperf/metrics/types/time_to_first_output_token_metric.py b/src/aiperf/metrics/types/time_to_first_output_token_metric.py index b59e9ac356..593450ff9f 100644 --- a/src/aiperf/metrics/types/time_to_first_output_token_metric.py +++ b/src/aiperf/metrics/types/time_to_first_output_token_metric.py @@ -1,5 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from aiperf.common.constants import STREAMED_REQUEST_TAG from aiperf.common.enums import MetricFlags, MetricTimeUnit from aiperf.common.exceptions import NoMetricValue from aiperf.common.models import ParsedResponseRecord @@ -26,7 +27,7 @@ class TimeToFirstOutputTokenMetric(BaseRecordMetric[int]): reasoning tokens. For models without reasoning, TTFO and TTFT are equivalent. - Non-reasoning tokens: Includes TextResponseData with non-empty text, ReasoningResponseData with non-empty content field (regardless of reasoning field), - or ToolCallResponseData with non-empty text. + or ToolCallResponseData with non-empty tool_call_text. Formula: Time to First Output = First Non-Reasoning Token Timestamp - Request Start Timestamp @@ -39,7 +40,7 @@ class TimeToFirstOutputTokenMetric(BaseRecordMetric[int]): display_unit = MetricTimeUnit.MILLISECONDS display_order = 210 flags = MetricFlags.STREAMING_TOKENS_ONLY | MetricFlags.SUPPORTS_REASONING - required_metrics = None + required_metrics = {STREAMED_REQUEST_TAG} def _parse_record( self, @@ -55,11 +56,15 @@ def _parse_record( Args: record: The parsed response record containing request and response timing data - record_metrics: Dictionary of previously computed metrics for this record (unused) + record_metrics: Dictionary of previously computed metrics for this record; must + contain the streamed-request predicate or the record is skipped Returns: The time to first output token in nanoseconds """ + if STREAMED_REQUEST_TAG not in record_metrics: + raise NoMetricValue("record did not stream; streaming metrics skipped") + try: # Try and find the first non-reasoning token output and extract the timestamp. # This is done by checking for the first response that is either a TextResponseData or a ReasoningResponseData diff --git a/src/aiperf/metrics/types/ttft_metric.py b/src/aiperf/metrics/types/ttft_metric.py index 6283af379f..dfdcdbe033 100644 --- a/src/aiperf/metrics/types/ttft_metric.py +++ b/src/aiperf/metrics/types/ttft_metric.py @@ -1,5 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from aiperf.common.constants import STREAMED_REQUEST_TAG from aiperf.common.enums import MetricFlags, MetricTimeUnit from aiperf.common.exceptions import NoMetricValue from aiperf.common.models import ParsedResponseRecord @@ -25,7 +26,7 @@ class TTFTMetric(BaseRecordMetric[int]): MetricFlags.STREAMING_TOKENS_ONLY | MetricFlags.PERCENTILE_INCLUDES_FAILED_REQUESTS ) - required_metrics = None + required_metrics = {STREAMED_REQUEST_TAG} def _parse_record( self, @@ -37,9 +38,11 @@ def _parse_record( RequestRecord object, computes the difference (TTFT), and returns the result. Raises: - NoMetricValue: If the record does not have at least one content response + NoMetricValue: If the record did not stream, or does not have at least one content response ValueError: If the first response is before the request start timestamp. """ + if STREAMED_REQUEST_TAG not in record_metrics: + raise NoMetricValue("record did not stream; streaming metrics skipped") if len(record.content_responses) < 1: raise NoMetricValue( diff --git a/src/aiperf/metrics/types/ttst_metric.py b/src/aiperf/metrics/types/ttst_metric.py index 5988967f4f..d9af486e71 100644 --- a/src/aiperf/metrics/types/ttst_metric.py +++ b/src/aiperf/metrics/types/ttst_metric.py @@ -1,5 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from aiperf.common.constants import STREAMED_REQUEST_TAG from aiperf.common.enums import MetricFlags, MetricTimeUnit from aiperf.common.exceptions import NoMetricValue from aiperf.common.models import ParsedResponseRecord @@ -22,7 +23,7 @@ class TTSTMetric(BaseRecordMetric[int]): display_unit = MetricTimeUnit.MILLISECONDS display_order = 200 flags = MetricFlags.STREAMING_TOKENS_ONLY - required_metrics = None + required_metrics = {STREAMED_REQUEST_TAG} def _parse_record( self, @@ -34,9 +35,11 @@ def _parse_record( RequestRecord object, computes the difference (TTST), and returns the result. Raises: - NoMetricValue: If the record does not have at least two content responses + NoMetricValue: If the record did not stream, or does not have at least two content responses ValueError: If the second response is before the first response. """ + if STREAMED_REQUEST_TAG not in record_metrics: + raise NoMetricValue("record did not stream; streaming metrics skipped") if len(record.content_responses) < 2: raise NoMetricValue( diff --git a/src/aiperf/network_latency/probe.py b/src/aiperf/network_latency/probe.py index 40d19076cd..3e0313dd9a 100644 --- a/src/aiperf/network_latency/probe.py +++ b/src/aiperf/network_latency/probe.py @@ -107,7 +107,9 @@ async def resolve(self) -> None: self._resolved_family = family self._resolved_host = sockaddr[0] self.debug( - lambda: f"Resolved {self._target_host}:{self._target_port} -> {self._resolved_host}" + lambda: ( + f"Resolved {self._target_host}:{self._target_port} -> {self._resolved_host}" + ) ) except OSError as e: self.warning( @@ -150,8 +152,10 @@ async def probe_once(self) -> None: await writer.wait_closed() except OSError as close_error: self.debug( - lambda err=close_error: f"Error closing probe connection to " - f"{self._target_host}:{self._target_port}: {err!r}" + lambda err=close_error: ( + f"Error closing probe connection to " + f"{self._target_host}:{self._target_port}: {err!r}" + ) ) except OSError as e: error = ErrorDetails.from_exception(e) diff --git a/src/aiperf/orchestrator/local_executor.py b/src/aiperf/orchestrator/local_executor.py index a7b9f4206a..96294b1c77 100644 --- a/src/aiperf/orchestrator/local_executor.py +++ b/src/aiperf/orchestrator/local_executor.py @@ -30,6 +30,58 @@ __all__ = ["LocalSubprocessExecutor"] +def _resolve_scenario_in_parent(run: BenchmarkRun) -> None: + """Apply the scenario lock ONCE in the parent, before the subprocess dump. + + The multi-run path serializes the run (``model_dump(exclude_none=True)``) + and re-loads it (``BenchmarkRun.model_validate``) inside the subprocess, + THEN runs the resolver chain there. ``model_validate`` re-marks every + non-None field as "set", so in the subprocess ``endpoint.model_fields_set`` + spuriously contains ``streaming`` / ``cache_bust`` even when the user never + set them -- which would make the in-subprocess scenario validator read the + auto-fillable defaults as explicit conflicts and raise a SPURIOUS + ``ScenarioLockError``. + + Resolving here, on the LIVE config (faithful ``model_fields_set``), bakes the + auto-fills (streaming=True, cache_bust=first_turn_prefix, duration) into + ``run.cfg`` before the dump. The in-subprocess re-resolution then becomes a + no-op: the baked values already satisfy the locks, so the stale + ``model_fields_set`` is never consulted. No-op when ``run.cfg.scenario`` is + None. A genuine user conflict still raises here (fail fast, pre-subprocess). + """ + if getattr(run.cfg, "scenario", None) is None: + return + from aiperf.common.scenario import apply_scenario + + apply_scenario(run) + + +def _carried_scenario_outcome(run: BenchmarkRun) -> tuple[bool | None, list[str]]: + """Extract the lock-only scenario verdict to carry onto ``RunResult``. + + Reads ``run.resolved.scenario_outcome`` -- the ``ScenarioOutcome`` stamped by + ``apply_scenario`` in ``_resolve_scenario_in_parent`` (parent process, before + the subprocess dump). This is the VALIDATOR-side outcome only; the aggregate + exporter folds cross-run runtime signals (context-overflow rate, cancellation) + on top of it. Returns ``(None, [])`` when no scenario was set (the outcome's + ``submission_valid`` is None then), so the aggregate reader omits the verdict. + + Args: + run: The resolved ``BenchmarkRun`` whose ``resolved.scenario_outcome`` + carries the lock verdict. + + Returns: + A ``(submission_valid, submission_invalid_reasons)`` tuple mirroring the + validator outcome; ``submission_valid`` is None for a non-scenario run. + """ + outcome = run.resolved.scenario_outcome + if outcome is None: + return None, [] + submission_valid = getattr(outcome, "submission_valid", None) + reasons = list(getattr(outcome, "submission_invalid_reasons", []) or []) + return submission_valid, reasons + + class LocalSubprocessExecutor(RunExecutor): """Run benchmarks via subprocess of aiperf.orchestrator.subprocess_runner.""" @@ -46,16 +98,31 @@ async def execute(self, run: BenchmarkRun) -> RunResult: def _execute_sync(self, run: BenchmarkRun) -> RunResult: artifacts_path = run.artifact_dir artifacts_path.mkdir(parents=True, exist_ok=True) + # Scenario locks mutate only THIS run's config (apply-or-lock on + # run.cfg); there are no process-global writes to scope, so nothing + # can leak into later runs' subprocess environments. + _resolve_scenario_in_parent(run) + submission_valid, submission_invalid_reasons = _carried_scenario_outcome(run) try: config_file = self._prepare_run_artifacts(run, artifacts_path) result = self._run_benchmark_subprocess(config_file, run) if result.returncode != 0: - return self._failure_from_subprocess(result, run.label, artifacts_path) + return self._failure_from_subprocess( + result, + run.label, + artifacts_path, + submission_valid=submission_valid, + submission_invalid_reasons=submission_invalid_reasons, + ) summary_metrics = self._extract_summary_metrics(run) return self._build_result_from_metrics( - summary_metrics, run.label, artifacts_path + summary_metrics, + run.label, + artifacts_path, + submission_valid=submission_valid, + submission_invalid_reasons=submission_invalid_reasons, ) except Exception as e: logger.exception(f"Error executing run {run.label}") @@ -64,6 +131,8 @@ def _execute_sync(self, run: BenchmarkRun) -> RunResult: success=False, error=str(e), artifacts_path=artifacts_path, + submission_valid=submission_valid, + submission_invalid_reasons=submission_invalid_reasons, ) @staticmethod @@ -161,6 +230,9 @@ def _failure_from_subprocess( result: subprocess.CompletedProcess[str], label: str, artifacts_path: Path, + *, + submission_valid: bool | None, + submission_invalid_reasons: list[str], ) -> RunResult: """Build a failed RunResult from a non-zero subprocess exit.""" error_msg = f"Benchmark failed with exit code {result.returncode}" @@ -172,6 +244,8 @@ def _failure_from_subprocess( success=False, error=error_msg, artifacts_path=artifacts_path, + submission_valid=submission_valid, + submission_invalid_reasons=submission_invalid_reasons, ) @staticmethod @@ -179,6 +253,9 @@ def _build_result_from_metrics( summary_metrics: dict[str, JsonMetricResult], label: str, artifacts_path: Path, + *, + submission_valid: bool | None, + submission_invalid_reasons: list[str], ) -> RunResult: """Classify success/failure from extracted summary metrics.""" if not summary_metrics: @@ -191,6 +268,8 @@ def _build_result_from_metrics( success=False, error=error_msg, artifacts_path=artifacts_path, + submission_valid=submission_valid, + submission_invalid_reasons=submission_invalid_reasons, ) request_count_metric = summary_metrics.get("request_count") @@ -207,6 +286,8 @@ def _build_result_from_metrics( success=False, error=error_msg, artifacts_path=artifacts_path, + submission_valid=submission_valid, + submission_invalid_reasons=submission_invalid_reasons, ) return RunResult( @@ -214,6 +295,8 @@ def _build_result_from_metrics( success=True, summary_metrics=summary_metrics, artifacts_path=artifacts_path, + submission_valid=submission_valid, + submission_invalid_reasons=submission_invalid_reasons, ) def _extract_summary_metrics( diff --git a/src/aiperf/orchestrator/models.py b/src/aiperf/orchestrator/models.py index afeb5af5f8..9fea87e57e 100644 --- a/src/aiperf/orchestrator/models.py +++ b/src/aiperf/orchestrator/models.py @@ -43,6 +43,23 @@ class RunResult(AIPerfBaseModel): trial_index: int = Field( default=0, description="Zero-based trial index within the variation." ) + submission_valid: bool | None = Field( + default=None, + description="Lock-only scenario verdict carried from " + "run.resolved.scenario_outcome.submission_valid (stamped by " + "apply_scenario in the parent before the subprocess dump). True for a " + "clean lock, False under --unsafe-override with violations, None when no " + "--scenario was set. This is the VALIDATOR-side value ONLY; the aggregate " + "exporter folds cross-run runtime signals (context-overflow rate, " + "cancellation) on top of it. Distinct from the per-run JSON's already-" + "runtime-folded submission_valid, which must not be reused as the seed.", + ) + submission_invalid_reasons: list[str] = Field( + default_factory=list, + description="Lock-only scenario invalid-reason tags carried from " + "run.resolved.scenario_outcome.submission_invalid_reasons (e.g. " + "['unsafe_override']). Empty for a clean lock or a non-scenario run.", + ) VariationKey = tuple[str, tuple[tuple[str, Any], ...]] diff --git a/src/aiperf/plugin/categories.yaml b/src/aiperf/plugin/categories.yaml index 6310d4eb56..ff64f25cc1 100644 --- a/src/aiperf/plugin/categories.yaml +++ b/src/aiperf/plugin/categories.yaml @@ -117,6 +117,22 @@ public_dataset_loader: Supports HTTP-downloaded and HuggingFace-hosted datasets. One-to-one mapping: selected based on --public-dataset configuration. +graph_adapter: + protocol: aiperf.dataset.graph.adapters.protocols:GraphAdapterProtocol + metadata_class: aiperf.plugin.schema.schemas:GraphAdapterMetadata + enum: GraphAdapterType + description: | + Graph adapters convert third-party trace/log files into the canonical + ParsedGraph used by the graph benchmark mode. Each adapter exposes + can_load(path) for format detection and parse(path, ctx) for full + conversion, where ctx is an optional GraphParseContext of run-derived + knobs; adapters forward a ctx field only when it is set, so a ctx-less + parse keeps the protocol-default entry behavior. Selected by + `--graph-format` or auto-detected from file extension + content sniff. + Detection ordering (when multiple adapters' can_load() return True for + the same file) comes from each plugin's metadata `detection_priority` + field. + # ============================================================================= # Endpoint and Transport Categories @@ -140,6 +156,23 @@ transport: Manages connection pooling, streaming, error handling, and TCP configuration. One-to-one mapping based on transport_type configuration. +# ============================================================================= +# Session Routing Category +# ============================================================================= + +session_routing: + protocol: aiperf.workers.session_routing:SessionRoutingBase + enum: SessionRoutingType + description: | + Session-routing transforms stamp per-session identity onto outbound + requests (headers or body metadata) so an external router in front of + multiple replicas can pin every turn of a session to one worker. + Selected via --session-routing; parameterized via repeatable + --session-routing-opt key=value validated against the plugin's + Options model. mutates_body is protocol metadata marking a mode as + body-mutating / incompatible with any verbatim-bytes request path; it is + not currently enforced by any gate on this codebase. + # ============================================================================= # Processing Categories # ============================================================================= diff --git a/src/aiperf/plugin/enums.py b/src/aiperf/plugin/enums.py index e094ca3a52..8125237ea7 100644 --- a/src/aiperf/plugin/enums.py +++ b/src/aiperf/plugin/enums.py @@ -12,7 +12,7 @@ from aiperf.plugin import plugins from aiperf.plugin.extensible_enums import create_enum -__all__ = ["APIRouterType", "APIRouterTypeStr", "AccumulatorType", "AccumulatorTypeStr", "AccuracyBenchmarkType", "AccuracyBenchmarkTypeStr", "AccuracyGraderType", "AccuracyGraderTypeStr", "AnalyzerType", "AnalyzerTypeStr", "ArrivalPattern", "ArrivalPatternStr", "CommClientType", "CommClientTypeStr", "CommunicationBackend", "CommunicationBackendStr", "ComposerType", "ComposerTypeStr", "ConsoleExporterType", "ConsoleExporterTypeStr", "ConvergenceCriterionType", "ConvergenceCriterionTypeStr", "CustomDatasetType", "CustomDatasetTypeStr", "DataExporterType", "DataExporterTypeStr", "DatasetBackingStoreType", "DatasetBackingStoreTypeStr", "DatasetClientStoreType", "DatasetClientStoreTypeStr", "DatasetFormat", "DatasetFormatStr", "DatasetSamplingStrategy", "DatasetSamplingStrategyStr", "EndpointType", "EndpointTypeStr", "GPUTelemetryCollectorType", "GPUTelemetryCollectorTypeStr", "PhaseType", "PhaseTypeStr", "PlotType", "PlotTypeStr", "PluginType", "PluginTypeStr", "PublicDatasetType", "PublicDatasetTypeStr", "RampType", "RampTypeStr", "RecordObserverType", "RecordObserverTypeStr", "RecordProcessorType", "RecordProcessorTypeStr", "SearchPlannerType", "SearchPlannerTypeStr", "SearchRecipePostProcessType", "SearchRecipePostProcessTypeStr", "SearchRecipeType", "SearchRecipeTypeStr", "ServiceRunType", "ServiceRunTypeStr", "ServiceType", "ServiceTypeStr", "StreamExporterType", "StreamExporterTypeStr", "TimingMode", "TimingModeStr", "TransportType", "TransportTypeStr", "UIType", "UITypeStr", "URLSelectionStrategy", "URLSelectionStrategyStr", "ZMQProxyType", "ZMQProxyTypeStr"] +__all__ = ["APIRouterType", "APIRouterTypeStr", "AccumulatorType", "AccumulatorTypeStr", "AccuracyBenchmarkType", "AccuracyBenchmarkTypeStr", "AccuracyGraderType", "AccuracyGraderTypeStr", "AnalyzerType", "AnalyzerTypeStr", "ArrivalPattern", "ArrivalPatternStr", "CommClientType", "CommClientTypeStr", "CommunicationBackend", "CommunicationBackendStr", "ComposerType", "ComposerTypeStr", "ConsoleExporterType", "ConsoleExporterTypeStr", "ConvergenceCriterionType", "ConvergenceCriterionTypeStr", "CustomDatasetType", "CustomDatasetTypeStr", "DataExporterType", "DataExporterTypeStr", "DatasetBackingStoreType", "DatasetBackingStoreTypeStr", "DatasetClientStoreType", "DatasetClientStoreTypeStr", "DatasetFormat", "DatasetFormatStr", "DatasetSamplingStrategy", "DatasetSamplingStrategyStr", "EndpointType", "EndpointTypeStr", "GPUTelemetryCollectorType", "GPUTelemetryCollectorTypeStr", "GraphAdapterType", "GraphAdapterTypeStr", "PhaseType", "PhaseTypeStr", "PlotType", "PlotTypeStr", "PluginType", "PluginTypeStr", "PublicDatasetType", "PublicDatasetTypeStr", "RampType", "RampTypeStr", "RecordObserverType", "RecordObserverTypeStr", "RecordProcessorType", "RecordProcessorTypeStr", "SearchPlannerType", "SearchPlannerTypeStr", "SearchRecipePostProcessType", "SearchRecipePostProcessTypeStr", "SearchRecipeType", "SearchRecipeTypeStr", "ServiceRunType", "ServiceRunTypeStr", "ServiceType", "ServiceTypeStr", "SessionRoutingType", "SessionRoutingTypeStr", "StreamExporterType", "StreamExporterTypeStr", "TimingMode", "TimingModeStr", "TransportType", "TransportTypeStr", "UIType", "UITypeStr", "URLSelectionStrategy", "URLSelectionStrategyStr", "ZMQProxyType", "ZMQProxyTypeStr"] # Plugin Protocol Categories if TYPE_CHECKING: @@ -31,7 +31,7 @@ TimingModeStr: TypeAlias = str TimingMode = plugins.create_enum(PluginType.TIMING_STRATEGY, "TimingMode", module=__name__) -"""Dynamic enum for timing strategy. Example: TimingMode.ADAPTIVE_SCALE, TimingMode.REQUEST_RATE, TimingMode.USER_CENTRIC_RATE""" +"""Dynamic enum for timing strategy. Example: TimingMode.ADAPTIVE_SCALE, TimingMode.GRAPH_IR, TimingMode.USER_CENTRIC_RATE""" ArrivalPatternStr: TypeAlias = str ArrivalPattern = plugins.create_enum(PluginType.ARRIVAL_PATTERN, "ArrivalPattern", module=__name__) @@ -47,7 +47,7 @@ DatasetClientStoreTypeStr: TypeAlias = str DatasetClientStoreType = plugins.create_enum(PluginType.DATASET_CLIENT_STORE, "DatasetClientStoreType", module=__name__) -"""Dynamic enum for dataset client store. Example: DatasetClientStoreType.MEMORY_MAP""" +"""Dynamic enum for dataset client store. Example: DatasetClientStoreType.GRAPH_SEGMENT, DatasetClientStoreType.MEMORY_MAP""" DatasetSamplingStrategyStr: TypeAlias = str DatasetSamplingStrategy = plugins.create_enum(PluginType.DATASET_SAMPLER, "DatasetSamplingStrategy", module=__name__) @@ -65,6 +65,10 @@ PublicDatasetType = plugins.create_enum(PluginType.PUBLIC_DATASET_LOADER, "PublicDatasetType", module=__name__) """Dynamic enum for public dataset loader. Example: PublicDatasetType.AIMO, PublicDatasetType.LLAVA_ONEVISION, PublicDatasetType.VOXPOPULI""" +GraphAdapterTypeStr: TypeAlias = str +GraphAdapterType = plugins.create_enum(PluginType.GRAPH_ADAPTER, "GraphAdapterType", module=__name__) +"""Dynamic enum for graph adapter. Example: GraphAdapterType.DAG_JSONL, GraphAdapterType.NATIVE, GraphAdapterType.WEKA_TRACE""" + EndpointTypeStr: TypeAlias = str EndpointType = plugins.create_enum(PluginType.ENDPOINT, "EndpointType", module=__name__) """Dynamic enum for endpoint. Example: EndpointType.CHAT, EndpointType.IMAGE_RETRIEVAL, EndpointType.VIDEO_GENERATION""" @@ -73,6 +77,10 @@ TransportType = plugins.create_enum(PluginType.TRANSPORT, "TransportType", module=__name__) """Dynamic enum for transport. Example: TransportType.HTTP""" +SessionRoutingTypeStr: TypeAlias = str +SessionRoutingType = plugins.create_enum(PluginType.SESSION_ROUTING, "SessionRoutingType", module=__name__) +"""Dynamic enum for session routing. Example: SessionRoutingType.DYNAMO_HEADERS, SessionRoutingType.SESSION_ID_HEADER, SessionRoutingType.SMG_ROUTING_KEY""" + RecordProcessorTypeStr: TypeAlias = str RecordProcessorType = plugins.create_enum(PluginType.RECORD_PROCESSOR, "RecordProcessorType", module=__name__) """Dynamic enum for record processor. Example: RecordProcessorType.ACCURACY_RECORD, RecordProcessorType.METRIC_RECORD""" @@ -83,7 +91,7 @@ AccumulatorTypeStr: TypeAlias = str AccumulatorType = plugins.create_enum(PluginType.ACCUMULATOR, "AccumulatorType", module=__name__) -"""Dynamic enum for accumulator. Example: AccumulatorType.ACCURACY, AccumulatorType.METRIC_RESULTS, AccumulatorType.SERVER_METRICS""" +"""Dynamic enum for accumulator. Example: AccumulatorType.ACCURACY, AccumulatorType.METRIC_RESULTS, AccumulatorType.THEORETICAL_PREFIX_CACHE""" StreamExporterTypeStr: TypeAlias = str StreamExporterType = plugins.create_enum(PluginType.STREAM_EXPORTER, "StreamExporterType", module=__name__) diff --git a/src/aiperf/plugin/plugins.py b/src/aiperf/plugin/plugins.py index 2d6e5b805f..2bfeef44f6 100644 --- a/src/aiperf/plugin/plugins.py +++ b/src/aiperf/plugin/plugins.py @@ -932,6 +932,7 @@ def _load_package_metadata( from aiperf.common.protocols import CommunicationClientProtocol, CommunicationProtocol, ServiceProtocol from aiperf.controller.protocols import ServiceManagerProtocol from aiperf.dataset.composer.base import BaseDatasetComposer + from aiperf.dataset.graph.adapters.protocols import GraphAdapterProtocol from aiperf.dataset.protocols import CustomDatasetLoaderProtocol, DatasetBackingStoreProtocol, DatasetClientStoreProtocol, DatasetSamplingStrategyProtocol, PublicDatasetLoaderProtocol from aiperf.endpoints.protocols import EndpointProtocol from aiperf.exporters.protocols import ConsoleExporterProtocol, DataExporterProtocol @@ -939,7 +940,7 @@ def _load_package_metadata( from aiperf.orchestrator.convergence.base import ConvergenceCriterion from aiperf.orchestrator.search_planner.base import SearchPlanner from aiperf.plot.core.plot_type_handlers import PlotTypeHandlerProtocol - from aiperf.plugin.enums import APIRouterType, AccumulatorType, AccuracyBenchmarkType, AccuracyGraderType, AnalyzerType, ArrivalPattern, CommClientType, CommunicationBackend, ComposerType, ConsoleExporterType, ConvergenceCriterionType, CustomDatasetType, DataExporterType, DatasetBackingStoreType, DatasetClientStoreType, DatasetSamplingStrategy, EndpointType, GPUTelemetryCollectorType, PlotType, PluginType, PluginTypeStr, PublicDatasetType, RampType, RecordObserverType, RecordProcessorType, SearchPlannerType, SearchRecipePostProcessType, SearchRecipeType, ServiceRunType, ServiceType, StreamExporterType, TimingMode, TransportType, UIType, URLSelectionStrategy, ZMQProxyType + from aiperf.plugin.enums import APIRouterType, AccumulatorType, AccuracyBenchmarkType, AccuracyGraderType, AnalyzerType, ArrivalPattern, CommClientType, CommunicationBackend, ComposerType, ConsoleExporterType, ConvergenceCriterionType, CustomDatasetType, DataExporterType, DatasetBackingStoreType, DatasetClientStoreType, DatasetSamplingStrategy, EndpointType, GPUTelemetryCollectorType, GraphAdapterType, PlotType, PluginType, PluginTypeStr, PublicDatasetType, RampType, RecordObserverType, RecordProcessorType, SearchPlannerType, SearchRecipePostProcessType, SearchRecipeType, ServiceRunType, ServiceType, SessionRoutingType, StreamExporterType, TimingMode, TransportType, UIType, URLSelectionStrategy, ZMQProxyType from aiperf.post_processors.protocols import RecordObserverProtocol, RecordProcessorProtocol from aiperf.search_recipes._base import SearchRecipe from aiperf.search_recipes.post_process import PostProcessHandler @@ -949,6 +950,7 @@ def _load_package_metadata( from aiperf.timing.url_samplers import URLSelectionStrategyProtocol from aiperf.transports.base_transports import TransportProtocol from aiperf.ui.protocols import AIPerfUIProtocol + from aiperf.workers.session_routing import SessionRoutingBase from aiperf.zmq.zmq_proxy_base import BaseZMQProxy from typing import Literal, overload # @@ -994,6 +996,10 @@ def get_class(category: Literal[PluginType.PUBLIC_DATASET_LOADER, "public_datase @overload def iter_all(category: Literal[PluginType.PUBLIC_DATASET_LOADER, "public_dataset_loader"]) -> Iterator[tuple[PluginEntry, type[PublicDatasetLoaderProtocol]]]: ... @overload + def get_class(category: Literal[PluginType.GRAPH_ADAPTER, "graph_adapter"], name_or_class_path: GraphAdapterType | str) -> type[GraphAdapterProtocol]: ... + @overload + def iter_all(category: Literal[PluginType.GRAPH_ADAPTER, "graph_adapter"]) -> Iterator[tuple[PluginEntry, type[GraphAdapterProtocol]]]: ... + @overload def get_class(category: Literal[PluginType.ENDPOINT, "endpoint"], name_or_class_path: EndpointType | str) -> type[EndpointProtocol]: ... @overload def iter_all(category: Literal[PluginType.ENDPOINT, "endpoint"]) -> Iterator[tuple[PluginEntry, type[EndpointProtocol]]]: ... @@ -1002,6 +1008,10 @@ def get_class(category: Literal[PluginType.TRANSPORT, "transport"], name_or_clas @overload def iter_all(category: Literal[PluginType.TRANSPORT, "transport"]) -> Iterator[tuple[PluginEntry, type[TransportProtocol]]]: ... @overload + def get_class(category: Literal[PluginType.SESSION_ROUTING, "session_routing"], name_or_class_path: SessionRoutingType | str) -> type[SessionRoutingBase]: ... + @overload + def iter_all(category: Literal[PluginType.SESSION_ROUTING, "session_routing"]) -> Iterator[tuple[PluginEntry, type[SessionRoutingBase]]]: ... + @overload def get_class(category: Literal[PluginType.RECORD_PROCESSOR, "record_processor"], name_or_class_path: RecordProcessorType | str) -> type[RecordProcessorProtocol]: ... @overload def iter_all(category: Literal[PluginType.RECORD_PROCESSOR, "record_processor"]) -> Iterator[tuple[PluginEntry, type[RecordProcessorProtocol]]]: ... diff --git a/src/aiperf/plugin/plugins.yaml b/src/aiperf/plugin/plugins.yaml index c22541b0d3..8c49d6aad2 100644 --- a/src/aiperf/plugin/plugins.yaml +++ b/src/aiperf/plugin/plugins.yaml @@ -452,6 +452,48 @@ transport: transport_type: http url_schemes: [http, https] +# ============================================================================= +# Session Routing +# ============================================================================= +# Session-routing transforms stamp per-session identity onto outbound +# requests so an external router can pin every turn to one worker. +# Selected via --session-routing; one instance per worker. +# ============================================================================= +session_routing: + dynamo_headers: + class: aiperf.workers.session_routing:DynamoHeadersRouting + description: | + Dynamo session affinity via X-Dynamo-Session-ID plus + X-Dynamo-Parent-Session-ID on subagent children. Pair with a Dynamo + frontend running --router-session-affinity-ttl-secs. No options. + + dynamo_nvext: + class: aiperf.workers.session_routing:DynamoNvextRouting + description: | + Dynamo session affinity via nvext.session_control request-body + metadata. Two wire contracts via --session-routing-opt contract=...: + 'bind' (default, Dynamo >= v1.3.0-dev) re-binds every non-final turn + and closes on the final; 'open' (released v1.2.x, whose SessionAction + accepts only open/close) opens once per worker-session then rides bare + session_id. Only for Dynamo builds that implement session_control + (current upstream main does not; use dynamo_headers there). Sets + mutates_body (body-mutating: skipped with a one-time warning on the + graph-IR verbatim-bytes path; headers still apply). + Option: timeout_seconds (default 300). + + smg_routing_key: + class: aiperf.workers.session_routing:SmgRoutingKeyRouting + description: | + SGLang Model Gateway stickiness: emits X-SMG-Routing-Key with the + session's correlation ID for the gateway's manual routing policy + (--policy manual). No options. + + session_id_header: + class: aiperf.workers.session_routing:SessionIdHeaderRouting + description: | + Generic additive session-affinity header carrying the session's + correlation ID. Option: header_name (default X-Session-ID). + # ============================================================================= # Dataset Composers # ============================================================================= @@ -877,6 +919,13 @@ dataset_client_store: files with O(1) lookup performance. In Kubernetes mode, WorkerGroupManager downloads the dataset once per pod and workers use local mmap access. + graph_segment: + class: aiperf.dataset.graph_client_store:GraphSegmentDatasetClientStore + description: | + Graph unified segment-store client metadata carrier for graph IR runs. + Workers materialize per-node payloads from the unified store addressed + by (trace_id, node_ordinal); conversation lookups are a hard error. + # ============================================================================= # Timing Strategies # ============================================================================= @@ -910,6 +959,16 @@ timing_strategy: Users block on their previous turn (no interleaving within a user). Matches LMBenchmark behavior for KV cache benchmarking. + graph_ir: + class: aiperf.timing.strategies.graph_ir_replay:GraphIRReplayStrategy + description: | + A mode that replays a weka conversation-DAG (graph IR) by driving a + dataflow TraceExecutor per trace over the v1 credit pipeline. The strategy + OWNS trace completion (a trace is done when its executor TaskGroup drains) + and trace-admission concurrency; graph credits bypass the linear + session-slot lifecycle via CreditIssuer.issue_graph_credit. Selected + automatically for graph (weka IR) workloads. + # ============================================================================= # Arrival Pattern (Interval Generators) # ============================================================================= @@ -1057,6 +1116,17 @@ accumulator: metadata: record_types: [accuracy] + theoretical_prefix_cache: + class: aiperf.post_processors.theoretical_prefix_cache:TheoreticalPrefixCacheAccumulator + description: | + Lightweight accumulator for trace loaders (WEKA graph-IR) that stamp + per-turn theoretical prefix-cache block counts on each Turn. Emits the + cumulative infinite-cache prefix hit rate (theoretical_prefix_cache_hit, + percent) during profiling without carrying hash_ids through the request + path. Always loaded; emits nothing when no prefix-cache metadata is present. + metadata: + record_types: [metric_records] + # ============================================================================= # Stream Exporters # ============================================================================= @@ -1778,7 +1848,9 @@ public_dataset_loader: Exgentic agent LLM traces from HuggingFace. Replays complete recorded message and tool snapshots with recorded output limits. Fixed-schedule mode preserves call timestamps and overlaps; other modes replay sessions - closed-loop with recorded inter-turn delays. + closed-loop with recorded inter-turn delays. Pair with `--session-routing + dynamo_headers` for Dynamo session affinity (the loader no longer stamps a + session header). metadata: hf_dataset_name: Exgentic/agent-llm-traces hf_split: train @@ -2051,6 +2123,48 @@ public_dataset_loader: audio_column: audio streaming: true +# ============================================================================= +# Graph Adapters +# ============================================================================= +# Graph adapters convert third-party trace/log files into the canonical +# ParsedGraph used by the graph benchmark mode. Detection order is driven by +# each entry's detection_priority metadata (higher wins). +# ============================================================================= +graph_adapter: + dynamo_trace: + class: aiperf.dataset.graph.adapters.dynamo.trace:DynamoTraceAdapter + description: | + Dynamo agent-trace v1 segmented JSONL.gz directories (or single .gz/.jsonl). + metadata: + detection_priority: 100 + + weka_trace: + class: aiperf.dataset.graph.adapters.weka.trace:WekaTraceAdapter + description: | + Weka KV-cache-tester agentic-coding trace JSON (single file or directory). + metadata: + detection_priority: 85 + + dag_jsonl: + class: aiperf.dataset.graph.adapters.dag_jsonl.trace:DagJsonlGraphAdapter + description: | + DAG-shaped conversation JSONL (the dag_jsonl custom-dataset format) lowered + onto the graph IR: forks/spawns become static edges + AND fan-in gates and + accumulated history becomes interned segments plus live-reply dynamic slots. + Selected explicitly via --graph-format dag_jsonl (excluded from autodetect + so legacy --custom-dataset-type dag_jsonl runs are untouched). + metadata: + detection_priority: 90 + + native: + class: aiperf.dataset.graph.adapters.native:NativeGraphAdapter + description: | + Canonical AIPerf graph workload format (.yaml / .yml / .jsonl conforming + to the graph schema). Lowest priority fallback; reached via explicit + --graph / parse_graph, not workload auto-detection. + metadata: + detection_priority: 1 + # ============================================================================= # ZMQ Proxies # ============================================================================= diff --git a/src/aiperf/plugin/schema/plugins.schema.json b/src/aiperf/plugin/schema/plugins.schema.json index 46c237a2a8..b844204250 100644 --- a/src/aiperf/plugin/schema/plugins.schema.json +++ b/src/aiperf/plugin/schema/plugins.schema.json @@ -90,6 +90,14 @@ "$ref": "#/$defs/PublicDatasetLoaderPlugin" } }, + "graph_adapter": { + "title": "Graph Adapter Plugins", + "type": "object", + "description": "Graph adapters convert third-party trace/log files into the canonical\nParsedGraph used by the graph benchmark mode. Each adapter exposes\ncan_load(path) for format detection and parse(path, ctx) for full\nconversion, where ctx is an optional GraphParseContext of run-derived\nknobs; adapters forward a ctx field only when it is set, so a ctx-less\nparse keeps the protocol-default entry behavior. Selected by\n`--graph-format` or auto-detected from file extension + content sniff.\nDetection ordering (when multiple adapters' can_load() return True for\nthe same file) comes from each plugin's metadata `detection_priority`\nfield.", + "additionalProperties": { + "$ref": "#/$defs/GraphAdapterPlugin" + } + }, "endpoint": { "title": "Endpoint Plugins", "type": "object", @@ -106,6 +114,14 @@ "$ref": "#/$defs/TransportPlugin" } }, + "session_routing": { + "title": "Session Routing Plugins", + "type": "object", + "description": "Session-routing transforms stamp per-session identity onto outbound\nrequests (headers or body metadata) so an external router in front of\nmultiple replicas can pin every turn of a session to one worker.\nSelected via --session-routing; parameterized via repeatable\n--session-routing-opt key=value validated against the plugin's\nOptions model. mutates_body is protocol metadata marking a mode as\nbody-mutating / incompatible with any verbatim-bytes request path; it is\nnot currently enforced by any gate on this codebase.", + "additionalProperties": { + "$ref": "#/$defs/SessionRoutingPlugin" + } + }, "record_processor": { "title": "Record Processor Plugins", "type": "object", @@ -635,7 +651,7 @@ "type": "integer" }, "metadata": { - "description": "Metadata schema for custom dataset loader plugins.\n\nDefines format-specific defaults for dataset loaders. When a loader specifies\n``block_size``, it overrides the user's ``--isl-block-size`` config default,\nensuring hash-based prompt generation uses the correct token block size for the\ntrace format (e.g. 16 for Bailian, 512 for Mooncake).\n\nReferenced by: categories.yaml custom_dataset_loader.metadata_class\nUsed in: plugins.yaml custom_dataset_loader entries", + "description": "Metadata schema for custom dataset loader plugins.\n\nDefines format-specific defaults for dataset loaders. When a loader specifies\n``default_block_size``, it overrides the user's ``--isl-block-size`` config default,\nensuring hash-based prompt generation uses the correct token block size for the\ntrace format (e.g. 16 for Bailian, 512 for Mooncake).\n\nReferenced by: categories.yaml custom_dataset_loader.metadata_class\nUsed in: plugins.yaml custom_dataset_loader entries", "properties": { "is_trace": { "default": false, @@ -855,6 +871,47 @@ "title": "Public Dataset Loader Plugin", "description": "Public dataset loaders fetch shared benchmark datasets without requiring a local file.\nSupports HTTP-downloaded and HuggingFace-hosted datasets.\nOne-to-one mapping: selected based on --public-dataset configuration." }, + "GraphAdapterPlugin": { + "type": "object", + "properties": { + "class": { + "description": "Python class that implements this plugin entry. Use 'module.path:ClassName' format, e.g., 'aiperf.endpoints.openai_chat:ChatEndpoint'.", + "title": "Class", + "type": "string" + }, + "description": { + "default": "", + "description": "Brief explanation of what this plugin type does and when to use it.", + "title": "Description", + "type": "string" + }, + "priority": { + "default": 0, + "description": "Conflict resolution priority. When multiple packages register the same type name, the one with higher priority wins. Use 0 for normal plugins, higher values to override built-in implementations.", + "title": "Priority", + "type": "integer" + }, + "metadata": { + "description": "Metadata for graph_adapter plugins.\n\nReferenced by: categories.yaml graph_adapter.metadata_class", + "properties": { + "detection_priority": { + "default": 0, + "description": "Detection-order priority for registered graph adapters. Higher values run first when multiple adapters' can_load() return True for the same file; configure concrete priorities in plugins.yaml.", + "minimum": 0, + "title": "Detection Priority", + "type": "integer" + } + }, + "title": "GraphAdapterMetadata", + "type": "object" + } + }, + "required": [ + "class" + ], + "title": "Graph Adapter Plugin", + "description": "Graph adapters convert third-party trace/log files into the canonical\nParsedGraph used by the graph benchmark mode. Each adapter exposes\ncan_load(path) for format detection and parse(path, ctx) for full\nconversion, where ctx is an optional GraphParseContext of run-derived\nknobs; adapters forward a ctx field only when it is set, so a ctx-less\nparse keeps the protocol-default entry behavior. Selected by\n`--graph-format` or auto-detected from file extension + content sniff.\nDetection ordering (when multiple adapters' can_load() return True for\nthe same file) comes from each plugin's metadata `detection_priority`\nfield." + }, "EndpointPlugin": { "type": "object", "properties": { @@ -1058,6 +1115,47 @@ "title": "Transport Plugin", "description": "Transports handle the network layer for sending requests to inference servers.\nManages connection pooling, streaming, error handling, and TCP configuration.\nOne-to-one mapping based on transport_type configuration." }, + "SessionRoutingPlugin": { + "type": "object", + "properties": { + "class": { + "description": "Python class that implements this plugin entry. Use 'module.path:ClassName' format, e.g., 'aiperf.endpoints.openai_chat:ChatEndpoint'.", + "title": "Class", + "type": "string" + }, + "description": { + "default": "", + "description": "Brief explanation of what this plugin type does and when to use it.", + "title": "Description", + "type": "string" + }, + "priority": { + "default": 0, + "description": "Conflict resolution priority. When multiple packages register the same type name, the one with higher priority wins. Use 0 for normal plugins, higher values to override built-in implementations.", + "title": "Priority", + "type": "integer" + }, + "metadata": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Category-specific configuration for this plugin type. The allowed fields depend on the category's metadata_class in categories.yaml.", + "title": "Metadata" + } + }, + "required": [ + "class" + ], + "title": "Session Routing Plugin", + "description": "Session-routing transforms stamp per-session identity onto outbound\nrequests (headers or body metadata) so an external router in front of\nmultiple replicas can pin every turn of a session to one worker.\nSelected via --session-routing; parameterized via repeatable\n--session-routing-opt key=value validated against the plugin's\nOptions model. mutates_body is protocol metadata marking a mode as\nbody-mutating / incompatible with any verbatim-bytes request path; it is\nnot currently enforced by any gate on this codebase." + }, "RecordProcessorPlugin": { "type": "object", "properties": { diff --git a/src/aiperf/plugin/schema/schemas.py b/src/aiperf/plugin/schema/schemas.py index a7518ad769..8fc360a03c 100644 --- a/src/aiperf/plugin/schema/schemas.py +++ b/src/aiperf/plugin/schema/schemas.py @@ -327,11 +327,28 @@ class PlotMetadata(BaseModel): ) +class GraphAdapterMetadata(BaseModel): + """Metadata for graph_adapter plugins. + + Referenced by: categories.yaml graph_adapter.metadata_class + """ + + detection_priority: int = Field( + default=0, + ge=0, + description=( + "Detection-order priority for registered graph adapters. Higher " + "values run first when multiple adapters' can_load() return True " + "for the same file; configure concrete priorities in plugins.yaml." + ), + ) + + class CustomDatasetLoaderMetadata(BaseModel): """Metadata schema for custom dataset loader plugins. Defines format-specific defaults for dataset loaders. When a loader specifies - ``block_size``, it overrides the user's ``--isl-block-size`` config default, + ``default_block_size``, it overrides the user's ``--isl-block-size`` config default, ensuring hash-based prompt generation uses the correct token block size for the trace format (e.g. 16 for Bailian, 512 for Mooncake). diff --git a/src/aiperf/plugin/types.py b/src/aiperf/plugin/types.py index 15efe7e8c0..8e58c0d35c 100644 --- a/src/aiperf/plugin/types.py +++ b/src/aiperf/plugin/types.py @@ -138,7 +138,7 @@ def load(self) -> type: if self.loaded_class is not None: return self.loaded_class - # Validate and parse class path using structural pattern matching + # Validate and parse the 'module.path:ClassName' class path module_path, _, class_name = self.class_path.rpartition(":") if not module_path or not class_name: raise ValueError( diff --git a/src/aiperf/post_processors/base_metrics_processor.py b/src/aiperf/post_processors/base_metrics_processor.py index 25ee7df170..0b69b5c49e 100644 --- a/src/aiperf/post_processors/base_metrics_processor.py +++ b/src/aiperf/post_processors/base_metrics_processor.py @@ -47,7 +47,13 @@ def get_filters(self) -> tuple[MetricFlags, MetricFlags]: for capability, flag in capability_flags: if not getattr(endpoint_metadata, capability): disallowed_flags |= flag - if not self.run.cfg.endpoint.streaming: + if not self.run.cfg.endpoint.streaming and not self._is_graph_workload(): + # Graph-IR replays (weka / dynamo / native graph) stream per-request + # from recorded node modes even when the global --streaming flag is + # off, so the run-level STREAMING_ONLY gate must not fire for them; + # per-record applicability is enforced by the streamed_request + # predicate metric instead. Only disable STREAMING_ONLY when nothing + # in the run can stream (global flag off AND not a graph workload). disallowed_flags |= MetricFlags.STREAMING_ONLY if self.run.cfg.endpoint.use_server_token_count: # Disable usage diff metrics if server token counts are used, because @@ -69,6 +75,20 @@ def get_filters(self) -> tuple[MetricFlags, MetricFlags]: return required_flags, disallowed_flags + def _is_graph_workload(self) -> bool: + """True when the run's input is a graph-IR workload. + + Graph-IR workloads carry per-node recorded stream modes, so individual + requests can stream on the wire even when ``endpoint.streaming`` is off + run-wide. Reads the memoized single-source resolution + (``workload_detect.resolve_graph_workload``), which itself degrades + detection failures to None -- keeping the run-level STREAMING_ONLY + gate for anything that is not a graph workload. + """ + from aiperf.dataset.graph.workload_detect import resolve_graph_workload + + return resolve_graph_workload(self.run) is not None + def _configure_goodput(self, applicable_tags: set[str]) -> None: """ If --goodput SLOs are provided, wire the SLOs into the GoodRequestCountMetric. diff --git a/src/aiperf/post_processors/theoretical_prefix_cache.py b/src/aiperf/post_processors/theoretical_prefix_cache.py new file mode 100644 index 0000000000..bafd460d4e --- /dev/null +++ b/src/aiperf/post_processors/theoretical_prefix_cache.py @@ -0,0 +1,193 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Theoretical prefix-cache hit accounting for trace replay. + +A metric-record accumulator (``on_dataset_configured`` / ``process_record`` / +``summarize``) that the ``RecordsManager`` drives through record-type routing, +summing loader-provided per-turn hit/total block counts into one +ratio-of-sums PERCENT metric (``theoretical_prefix_cache_hit``). + +Lookup key +---------- +ONE join key for both planes: ``(conversation_id, turn_index)``. Non-graph +trace loaders stamp per-``TurnMetadata`` counts keyed by it directly. Graph +records carry the SAME legacy-shaped identity -- ``conversation_id`` is the +trajectory TEMPLATE id (``{trace}`` / ``{trace}::{scope}``) and +``turn_index`` the node's 0-based turn -- so the ``{scope}:{turn}`` node id +into the graph facet ``DatasetMetadata.graph.prefix_cache_by_trace`` +(``{trace_id: {node_id: [hit, total]}}``) is a pure function of the two +record fields. No correlation-id parsing anywhere. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from aiperf.common.models import MetricResult +from aiperf.metrics.types.theoretical_prefix_cache_metric import ( + TheoreticalPrefixCacheHitMetric, +) +from aiperf.post_processors.base_metrics_processor import BaseMetricsProcessor + +if TYPE_CHECKING: + from aiperf.common.messages.inference_messages import MetricRecordsData + from aiperf.common.models.dataset_models import DatasetMetadata + from aiperf.config.resolution.plan import BenchmarkRun + + +# Back-compat alias: the display metadata (header, unit, console group) lives on +# the registered TheoreticalPrefixCacheHitMetric class so the realtime dashboard +# and console exporter can resolve the tag from the MetricRegistry. +THEORETICAL_PREFIX_CACHE_HIT_TAG = TheoreticalPrefixCacheHitMetric.tag + + +def _graph_node_identity( + conversation_id: str | None, turn_index: int | None +) -> tuple[str, str] | None: + """Recover ``(template_trace_id, node_id)`` from a record's own fields. + + Graph records carry LEGACY-shaped identity: ``conversation_id`` is the + trajectory TEMPLATE id (``{trace}`` for the root scope, + ``{trace}::{scope}`` for children) and ``turn_index`` is the node's + 0-based turn within its trajectory -- so the ``{scope}:{turn}`` node id is + a pure function of the two record fields, no correlation-id parsing. + """ + if not conversation_id or turn_index is None: + return None + trace, sep, scope = conversation_id.partition("::") + return trace, f"{scope if sep else trace}:{turn_index}" + + +class TheoreticalPrefixCacheAccumulator(BaseMetricsProcessor): + """Track infinite-cache prefix hits from loader-provided per-turn counts. + + Consumes leading prefix-cache ``hit_blocks`` and ``total_blocks`` integers + stamped per node on the native ``LlmNode.theoretical_prefix_cache_hit_blocks`` + / ``_total_blocks`` fields by the shared trie build + (``segment_ir.prefix_cache.stamp_theoretical_prefix_cache``, weka and + dynamo), or per turn on ``Turn`` (non-graph trace loaders). Runtime accounting + therefore avoids carrying hash_ids or re-tokenizing prompts; each completed + record is only a metadata lookup plus two integer additions. + + Emits the cumulative ``theoretical_prefix_cache_hit`` (percent) = ``100 * + sum(hit_blocks) / sum(total_blocks)`` over valid profiling records. The + ``RecordsManager`` only forwards profiling-phase records to results + processors, so no explicit phase gate is needed here. + """ + + def __init__(self, run: BenchmarkRun, **kwargs: Any) -> None: + super().__init__(run=run, **kwargs) + # Per-node counts keyed by (trace_id, node_id) for the graph dispatch + # path. The trace_id qualifier keeps the bare per-trace node ids + # (``parent_0`` repeats in every trace of a corpus) disjoint. + self._blocks_by_node: dict[tuple[str, str], tuple[int, int]] = {} + # Per-turn counts keyed by (conversation_id, turn_index), for non-graph + # trace loaders that stamp the TurnMetadata fields. + self._turn_blocks_by_conversation: dict[ + str, tuple[tuple[int, int] | None, ...] + ] = {} + self._hit_blocks = 0 + self._total_blocks = 0 + self._enabled = False + + def on_dataset_configured(self, metadata: DatasetMetadata) -> None: + """Receive per-node (graph facet) / per-turn prefix-cache metadata.""" + by_node: dict[tuple[str, str], tuple[int, int]] = {} + # Graph facet: per-trace {node_id: [hit, total]}, keyed by the BASE + # template trace id -- the same template identity records carry in + # conversation_id (instance identity rides x_correlation_id). + graph = metadata.graph + if graph is not None and graph.prefix_cache_by_trace: + for trace_id, node_map in graph.prefix_cache_by_trace.items(): + for node_id, counts in node_map.items(): + if len(counts) >= 2: + by_node[(trace_id, node_id)] = ( + int(counts[0]), + int(counts[1]), + ) + # Per-turn counts for non-graph trace loaders that stamp the + # TurnMetadata fields. + by_conv: dict[str, tuple[tuple[int, int] | None, ...]] = {} + for conv in metadata.conversations: + per_turn: list[tuple[int, int] | None] = [] + has_turn_metadata = False + for turn in conv.turns: + hit_blocks = turn.theoretical_prefix_cache_hit_blocks + total_blocks = turn.theoretical_prefix_cache_total_blocks + if hit_blocks is None or total_blocks is None: + per_turn.append(None) + continue + has_turn_metadata = True + per_turn.append((hit_blocks, total_blocks)) + if has_turn_metadata: + by_conv[conv.conversation_id] = tuple(per_turn) + self._blocks_by_node = by_node + self._turn_blocks_by_conversation = by_conv + self._enabled = bool(by_node) or bool(by_conv) + + def _lookup_counts( + self, conversation_id: str | None, turn_index: int | None + ) -> tuple[int, int] | None: + """Resolve (hit_blocks, total_blocks) for one record, node map first. + + The graph node map keys by ``(template_trace_id, node_id)``; both are + derived from the record's template-level ``(conversation_id, + turn_index)`` -- the SAME join key the linear per-turn fallback uses, + so recycled instances of one template correctly re-apply the template + counts (duplication is desired, exactly like the linear path). + """ + identity = _graph_node_identity(conversation_id, turn_index) + if identity is not None: + counts = self._blocks_by_node.get(identity) + if counts is not None: + return counts + if conversation_id is None or turn_index is None: + return None + per_turn = self._turn_blocks_by_conversation.get(conversation_id) + if per_turn is None or turn_index < 0 or turn_index >= len(per_turn): + return None + return per_turn[turn_index] + + async def process_record(self, record_data: MetricRecordsData) -> None: + """Accumulate block counts for one successful profiling request.""" + if not self._enabled or not record_data.valid: + return + metadata = record_data.metadata + # A context-overflow skip record reaches this accumulator as a + # trimmed, error-free carrier for the overflow count (see + # RecordsManager._send_overflow_count_only): the request never ran to + # completion, so counting its planned blocks would pollute the hit rate. + if metadata.context_overflow_skip: + return + counts = self._lookup_counts(metadata.conversation_id, metadata.turn_index) + if counts is None: + return + hit_blocks, total_blocks = counts + if total_blocks <= 0: + return + # Clamp the hit count into [0, total_blocks]: a loader miscount must not + # drive the cumulative hit rate above 100% (or below 0%). + hit_blocks = max(0, min(hit_blocks, total_blocks)) + self._hit_blocks += hit_blocks + self._total_blocks += total_blocks + + async def export_results(self, ctx: object) -> list[MetricResult]: + """Export final theoretical prefix-cache metrics.""" + return await self.summarize() + + async def summarize(self) -> list[MetricResult]: + """Return the current cumulative theoretical prefix-cache hit rate.""" + if self._total_blocks <= 0: + return [] + hit_rate_pct = 100.0 * self._hit_blocks / self._total_blocks + return [ + MetricResult( + tag=TheoreticalPrefixCacheHitMetric.tag, + header=TheoreticalPrefixCacheHitMetric.header, + unit=str(TheoreticalPrefixCacheHitMetric.unit), + count=self._total_blocks, + current=hit_rate_pct, + avg=hit_rate_pct, + sum=self._hit_blocks, + ) + ] diff --git a/src/aiperf/records/inference_result_parser.py b/src/aiperf/records/inference_result_parser.py index f9b340de00..4775046c15 100644 --- a/src/aiperf/records/inference_result_parser.py +++ b/src/aiperf/records/inference_result_parser.py @@ -25,6 +25,7 @@ ToolCallResponseData, find_last_non_empty_usage, ) +from aiperf.common.scenario import is_context_overflow_response from aiperf.common.tokenizer import Tokenizer from aiperf.plugin import plugins from aiperf.plugin.enums import PluginType @@ -146,6 +147,19 @@ async def parse_request_record( # Make sure any invalid request records are converted to error records for combined processing. request_record.create_error_from_invalid() + # Classify context-overflow errors per InferenceX AgentX RFC §7. Runs + # unconditionally (cheap substring scan, allowlist-driven) so the + # ``context_overflow_count`` metric can aggregate even outside scenario + # mode. The boolean is consumed by ``ContextOverflowCountMetric`` via + # ``record.request.context_overflow``. + if request_record.has_error and request_record.error is not None: + try: + request_record.context_overflow = is_context_overflow_response( + body=request_record.error.message, + ) + except Exception: + # Detection is best-effort -- never surface as a parse error. + request_record.context_overflow = False # One payload decode + walk per record, shared by the ISL tokeniser and # the MediaCounts builder. Both valid and error records go through this. inputs = self._extract_payload_inputs_for_record(request_record) @@ -315,29 +329,34 @@ async def compute_input_token_count( ) -> int | None: """Compute the number of input (prompt) tokens for a request record. - Source of truth is ``request_info.payload_bytes`` -- the exact JSON the - endpoint sent on the wire -- decoded once and walked by + Primary source of truth is ``request_info.payload_bytes`` -- the exact + JSON the endpoint sent on the wire -- decoded once and walked by ``extract_payload_inputs`` into tokenisable ``texts`` (plus a chat-shape - ``messages`` view and any ``pretokenised_token_count``). ``turns`` are - never read here: on the payload-bytes fast path they are a content-free - stub, and the canonical body is the bytes. + ``messages`` view and any ``pretokenised_token_count``). - ``system_message`` / ``user_context_message`` are NOT added here -- the - endpoint's ``format_payload`` already inlined them into ``payload_bytes``, - so re-adding would double-count. + ``system_message`` / ``user_context_message`` are NOT added on the + payload-bytes path -- the endpoint's ``format_payload`` already inlined + them into ``payload_bytes``, so re-adding would double-count. Direct + callers without payload bytes fall back to the legacy turn-text walk. When ``--apply-chat-template`` is set AND the payload carries a chat ``messages`` array AND the HF tokenizer has a chat template, the templated count (role/header wrapping + generation prompt) is returned; otherwise the bare ``texts`` are encoded with a space separator. Any ``pretokenised_token_count`` (e.g. token-id embeddings input) is added on - top. Returns ``None`` when no ``payload_bytes`` is available or the + top. Returns ``None`` when no payload/legacy text is available or the endpoint reports nothing tokenisable. """ if inputs is None: inputs = self._extract_payload_inputs_for_record(request_record) if inputs is None: + prompt_texts = self._legacy_prompt_texts_from_record(request_record) + if prompt_texts: + tokenizer = await self.get_tokenizer(request_record.model_name) + return await self._compute_token_count( + tokenizer, prompt_texts, separator=" " + ) if not ( request_record.request_info and request_record.request_info.payload_bytes @@ -389,6 +408,26 @@ async def compute_input_token_count( # Pure pre-tokenised input (e.g. token-id embeddings) carries no text. return pretokenised if pretokenised > 0 else None + def _legacy_prompt_texts_from_record( + self, request_record: RequestRecord + ) -> list[str]: + """Best-effort prompt texts for callers that lack payload bytes.""" + prompt_texts: list[str] = [] + request_info = request_record.request_info + if request_info is not None: + if request_info.system_message: + prompt_texts.append(request_info.system_message) + if request_info.user_context_message: + prompt_texts.append(request_info.user_context_message) + turns = request_info.turns or request_record.turns + else: + turns = request_record.turns + + for turn in turns: + for text in turn.texts: + prompt_texts.append("".join(text.contents)) + return prompt_texts + async def _compute_chat_template_token_count( self, tokenizer: Tokenizer, diff --git a/src/aiperf/records/record_processor_service.py b/src/aiperf/records/record_processor_service.py index 56644e57fd..c03cdd6e7f 100644 --- a/src/aiperf/records/record_processor_service.py +++ b/src/aiperf/records/record_processor_service.py @@ -71,6 +71,19 @@ def __init__( self.records_push_client: PushClientProtocol = self.comms.create_push_client( CommAddress.RECORDS, ) + # Drop context-overflow records from the user-facing metric pipeline when + # the active workload is a graph-IR (weka / dynamo / native) run. The + # GraphIRReplayStrategy already terminates the overflowed trajectory early + # (the overflowed node stops dispatching downstream turns), so surfacing + # the overflow as an error record would pollute the error tracker and the + # performance-metric accumulators with an event we intentionally tolerate. + # The accessor degrades detection failures to None itself, so a + # non-graph (or undetectable) input keeps default error emission. + from aiperf.dataset.graph.workload_detect import resolve_graph_workload + + self._drop_graph_overflow_records: bool = ( + resolve_graph_workload(self.run) is not None + ) self.tokenizers: dict[str, Tokenizer] = {} self.tokenizer_lock: asyncio.Lock = asyncio.Lock() self.model_endpoint: ModelEndpointInfo = ModelEndpointInfo.from_run(self.run) @@ -308,6 +321,26 @@ async def _process_and_forward_record( record, message.service_id, last_response_perf_ns ) + # Flag context-overflow records for the records-side "skip" path on a + # graph-IR run. RecordsManager counts the record toward its success + # counter (so the completion barrier converges) and still aggregates + # ``context_overflow_count`` (for the submission-rate gate), but excludes + # it from the error tracker and the performance-metric accumulators so the + # overflow never shows up in a user-facing latency/ISL/OSL/TTFT/ITL + # metric. + if getattr(self, "_drop_graph_overflow_records", False) and getattr( + record, "context_overflow", False + ): + metadata.context_overflow_skip = True + self.debug( + lambda r=record: ( + "graph-IR: flagging context-overflow record as metrics-skip " + f"(credit={r.request_info.credit_num} " + f"conv={r.request_info.conversation_id} " + f"turn={r.request_info.turn_index})" + ) + ) + # Stage 1 - producers: run concurrently, group outputs by declared channel. by_type: dict[str, list[Any]] = {} producer_results = await asyncio.gather( diff --git a/src/aiperf/records/records_manager.py b/src/aiperf/records/records_manager.py index cd9fdb6f83..720654ac07 100644 --- a/src/aiperf/records/records_manager.py +++ b/src/aiperf/records/records_manager.py @@ -22,6 +22,7 @@ CommandType, CreditPhase, MessageType, + MetricValueTypeT, make_result_producer_capability, ) from aiperf.common.environment import Environment @@ -29,6 +30,7 @@ from aiperf.common.messages import ( AllRecordsReceivedMessage, DatasetConfiguredNotification, + MetricRecordsData, NetworkLatencyRecordMessage, ProcessAccuracyResultMessage, ProcessAllResultsMessage, @@ -96,6 +98,11 @@ from aiperf.config.config import BenchmarkConfig from aiperf.config.resolution.plan import BenchmarkRun +# Metric tag aggregated for a context-overflow-skip record so the submission-rate +# gate keeps a correct numerator (see ``_send_overflow_count_only``). Kept as a +# constant -- the tag is the metric class's stable identifier. +CONTEXT_OVERFLOW_COUNT_TAG = "context_overflow_count" + _LATENCY_LINE_LABELS: tuple[tuple[str, str], ...] = ( ("ttft", "time_to_first_token"), @@ -390,6 +397,13 @@ def __init__( self._records_tracker = RecordsTracker() self._error_tracker = ErrorTracker() + # Count of context-overflow records routed through the graph-IR + # metrics-skip path (``MetricRecordMetadata.context_overflow_skip``). + # These advance the records-side success counter (barrier lockstep) and + # still aggregate ``context_overflow_count`` for the submission-rate gate, + # but are excluded from the error tracker and the performance-metric + # accumulators. + self._skipped_context_overflow_count: int = 0 # DatasetConfiguredNotification (SUB) and metric records (PULL) arrive on # independent channels with no ordering guarantee. Gate record processing on # this event so results processors are configured (e.g. accuracy task names) @@ -581,11 +595,34 @@ async def _on_records(self, message: RecordsMessage) -> None: if self.is_trace_enabled: self.trace(f"Received records: {message}") + phase = message.metadata.benchmark_phase + + # Context-overflow records in a graph-IR run bypass the normal user-facing + # per-record processing. They are EXCLUDED from the error tracker and from + # the performance-metric accumulators (ISL/OSL/TTFT/ITL/latency), but still + # (a) advance the records-side SUCCESS counter so the completion barrier + # converges in lockstep with the credit side, and (b) aggregate ONLY the + # ``context_overflow_count`` metric so the submission-rate gate stays + # correct. The ``context_overflow_count`` aggregate (the submission-rate + # input) is kept alive by forwarding a trimmed record carrying that + # single metric. + if getattr(message.metadata, "context_overflow_skip", False): + self._skipped_context_overflow_count += 1 + await self._send_overflow_count_only(message) + self._records_tracker.update_from_request(message.metadata, None) + if ( + phase in self._complete_credit_phases + and self._records_tracker.check_and_set_all_records_received_for_phase( + phase + ) + ): + await self._handle_all_records_received(phase) + return + dispatch_errors: list[BaseException] = [] for record in message.records: dispatch_errors.extend(await self._dispatch_record(record)) - phase = message.metadata.benchmark_phase self._records_tracker.update_from_request(message.metadata, message.error) if message.error: self._error_tracker.increment_error_count_for_phase(phase, message.error) @@ -717,6 +754,39 @@ async def _finalize_and_process_results( await self._process_results(phase=phase, cancelled=cancelled) self.info("_finalize_and_process_results completed") + async def _send_overflow_count_only(self, message: RecordsMessage) -> None: + """Forward ONLY the ``context_overflow_count`` metric for a skip record. + + A context-overflow record is excluded from every performance-metric + accumulator (latency/ISL/OSL/TTFT/ITL) and from the error tracker, but + its overflow count must still reach the ``context_overflow_count`` + aggregate that the InferenceX submission-rate gate reads. We locate the + ``MetricRecordsData`` in the request envelope, build a trimmed one + carrying that single metric (and no error), and route it through the + metadata-derived record dispatcher. + + The trimmed record reuses the original ``metadata`` (whose + ``context_overflow_skip`` flag stays True): processors that must not + count overflow events at all (e.g. ``TheoreticalPrefixCacheAccumulator``, + which would otherwise fold the node's planned blocks into the hit rate) + key off that flag to skip it. + """ + overflow_value: MetricValueTypeT | None = None + for record in message.records: + metrics = getattr(record, "metrics", None) + if metrics and CONTEXT_OVERFLOW_COUNT_TAG in metrics: + overflow_value = metrics[CONTEXT_OVERFLOW_COUNT_TAG] + break + if overflow_value is None: + return + trimmed = MetricRecordsData( + metadata=message.metadata, + metrics={CONTEXT_OVERFLOW_COUNT_TAG: overflow_value}, + trace_data=None, + error=None, + ) + await self._dispatch_record(trimmed) + @on_message(MessageType.DATASET_CONFIGURED_NOTIFICATION) async def _on_dataset_configured( self, message: DatasetConfiguredNotification diff --git a/src/aiperf/search_recipes/_max_concurrency_under_sla.py b/src/aiperf/search_recipes/_max_concurrency_under_sla.py index 2fe8463166..7b024230f7 100644 --- a/src/aiperf/search_recipes/_max_concurrency_under_sla.py +++ b/src/aiperf/search_recipes/_max_concurrency_under_sla.py @@ -79,7 +79,7 @@ class MaxConcurrencyUnderSLA(SearchRecipe): x_metric="request_latency", x_stat="p95", y_metric="concurrency", # parameter-as-axis: read from variation params - y_stat="value", # sentinel: see _extract_axis_value in Task 6 + y_stat="value", # sentinel: see _extract_axis_value in cli_runner/_pareto.py ) _CONCURRENCY_PATH: ClassVar[str] = "phases.profiling.concurrency" diff --git a/src/aiperf/server_metrics/csv_exporter.py b/src/aiperf/server_metrics/csv_exporter.py index 932e6f75c1..21cdec7a0c 100644 --- a/src/aiperf/server_metrics/csv_exporter.py +++ b/src/aiperf/server_metrics/csv_exporter.py @@ -58,7 +58,6 @@ class CsvMetricInfo(NamedTuple): description: Metric description from HELP text unit: Inferred unit string (e.g., "seconds", "bytes") or None stats: Type-specific statistics (Gauge, Counter, or Histogram series) - is_derived: Whether this is a derived metric """ endpoint: str @@ -111,8 +110,12 @@ def __init__(self, exporter_config: ExporterConfig, **kwargs) -> None: super().__init__(exporter_config, **kwargs) self._file_path = exporter_config.cfg.artifacts.server_metrics_export_csv_file self.trace_or_debug( - lambda: f"Initializing ServerMetricsCsvExporter with config: {exporter_config}", - lambda: f"Initializing ServerMetricsCsvExporter with file path: {self._file_path}", + lambda: ( + f"Initializing ServerMetricsCsvExporter with config: {exporter_config}" + ), + lambda: ( + f"Initializing ServerMetricsCsvExporter with file path: {self._file_path}" + ), ) def get_export_info(self) -> FileExportInfo: diff --git a/src/aiperf/server_metrics/data_collector.py b/src/aiperf/server_metrics/data_collector.py index b484f6b0ad..36a91fb2c9 100644 --- a/src/aiperf/server_metrics/data_collector.py +++ b/src/aiperf/server_metrics/data_collector.py @@ -306,7 +306,9 @@ def _parse_metrics_to_records( # cumulatively over server lifetime, not per-benchmark period if family.name not in self._seen_summary_metrics: self.info( - lambda name=family.name: f"Skipping unsupported summary metric: {name}" + lambda name=family.name: ( + f"Skipping unsupported summary metric: {name}" + ) ) self._seen_summary_metrics.add(family.name) continue diff --git a/src/aiperf/server_metrics/json_exporter.py b/src/aiperf/server_metrics/json_exporter.py index 00a290bca6..f046ebd760 100644 --- a/src/aiperf/server_metrics/json_exporter.py +++ b/src/aiperf/server_metrics/json_exporter.py @@ -65,8 +65,12 @@ def __init__(self, exporter_config: ExporterConfig, **kwargs) -> None: super().__init__(exporter_config, **kwargs) self._file_path = exporter_config.cfg.artifacts.server_metrics_export_json_file self.trace_or_debug( - lambda: f"Initializing ServerMetricsJsonExporter with config: {exporter_config}", - lambda: f"Initializing ServerMetricsJsonExporter with file path: {self._file_path}", + lambda: ( + f"Initializing ServerMetricsJsonExporter with config: {exporter_config}" + ), + lambda: ( + f"Initializing ServerMetricsJsonExporter with file path: {self._file_path}" + ), ) def get_export_info(self) -> FileExportInfo: diff --git a/src/aiperf/server_metrics/manager.py b/src/aiperf/server_metrics/manager.py index efbe53edff..d89187c15d 100644 --- a/src/aiperf/server_metrics/manager.py +++ b/src/aiperf/server_metrics/manager.py @@ -42,13 +42,13 @@ class ServerMetricsManager(BaseComponentService): """Coordinates multiple ServerMetricsDataCollector instances for server metrics collection. The ServerMetricsManager coordinates multiple ServerMetricsDataCollector instances - to collect server metrics from multiple Prometheus endpoints and send unified - ServerMetricsRecordsMessage to RecordsManager. + to collect server metrics from multiple Prometheus endpoints and send + ServerMetricsRecordMessage to RecordsManager. This service: - Manages lifecycle of ServerMetricsDataCollector instances - Collects metrics from multiple Prometheus endpoints - - Sends ServerMetricsRecordsMessage to RecordsManager via message system + - Sends ServerMetricsRecordMessage to RecordsManager via message system - Handles errors gracefully with ErrorDetails - Follows centralized architecture patterns @@ -133,7 +133,9 @@ async def _profile_configure_command( for endpoint_url in self._server_metrics_endpoints: self.debug( - lambda url=endpoint_url: f"Server Metrics: Testing reachability of {url}" + lambda url=endpoint_url: ( + f"Server Metrics: Testing reachability of {url}" + ) ) collector = ServerMetricsDataCollector( endpoint_url=endpoint_url, @@ -148,11 +150,15 @@ async def _profile_configure_command( if is_reachable: self._collectors[endpoint_url] = collector self.debug( - lambda url=endpoint_url: f"Server Metrics: Prometheus endpoint {url} is reachable" + lambda url=endpoint_url: ( + f"Server Metrics: Prometheus endpoint {url} is reachable" + ) ) else: self.debug( - lambda url=endpoint_url: f"Server Metrics: Prometheus endpoint {url} is not reachable" + lambda url=endpoint_url: ( + f"Server Metrics: Prometheus endpoint {url} is not reachable" + ) ) except Exception as e: self.error(f"Server Metrics: Exception testing {endpoint_url}: {e}") @@ -178,7 +184,9 @@ async def _profile_configure_command( await collector.initialize() await collector.collect_and_process_metrics() self.debug( - lambda url=endpoint_url: f"Server Metrics: Captured baseline from {url}" + lambda url=endpoint_url: ( + f"Server Metrics: Captured baseline from {url}" + ) ) except Exception as e: self.warning( @@ -326,7 +334,9 @@ async def _handle_profile_complete_command( try: await collector.collect_and_process_metrics() self.debug( - lambda url=endpoint_url: f"Server Metrics: Captured final state from {url}" + lambda url=endpoint_url: ( + f"Server Metrics: Captured final state from {url}" + ) ) except Exception as e: self.warning( diff --git a/src/aiperf/server_metrics/parquet_exporter.py b/src/aiperf/server_metrics/parquet_exporter.py index 4431543a37..cc8a773dc0 100644 --- a/src/aiperf/server_metrics/parquet_exporter.py +++ b/src/aiperf/server_metrics/parquet_exporter.py @@ -112,8 +112,12 @@ def __init__( self._accumulator = server_metrics_accumulator self._time_filter = time_filter self.trace_or_debug( - lambda: f"Initializing ServerMetricsParquetExporter with run cfg: {self.run.cfg}", - lambda: f"Initializing ServerMetricsParquetExporter with file path: {self._file_path}", + lambda: ( + f"Initializing ServerMetricsParquetExporter with run cfg: {self.run.cfg}" + ), + lambda: ( + f"Initializing ServerMetricsParquetExporter with file path: {self._file_path}" + ), ) def get_export_info(self) -> FileExportInfo: @@ -187,7 +191,9 @@ async def export(self) -> FileExportInfo: writer.write_table(table) total_rows += len(batch_rows) self.trace( - lambda: f"Wrote batch of {len(batch_rows):,} rows (total: {total_rows:,})" + lambda: ( + f"Wrote batch of {len(batch_rows):,} rows (total: {total_rows:,})" + ) ) batch_rows = [] @@ -206,8 +212,9 @@ async def export(self) -> FileExportInfo: batch_count = len(batch_rows) current_total = total_rows self.trace( - lambda batch_count=batch_count, - current_total=current_total: f"Wrote batch of {batch_count:,} rows (total: {current_total:,})" + lambda batch_count=batch_count, current_total=current_total: ( + f"Wrote batch of {batch_count:,} rows (total: {current_total:,})" + ) ) batch_rows = [] diff --git a/src/aiperf/server_metrics/storage.py b/src/aiperf/server_metrics/storage.py index 0e0e8396f2..fc1dcc5903 100644 --- a/src/aiperf/server_metrics/storage.py +++ b/src/aiperf/server_metrics/storage.py @@ -54,8 +54,8 @@ class ServerMetricsTimeSeries: >>> # Add first record >>> record1 = ServerMetricsRecord(timestamp_ns=1000, metrics={...}) >>> ts.append_snapshot(record1) - >>> # Add duplicate (same metrics) - >>> record2 = ServerMetricsRecord(timestamp_ns=2000, is_duplicate=True) + >>> # Add duplicate (same metrics; empty-metrics records are skipped entirely) + >>> record2 = ServerMetricsRecord(timestamp_ns=2000, metrics={...}, is_duplicate=True) >>> ts.append_snapshot(record2) >>> ts._total_fetch_count # 2 fetches 2 diff --git a/src/aiperf/timing/_branch_orchestrator_spawn.py b/src/aiperf/timing/_branch_orchestrator_spawn.py index 1c79b2a51a..8612b66f3b 100644 --- a/src/aiperf/timing/_branch_orchestrator_spawn.py +++ b/src/aiperf/timing/_branch_orchestrator_spawn.py @@ -42,8 +42,20 @@ class BranchOrchestratorSpawnMixin: which spawns worker subprocesses. """ - async def _fire_pre_session_children(self, branch: ConversationBranchInfo) -> None: - """Issue turn-0 for every child of a pre-session SPAWN branch.""" + async def _fire_pre_session_children( + self, branch: ConversationBranchInfo, parent_conversation_id: str + ) -> None: + """Issue turn-0 for every child of a pre-session SPAWN branch. + + Each issued child is tracked as a LIVE pre-session descendant of + ``parent_conversation_id``. No root session exists yet (pre-dispatch + runs before sampling mints any root correlation id), so the children + cannot be registered with the session-tree registry here; the + orchestrator folds the still-live set into the first root instance's + tree at that root's turn-0 return + (``_fold_pre_session_descendants``), keeping the root's final-turn + ``is_tree_final`` stamp conservative while these children run. + """ for child_cid in branch.child_conversation_ids: try: child_session = self._cs.start_pre_session_child(child_cid) @@ -54,6 +66,11 @@ async def _fire_pre_session_children(self, branch: ConversationBranchInfo) -> No issued = await self._issuer.dispatch_first_turn(child_session) if issued: self.stats.children_spawned += 1 + child_corr = child_session.x_correlation_id + self._pre_child_conv[child_corr] = parent_conversation_id + self._pre_live_by_conv.setdefault(parent_conversation_id, set()).add( + child_corr + ) else: # ``dispatch_first_turn`` only returns False under # stop-condition refusal (cap reached). Tally as truncated. @@ -84,6 +101,14 @@ async def _spawn_children_and_register_gates( if all_children: self._descendant_counts.setdefault(parent_corr, 0) self._descendant_counts[parent_corr] += len(all_children) + # Per-tree finality ledger keys on the ROOT (not the parent), so a + # depth>1 spawn still folds into the depth-0 root's tree. Registered + # BEFORE the dispatch gather; a child whose dispatch does not land is + # decremented again in _rollback_failed_child. + if self._session_tree_registry is not None: + self._session_tree_registry.register_descendants( + credit.effective_root_correlation_id, len(all_children) + ) # If any expected gate had zero children actually register, still # create a future-join entry so the drain logic sees it and fires. @@ -180,6 +205,7 @@ def _start_one_child( parent_correlation_id=parent_corr, child_conversation_id=child_conv_id, agent_depth=parent_depth + 1, + root_correlation_id=credit.effective_root_correlation_id, branch_mode=branch.mode, ) except Exception: @@ -187,6 +213,9 @@ def _start_one_child( self.stats.children_errored += 1 return None child_corr = child.x_correlation_id + # Capture the child's tree root so descendant-done can key the registry + # by root even for depth>1 descendants. + self._child_root[child_corr] = credit.effective_root_correlation_id self._child_modes[child_corr] = branch.mode # FORK-mode children sticky-route to the parent's worker; SPAWN-mode # children do not register a refcount. @@ -237,6 +266,10 @@ def _rollback_failed_child( child_corr = child.x_correlation_id child_mode = per_child_branch_mode.get(child_corr) self._child_modes.pop(child_corr, None) + # Undo the register_descendants that ran for this child at spawn time. + rolled_root = self._child_root.pop(child_corr, None) + if self._session_tree_registry is not None and rolled_root is not None: + self._session_tree_registry.on_descendant_done(rolled_root) entries = self._child_to_join.pop(child_corr, []) for entry in entries: if entry.prereq_key is None: @@ -257,10 +290,17 @@ def _rollback_failed_child( self._sticky_router.release_child_routing(parent_corr) if parent_corr in self._descendant_counts: self._descendant_counts[parent_corr] -= 1 - # Three-way classification of non-True gather results: - # * BaseException -> genuine error. - # * False -> stop-condition refusal; tally as truncated. - # * None -> issuer suppressed silently; observable no-op. + self._tally_failed_child_result(result, child_corr) + self.stats.children_spawned -= 1 + + def _tally_failed_child_result(self, result: object, child_corr: str) -> None: + """Fold a non-True ``dispatch_first_turn`` gather result into stats. + + Three-way classification of non-True gather results: + * BaseException -> genuine error. + * False -> stop-condition refusal; tally as truncated. + * None -> issuer suppressed silently; observable no-op. + """ if isinstance(result, BaseException): logger.error( "dispatch_first_turn failed for child %s", @@ -279,7 +319,6 @@ def _rollback_failed_child( child_corr, ) self.stats.children_errored += 1 - self.stats.children_spawned -= 1 async def _drain_vestigial_gates(self, parent_corr: str) -> None: """Drain any zero-outstanding gates created in this spawn cycle. diff --git a/src/aiperf/timing/_branch_orchestrator_state.py b/src/aiperf/timing/_branch_orchestrator_state.py index 941d8a6657..62dee76e4d 100644 --- a/src/aiperf/timing/_branch_orchestrator_state.py +++ b/src/aiperf/timing/_branch_orchestrator_state.py @@ -72,10 +72,20 @@ class PendingBranchJoin: parent_num_turns: int parent_agent_depth: int = 0 parent_parent_correlation_id: str | None = None + parent_root_correlation_id: str | None = None + """Root of the parent's session tree (the parent credit's own + ``root_correlation_id``): None when the parent IS the tree root, set for a + nested parent. Re-stamped onto the gated join ``TurnToSend`` so the resumed + parent keeps its tree lineage for finality accounting.""" gated_turn_index: int | None = None outstanding: dict[str, PrereqState] = field(default_factory=dict) parent_branch_mode: ConversationBranchMode = ConversationBranchMode.FORK parent_has_forks_on_gated_turn: bool = False + parent_has_branches_on_gated_turn: bool = False + """True iff the gated turn itself declares ANY branch (FORK or SPAWN) in its + ``branch_ids``. Superset of the FORK-only flag above; re-stamped onto the + join ``TurnToSend`` so finality stamping stays conservative when the + resumed turn will spawn further descendants.""" is_blocked: bool = False created_at_ns: int = field(default_factory=time.monotonic_ns) diff --git a/src/aiperf/timing/branch_orchestrator.py b/src/aiperf/timing/branch_orchestrator.py index 1ec3f64da1..d8320386f2 100644 --- a/src/aiperf/timing/branch_orchestrator.py +++ b/src/aiperf/timing/branch_orchestrator.py @@ -34,6 +34,7 @@ from aiperf.credit.sticky_router import StickyCreditRouter from aiperf.credit.structs import Credit from aiperf.timing.conversation_source import ConversationSource, SampledSession + from aiperf.timing.session_tree import SessionTreeRegistry __all__ = [ "BranchOrchestrator", @@ -84,11 +85,32 @@ def __init__( sticky_router: StickyCreditRouter | None = None, *, benchmark_id: str = "unknown", + session_tree_registry: SessionTreeRegistry | None = None, ) -> None: self._cs = conversation_source self._issuer = credit_issuer self._sticky_router = sticky_router self._benchmark_id = benchmark_id + # Per-tree finality ledger (DAG datasets only). The orchestrator drives + # descendant registration + descendant/root terminal transitions; the + # issuer opens the tree and reads it for finality stamping. None on + # non-DAG paths. Finality bookkeeping only -- never a slot ledger. + self._session_tree_registry = session_tree_registry + # child_x_correlation_id -> its tree's root_correlation_id, captured at + # spawn so descendant-done can key the registry by root (parent != root + # for depth>1). Popped on completion/rollback. + self._child_root: dict[str, str] = {} + # Pre-session SPAWN children fire once per phase BEFORE any root + # session is sampled, so no root correlation id exists yet at dispatch. + # Track them per parent conversation until the first root instance's + # turn-0 return folds them into that root's session tree + # (register_descendants), so the root's final turn cannot stamp + # is_tree_final=True while a pre-session child is still live. + # _pre_child_conv: pre-child corr -> parent conversation_id (popped at + # the child's terminal). _pre_live_by_conv: conversation_id -> live + # pre-child corrs not yet folded into a tree (popped at fold). + self._pre_child_conv: dict[str, str] = {} + self._pre_live_by_conv: dict[str, set[str]] = {} self._child_modes: dict[str, ConversationBranchMode] = {} # Two-level pending-join state: a "future" join is registered at # spawn time and promoted to "active" once the parent reaches the @@ -157,7 +179,7 @@ async def dispatch_pre_session_branches(self) -> None: continue if branch.branch_id not in turn0_branch_ids: continue - await self._fire_pre_session_children(branch) + await self._fire_pre_session_children(branch, conv.conversation_id) self._pre_dispatched_branches.add( (conv.conversation_id, branch.branch_id) ) @@ -186,8 +208,65 @@ async def intercept(self, credit: Credit) -> bool: branch_ids = self.get_branch_ids(credit) if branch_ids: await self._spawn_children_and_register_gates(credit, branch_ids) + # Fold this conversation's live pre-session SPAWN children into the + # root instance's tree at its turn-0 return -- the earliest moment + # the root's correlation id meets the orchestrator (pre-dispatch + # fires before sampling, so no root id was knowable then). Must run + # BEFORE the root-terminal transition below so a single-turn root's + # tree does not retire while its pre-session children are live. + if credit.agent_depth == 0 and credit.turn_index == 0: + self._fold_pre_session_descendants(credit) + # Root's terminal turn: clear tree root-pending AFTER final-turn + # spawns have registered their descendants (register runs inside + # _spawn_children_and_register_gates above), so the tree does not + # transiently read as drained. + if ( + self._session_tree_registry is not None + and credit.agent_depth == 0 + and credit.is_final_turn + ): + self._session_tree_registry.on_root_terminal( + credit.effective_root_correlation_id + ) return self._maybe_suspend_parent(credit) + def _fold_pre_session_descendants(self, credit: Credit) -> None: + """Attach live pre-session SPAWN children to a root instance's tree. + + Runs at the root's turn-0 return. Folds ONCE, into the first root + instance of the conversation (the live set is popped); children that + already terminated were removed from the set on their return and are + not counted. Multi-instance runs of the same template attach the + one-shot pre children to the first instance's tree only. + """ + live = self._pre_live_by_conv.pop(credit.conversation_id, None) + if not live: + return + root_corr = credit.effective_root_correlation_id + for child_corr in live: + self._child_root[child_corr] = root_corr + if self._session_tree_registry is not None: + self._session_tree_registry.register_descendants(root_corr, len(live)) + + def _resolve_pre_session_descendant(self, child_corr: str) -> None: + """Terminal accounting for a pre-session SPAWN child. + + Before its conversation's fold: drop it from the live set so the fold + does not count an already-finished child. After the fold: decrement its + tree's outstanding count. No-op for non-pre-session children. + """ + conv_id = self._pre_child_conv.pop(child_corr, None) + if conv_id is None: + return + live = self._pre_live_by_conv.get(conv_id) + if live is not None: + live.discard(child_corr) + if not live: + self._pre_live_by_conv.pop(conv_id, None) + root_corr = self._child_root.pop(child_corr, None) + if self._session_tree_registry is not None and root_corr is not None: + self._session_tree_registry.on_descendant_done(root_corr) + def _ensure_future_join( self, credit: Credit, @@ -200,21 +279,25 @@ def _ensure_future_join( pending = gates_for_parent.get(gated_idx) if pending is None: has_forks = False + has_branches = False if 0 <= gated_idx < len(parent_meta.turns): has_forks = bool( getattr(parent_meta.turns[gated_idx], "has_forks", False) ) + has_branches = bool(parent_meta.turns[gated_idx].branch_ids) pending = PendingBranchJoin( parent_x_correlation_id=parent_corr, parent_conversation_id=credit.conversation_id, parent_num_turns=len(parent_meta.turns), parent_agent_depth=credit.agent_depth, parent_parent_correlation_id=credit.parent_correlation_id, + parent_root_correlation_id=credit.root_correlation_id, gated_turn_index=gated_idx, parent_branch_mode=getattr( credit, "branch_mode", ConversationBranchMode.FORK ), parent_has_forks_on_gated_turn=has_forks, + parent_has_branches_on_gated_turn=has_branches, ) # Phase 3 fan-in seed: pre-populate every prereq_key declared # by the gated turn with an unregistered PrereqState so the @@ -355,6 +438,7 @@ async def on_child_leaf_reached(self, child_x_correlation_id: str) -> None: """Called when a child session reaches its final turn.""" if self._cleaning_up: return + self._resolve_pre_session_descendant(child_x_correlation_id) entries = self._child_to_join.get(child_x_correlation_id) if not entries: return @@ -370,6 +454,7 @@ async def on_child_stopped(self, child_x_correlation_id: str) -> None: """ if self._cleaning_up: return + self._resolve_pre_session_descendant(child_x_correlation_id) entries = self._child_to_join.get(child_x_correlation_id) if not entries: return @@ -383,6 +468,12 @@ async def _handle_child_done( self._child_to_join.pop(child_corr, None) # Every entry shares the same parent_correlation_id by construction. parent = entries[0].parent_correlation_id + # Registry descendant-done keys on the tree ROOT (not the parent), so + # depth>1 descendants still decrement the right tree. Retiring a drained + # tree here is finality bookkeeping only -- no slot is released. + child_root = self._child_root.pop(child_corr, None) + if self._session_tree_registry is not None and child_root is not None: + self._session_tree_registry.on_descendant_done(child_root) child_mode = self._child_modes.pop(child_corr, None) if ( child_mode == ConversationBranchMode.FORK @@ -421,6 +512,7 @@ async def on_child_errored(self, child_x_correlation_id: str) -> None: """ if self._cleaning_up: return + self._resolve_pre_session_descendant(child_x_correlation_id) entries = self._child_to_join.get(child_x_correlation_id) if not entries: return @@ -436,6 +528,9 @@ async def _handle_child_errored_fail_fast( parent = entries[0].parent_correlation_id errored_mode = self._child_modes.pop(child_corr, None) self._child_to_join.pop(child_corr, None) + errored_root = self._child_root.pop(child_corr, None) + if self._session_tree_registry is not None and errored_root is not None: + self._session_tree_registry.on_descendant_done(errored_root) # Collect all tracked children for this parent as potential orphans. orphans = [ @@ -459,6 +554,9 @@ async def _handle_child_errored_fail_fast( for orphan in orphans: self._child_to_join.pop(orphan, None) + orphan_root = self._child_root.pop(orphan, None) + if self._session_tree_registry is not None and orphan_root is not None: + self._session_tree_registry.on_descendant_done(orphan_root) orphan_mode = self._child_modes.pop(orphan, None) if ( orphan_mode == ConversationBranchMode.FORK @@ -508,9 +606,14 @@ def cleanup(self) -> None: self._abort_observer = None self._log_stats() self._log_leaks() + if self._session_tree_registry is not None: + self._session_tree_registry.release_all() self._active_joins.clear() self._future_joins.clear() self._child_to_join.clear() + self._child_root.clear() + self._pre_child_conv.clear() + self._pre_live_by_conv.clear() self._child_modes.clear() self._descendant_counts.clear() self._parent_locks.clear() diff --git a/src/aiperf/timing/concurrency.py b/src/aiperf/timing/concurrency.py index 4adcc2b7da..e2cad96837 100644 --- a/src/aiperf/timing/concurrency.py +++ b/src/aiperf/timing/concurrency.py @@ -538,6 +538,18 @@ def try_acquire_prefill_slot( return can_proceed_fn() return self._prefill_limiter.try_acquire(phase, can_proceed_fn) + @property + def prefill_limiting_enabled(self) -> bool: + """True iff prefill-concurrency limiting is active (a prefill limit is + configured for at least one phase). + + The callback handler consults this so a ``FirstToken`` arriving with no + prefill limiting in force (post-TTFT first-token anchoring) does not + inflate the prefill-released counter. Slot release itself is a no-op + when disabled, so only the accounting needs guarding. + """ + return self._prefill_limiter.enabled + def release_prefill_slot(self, phase: CreditPhase) -> None: """Release a prefill concurrency slot. diff --git a/src/aiperf/timing/config.py b/src/aiperf/timing/config.py index 86d7514c45..3c5d84bc51 100644 --- a/src/aiperf/timing/config.py +++ b/src/aiperf/timing/config.py @@ -8,12 +8,14 @@ from pydantic import ConfigDict, Field, model_validator -from aiperf.common.enums import CreditPhase +from aiperf.common.aiperf_logger import AIPerfLogger +from aiperf.common.enums import CacheBustTarget, CreditPhase from aiperf.common.models.base_models import AIPerfBaseModel from aiperf.config.dataset.defaults import InputDefaults from aiperf.config.sweep.adaptive import SLAFilter from aiperf.plugin.enums import ( ArrivalPattern, + DatasetSamplingStrategy, PhaseType, TimingMode, URLSelectionStrategy, @@ -29,6 +31,8 @@ from aiperf.config.phases import PhaseConfig from aiperf.config.resolution.plan import BenchmarkRun +_logger = AIPerfLogger(__name__) + # Map ``PhaseType`` values onto the ``ArrivalPattern`` values consumed by the # timing strategies. Concurrency / fixed_schedule phases don't use an arrival @@ -58,9 +62,9 @@ def _phase_timing_mode(phase: PhaseConfig) -> TimingMode: class TimingConfig(AIPerfBaseModel): """Configuration for TimingManager and timing strategies. - Controls timing mode (REQUEST_RATE, FIXED_SCHEDULE, or USER_CENTRIC_RATE), - rate/concurrency settings, warmup/profiling phase stop conditions, and - request cancellation behavior. + Controls timing mode (REQUEST_RATE, FIXED_SCHEDULE, USER_CENTRIC_RATE, + ADAPTIVE_SCALE, or GRAPH_IR), rate/concurrency settings, warmup/profiling + phase stop conditions, and request cancellation behavior. """ model_config = ConfigDict(frozen=True) @@ -95,13 +99,181 @@ def from_run(cls, run: BenchmarkRun) -> TimingConfig: """ cfg = run.cfg + # Resolved dataset-selection policy (populated by the post-scenario + # graph-dispatch resolver for graph runs; None otherwise). Carried onto + # each CreditPhaseConfig so ``GraphIRReplayStrategy`` can consume them; + # non-graph phases just pass the Nones through. + dataset_sampling_strategy = run.resolved.dataset_sampling_strategy + allow_dataset_wrap = run.resolved.allow_dataset_wrap + + # Weka graph IR workloads replay a conversation DAG via + # ``GraphIRReplayStrategy`` (TimingMode.GRAPH_IR), not the linear + # per-turn timing modes. Read the memoized single-source resolution so + # BOTH the warmup and profiling phases select the graph strategy; + # non-graph workloads are unchanged. An explicit, graph-incompatible + # ``--custom-dataset-type`` takes PRECEDENCE over the detection -- a + # user who pinned a different loader is never silently rerouted to the + # graph pipeline. The veto composes OVER the accessor (behavior this + # module alone owns), it is not absorbed into it. + from aiperf.dataset.graph.workload_detect import ( + _resolve_graph_max_context, + _resolve_graph_num_entries, + resolve_graph_workload, + ) + + is_graph = resolve_graph_workload( + run + ) is not None and not _explicit_non_graph_format(run) + + # Resolved per-trace-instance cache-bust target (scenario-auto-filled on + # ``endpoint.cache_bust`` before this runs). Threaded onto each graph + # CreditPhaseConfig so ``GraphIRReplayStrategy``'s dispatch duplication + # report can decide whether recycle duplication is safe. None for + # non-graph runs (the strategy is never built there). + cache_bust = cfg.endpoint.cache_bust if is_graph else None + + # Resolved graph-plane corpus-selection caps, threaded onto each graph + # CreditPhaseConfig so ``GraphIRReplayStrategy``'s wrap-guard can phrase + # an over-subscription shortfall as CAPPED (a knob shrank the corpus) + # rather than EXHAUSTED. REUSE the build plane's OWN resolvers so the + # dispatch-side caps match the build-side exactly (single source of + # truth). None for non-graph runs. + num_dataset_entries = _resolve_graph_num_entries(run) if is_graph else None + max_context_length = _resolve_graph_max_context(run) if is_graph else None + + # Resolved t* snapshot window + phase-start burst mode, threaded from + # the run config (--trajectory-start-min/max-ratio, scenario-auto- + # applied when unset; --burst-phase-starts) onto each graph + # CreditPhaseConfig so the strategy samples the same window the + # auto-warmup decision used. None for non-graph runs. + trajectory_start_min_ratio = ( + (cfg.trajectory_start_min_ratio or 0.0) if is_graph else None + ) + trajectory_start_max_ratio = ( + (cfg.trajectory_start_max_ratio or 0.0) if is_graph else None + ) + burst_phase_starts = cfg.burst_phase_starts if is_graph else None + + # Run seed for deterministic per-trace t* sampling (salted per trace/ + # lane inside the strategy) -- the SAME --random-seed that seeds + # synthesized content, so one seed reproduces the whole run. + random_seed = run.random_seed + + # Extended (cache-pressure) warmup duration, per-run config from + # --agentic-cache-warmup-duration. None = no pressure stage. + cache_pressure_duration = ( + cfg.agentic_cache_warmup_duration if is_graph else None + ) + artifact_dir = cfg.artifacts.dir configs: list[CreditPhaseConfig] = [] - for phase in cfg.get_warmup_phases(): - configs.append(_build_warmup_config(phase, artifact_dir=artifact_dir)) - for phase in cfg.get_profiling_phases(): - configs.append(_build_profiling_config(phase, artifact_dir=artifact_dir)) + warmup_phases = list(cfg.get_warmup_phases()) + profiling_phases = list(cfg.get_profiling_phases()) + # Explicit graph-incompatible phase choices take precedence over the + # detection (same rule as --custom-dataset-type above): reject loudly + # instead of silently rerouting. Checked BEFORE the pressure supersede + # so a rate-typed user warmup is rejected, not absorbed. + if is_graph: + _reject_graph_incompatible_phases(warmup_phases, profiling_phases) + # Pressure-mode graph warmups are MODE-OWNED (agentx parity: the + # agentic warmup pins expected_duration_sec=None and its own grace + # regardless of user warmup settings). A user warmup phase here would + # bound priming+pressure combined and its count caps would starve the + # pressure budget, so supersede it loudly with the auto shape -- but + # carry an EXPLICIT user grace through (agentx honors it verbatim). + graph_pressure_user_grace: float | None = None + if is_graph and cache_pressure_duration is not None and warmup_phases: + graph_pressure_user_grace = next( + ( + p.grace_period + for p in warmup_phases + if getattr(p, "grace_period", None) is not None + ), + None, + ) + _logger.notice( + "graph cache-pressure warmup supersedes the user-configured " + f"warmup phase(s) ({len(warmup_phases)}): boundary priming + " + "the pressure stage own the warmup shape (duration=None, " + "grace=user grace if set, else min(pressure duration, cap)); " + "profiling phases are unchanged" + ) + warmup_phases = [] + for phase in warmup_phases: + configs.append( + _build_warmup_config( + phase, + is_graph=is_graph, + artifact_dir=artifact_dir, + dataset_sampling_strategy=dataset_sampling_strategy, + allow_dataset_wrap=allow_dataset_wrap, + cache_bust=cache_bust, + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + trajectory_start_min_ratio=trajectory_start_min_ratio, + trajectory_start_max_ratio=trajectory_start_max_ratio, + burst_phase_starts=burst_phase_starts, + random_seed=random_seed, + cache_pressure_duration=cache_pressure_duration, + ) + ) + # AgentX parity (timing/config.py::_build_warmup_config graph-replay + # branch): a t*-snapshot weka graph run ALWAYS runs a WARMUP phase to + # prime each live chain's pre-t* boundary turn into the server KV cache, + # so the profiled (at/after-t*) turns measure a warm cache -- NOT a cold + # start. AgentX auto-creates that warmup unconditionally for the agentic + # replay mode; our graph path only built one when the user passed an + # explicit --warmup-* trigger, so a default t*>0 run started PROFILING + # cold (no warmup phase). Inject the auto-warmup when the t* window is + # active (trajectory_start_max_ratio > 0) and the user supplied no explicit warmup + # phase. With t*=0 (full native replay) there is no pre-t* prefix to + # prime, so rewrite_for_warmup returns an empty graph and the phase + # finalizes immediately -- harmless -- but we still skip it to keep the + # full-replay phase list byte-identical (one PROFILING phase). + # The pressure stage (extended warmup) also requires the phase even at t*=0. + if ( + is_graph + and not warmup_phases + and ( + _graph_tstar_active(trajectory_start_max_ratio) + or cache_pressure_duration is not None + ) + ): + auto = _build_graph_auto_warmup_config( + profiling_phases, + user_grace=graph_pressure_user_grace, + dataset_sampling_strategy=dataset_sampling_strategy, + allow_dataset_wrap=allow_dataset_wrap, + cache_bust=cache_bust, + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + trajectory_start_min_ratio=trajectory_start_min_ratio, + trajectory_start_max_ratio=trajectory_start_max_ratio, + burst_phase_starts=burst_phase_starts, + random_seed=random_seed, + cache_pressure_duration=cache_pressure_duration, + ) + if auto is not None: + configs.append(auto) + for phase in profiling_phases: + configs.append( + _build_profiling_config( + phase, + is_graph=is_graph, + artifact_dir=artifact_dir, + dataset_sampling_strategy=dataset_sampling_strategy, + allow_dataset_wrap=allow_dataset_wrap, + cache_bust=cache_bust, + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + trajectory_start_min_ratio=trajectory_start_min_ratio, + trajectory_start_max_ratio=trajectory_start_max_ratio, + burst_phase_starts=burst_phase_starts, + random_seed=random_seed, + cache_pressure_duration=cache_pressure_duration, + ) + ) cancellation_config: RequestCancellationConfig = RequestCancellationConfig() for phase in cfg.get_profiling_phases(): @@ -240,6 +412,83 @@ class CreditPhaseConfig(AIPerfBaseModel): default_factory=AdaptiveTimingConfig, description="Adaptive scale timing settings.", ) + dataset_sampling_strategy: DatasetSamplingStrategy | None = Field( + default=None, + description="Resolved run-level dataset sampling strategy for graph " + "workloads, carried from `run.resolved.dataset_sampling_strategy` so the " + "GraphIRReplayStrategy can consume it. None for non-graph phases and " + "until resolution derives it.", + ) + allow_dataset_wrap: bool | None = Field( + default=None, + description="Resolved graph-plane dataset-wrap policy, carried from " + "`run.resolved.allow_dataset_wrap` so the GraphIRReplayStrategy can " + "consume it. None for non-graph phases and until resolution derives it.", + ) + cache_bust: CacheBustTarget | None = Field( + default=None, + description="Resolved per-trace-instance cache-bust target, carried from " + "`endpoint.cache_bust` so the GraphIRReplayStrategy's dispatch " + "duplication report can decide whether recycle duplication is safe " + "(cache-bust ON) or warns (OFF). None for non-graph phases.", + ) + num_dataset_entries: int | None = Field( + default=None, + ge=1, + description="Resolved explicit --num-dataset-entries corpus cap for graph " + "workloads (the run's default-dataset `entries`), carried so the " + "GraphIRReplayStrategy wrap-guard phrases an over-subscription shortfall " + "as CAPPED rather than EXHAUSTED. None for non-graph phases and when unset.", + ) + max_context_length: int | None = Field( + default=None, + ge=1, + description="Resolved --max-context-length per-trace context cap for graph " + "workloads (`synthesis.max_context_length`), carried so the " + "GraphIRReplayStrategy wrap-guard phrases an over-subscription shortfall " + "as CAPPED rather than EXHAUSTED. None for non-graph phases and when unset.", + ) + trajectory_start_min_ratio: float | None = Field( + default=None, + ge=0.0, + le=1.0, + description="Resolved t* snapshot-window lower bound for graph " + "workloads, carried from `cfg.trajectory_start_min_ratio` " + "(`--trajectory-start-min-ratio`, scenario-auto-applied when unset) so " + "the GraphIRReplayStrategy samples the same window the auto-warmup " + "decision used. None for non-graph phases; unset resolves to 0.0.", + ) + trajectory_start_max_ratio: float | None = Field( + default=None, + ge=0.0, + le=1.0, + description="Resolved t* snapshot-window upper bound for graph " + "workloads, carried from `cfg.trajectory_start_max_ratio` " + "(`--trajectory-start-max-ratio`, scenario-auto-applied when unset). " + "None for non-graph phases; unset resolves to 0.0 (window OFF).", + ) + burst_phase_starts: bool | None = Field( + default=None, + description="Resolved --burst-phase-starts phase-start dispatch mode " + "for graph workloads, carried from `cfg.burst_phase_starts`. " + "None for non-graph phases.", + ) + random_seed: int | None = Field( + default=None, + ge=0, + description="The run's resolved --random-seed, threaded so the graph " + "strategy's t* sampling derives from the SAME seed as synthesized " + "content (one seed reproduces the whole run; sweep cells decorrelate " + "via the orchestrator's per-variation seed derivation).", + ) + cache_pressure_duration: float | None = Field( + default=None, + gt=0, + description="Extended (cache-pressure) warmup duration in seconds, " + "carried from `cfg.agentic_cache_warmup_duration` " + "(--agentic-cache-warmup-duration). None = no pressure stage; also " + "None for non-graph phases.", + ) @model_validator(mode="before") @classmethod @@ -331,6 +580,107 @@ def _ramp_duration(ramp: object | None) -> float | None: return getattr(ramp, "duration", None) +def _explicit_non_graph_format(run: BenchmarkRun) -> bool: + """True when the run pins an explicit, graph-incompatible dataset format. + + The weka graph workload is auto-detected by a file-content sniff + (``workload_detect.resolve_graph_workload``); there is no ``weka`` value in + ``DatasetFormat``. A bare ``--input-file weka.json`` run resolves to the + default ``DatasetFormat.SINGLE_TURN``, so we treat SINGLE_TURN as "unset -- + let the sniff decide". Any OTHER resolved format means the user explicitly + pinned a different loader (``--custom-dataset-type multi_turn`` / + ``mooncake_trace`` / ...), which must take PRECEDENCE over the sniff -- the + run is NOT silently rerouted to the graph IR pipeline. Returns False for any + non-file dataset (no ``format`` attribute) so synthetic / public runs are + unaffected. + """ + from aiperf.common.enums import DatasetFormat + + try: + dataset = run.cfg.get_default_dataset() + except (IndexError, AttributeError): + return False + fmt = getattr(dataset, "format", None) + if fmt is None: + return False + return fmt != DatasetFormat.SINGLE_TURN + + +def _reject_graph_incompatible_phases( + warmup_phases: list[PhaseConfig], profiling_phases: list[PhaseConfig] +) -> None: + """Reject phase configs whose explicit pacing a graph replay would discard. + + The recorded graph replay (TimingMode.GRAPH_IR) owns pacing and + concurrency, so rerouting these phases to GRAPH_IR would + silently discard the user's explicit choice: + + * ``adaptive_scale`` profiling phases drive their own concurrency ladder. + * Rate-controlled and fixed-schedule phase TYPES encode explicit user + pacing (``--request-rate`` / ``--user-centric-rate`` / + ``--fixed-schedule``). + """ + if any(getattr(phase, "adaptive_scale", False) for phase in profiling_phases): + raise ValueError( + "adaptive_scale is not supported for graph workloads: the " + "recorded graph replay (TimingMode.GRAPH_IR) owns pacing and " + "concurrency. Remove adaptive_scale from the phase config, or " + "pin a non-graph loader with --custom-dataset-type to run " + "this input through the linear pipeline." + ) + for phase in (*warmup_phases, *profiling_phases): + if phase.type != PhaseType.CONCURRENCY: + raise ValueError( + f"phase '{phase.name}' (type={phase.type}) is not " + "supported for graph workloads: the recorded graph " + "replay (TimingMode.GRAPH_IR) owns pacing, so " + "rate-controlled arrivals (--request-rate / " + "--user-centric-rate) and --fixed-schedule " + "timestamps would be silently discarded. Remove the " + "rate/schedule options (--concurrency bounds the " + "replay lanes), or pin a non-graph loader with " + "--custom-dataset-type to run this input through " + "the linear pipeline." + ) + + +def resolve_graph_content_seed(run: BenchmarkRun) -> int | None: + """Return the run's seed for Weka content synthesis -- the AIPerf seed. + + Just ``run.random_seed`` (``--random-seed``), the same seed schedule / + topology / t* derive from. The DatasetManager is the only parser; the + TimingManager ingests the graph_meta sidecar from the graph-typed dataset + broadcast. Resolving from the run config alone keeps every parse of the same + run (in-process or spawn-started pool worker) byte-identical. ``None`` (no + ``--random-seed``) keeps the ambient global RNG -- there is no weka-specific + seed fallback. + """ + return run.random_seed + + +def resolve_graph_content_tokenizer(run: BenchmarkRun) -> str: + """Resolve the tokenizer the Weka content synthesizer must use. + + The synthesized message content (block + partial-tail token IDs decoded to + wire text) is only valid if it is decoded with the SAME tokenizer the run + dispatches and token-counts against. So this delegates to the ONE existing + resolver -- ``TokenizerConfig.get_tokenizer_name_for_model`` -- the exact + call :meth:`DatasetManager._configure_tokenizer` uses to load the dispatch + tokenizer (CLI-resolved name, else explicit ``--tokenizer``, else the model + name). No separate resolution logic lives here. + + The DatasetManager is the only parser; the TimingManager ingests the + graph_meta sidecar from the graph-typed dataset broadcast. Resolving from + the run config alone keeps every parse of the same run (in-process or + spawn-started pool worker) byte-identical (the same content contract as + :func:`resolve_graph_content_seed`). + """ + cfg = run.cfg + model_names = cfg.get_model_names() + model = model_names[0] if model_names else "" + return cfg.tokenizer.get_tokenizer_name_for_model(model) + + def _phase_request_rate(phase: PhaseConfig) -> float | None: """Return the configured request rate for a phase, if any.""" # Lazy import: aiperf.config.phases is only a TYPE_CHECKING import here. @@ -345,7 +695,20 @@ def _phase_arrival_pattern(phase: PhaseConfig) -> ArrivalPattern: def _build_warmup_config( - phase: PhaseConfig, *, artifact_dir: Path | None = None + phase: PhaseConfig, + *, + is_graph: bool = False, + artifact_dir: Path | None = None, + dataset_sampling_strategy: DatasetSamplingStrategy | None = None, + allow_dataset_wrap: bool | None = None, + cache_bust: CacheBustTarget | None = None, + num_dataset_entries: int | None = None, + max_context_length: int | None = None, + trajectory_start_min_ratio: float | None = None, + trajectory_start_max_ratio: float | None = None, + burst_phase_starts: bool | None = None, + random_seed: int | None = None, + cache_pressure_duration: float | None = None, ) -> CreditPhaseConfig: """Build a warmup CreditPhaseConfig from a warmup PhaseConfig. @@ -356,6 +719,13 @@ def _build_warmup_config( forever for in-flight requests). This differs from the CreditPhaseConfig field default of None (disabled) because warmup should always complete all in-flight requests before transitioning to profiling. + + ``is_graph`` forces ``TimingMode.GRAPH_IR`` so a weka graph workload's WARMUP + phase runs the graph strategy (re-seeding the boundary prefix via warmup + materialization over the graph store's profiling bytes) rather than the linear + ``RequestRateStrategy`` over zero-turn ``Conversation`` stubs -- which would + send nothing useful (the per-node payloads live in the graph store mmap, not + the stub conversations). Mirrors ``_build_profiling_config``'s ``is_graph``. """ grace_period = phase.grace_period if grace_period is None: @@ -363,8 +733,7 @@ def _build_warmup_config( return CreditPhaseConfig( phase=CreditPhase.WARMUP, - # Warmup phase is always request rate timing mode - timing_mode=TimingMode.REQUEST_RATE, + timing_mode=TimingMode.GRAPH_IR if is_graph else TimingMode.REQUEST_RATE, total_expected_requests=phase.requests, expected_duration_sec=phase.duration, expected_num_sessions=phase.sessions, @@ -381,21 +750,148 @@ def _build_warmup_config( getattr(phase, "rate_ramp", None) ), artifact_dir=artifact_dir, + dataset_sampling_strategy=dataset_sampling_strategy, + allow_dataset_wrap=allow_dataset_wrap, + cache_bust=cache_bust, + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + trajectory_start_min_ratio=trajectory_start_min_ratio, + trajectory_start_max_ratio=trajectory_start_max_ratio, + burst_phase_starts=burst_phase_starts, + random_seed=random_seed, + cache_pressure_duration=cache_pressure_duration, + ) + + +def _graph_tstar_active(trajectory_start_max_ratio: float | None) -> bool: + """True iff the graph t* snapshot window is engaged (max ratio > 0). + + Takes the SAME resolved config value ``from_run`` threads onto each graph + ``CreditPhaseConfig``, so the auto-warmup decision and the strategy's t* + sampling agree: a positive upper ratio means at least some traces sample + ``t* > 0`` and therefore have a pre-``t*`` prefix worth priming. ``[0, 0]`` + (full native replay) leaves it inactive -- no warmup needed. + """ + return (trajectory_start_max_ratio or 0.0) > 0.0 + + +def _graph_pressure_grace_sec(user_grace: float | None, duration: float) -> float: + """Drain grace for a pressure-mode warmup. + + AgentX parity (`_agentic_warmup_grace_period`, its timing/config.py:226-240): + an EXPLICIT user warmup grace is honored verbatim (the operator's escape + hatch when a healthy drain outlives the pressure duration -- e.g. a + 45s prefill in flight at a 30s deadline); otherwise min(duration, cap) + bounds the drain so a wedged or lost return cannot hang the run. We do + not port agentx's benchmark-grace floor branch. The strategy's stash + completeness gate converts a drain that force-completes with unreturned + credits into a safe handoff skip. Callers gate on a set + pressure duration first. + """ + from aiperf.common.environment import Environment + + if user_grace is not None: + return user_grace + + cap = Environment.GRAPH.PRESSURE_DRAIN_GRACE_CAP + return min(duration, cap) + + +def _build_graph_auto_warmup_config( + profiling_phases: list[PhaseConfig], + *, + user_grace: float | None = None, + dataset_sampling_strategy: DatasetSamplingStrategy | None = None, + allow_dataset_wrap: bool | None = None, + cache_bust: CacheBustTarget | None = None, + num_dataset_entries: int | None = None, + max_context_length: int | None = None, + trajectory_start_min_ratio: float | None = None, + trajectory_start_max_ratio: float | None = None, + burst_phase_starts: bool | None = None, + random_seed: int | None = None, + cache_pressure_duration: float | None = None, +) -> CreditPhaseConfig | None: + """Auto-build the t*-snapshot WARMUP phase for a weka graph-IR run. + + AgentX parity: the graph-replay scenario path always prepends a WARMUP phase priming the + boundary (``k_i-1``) turn of every chain live at ``t*`` (one priming credit + per live chain), dispatched as a ``CONCURRENCY_BURST`` with an infinite grace + period so the warmup barrier holds until every priming credit returns. The + GraphIRReplayStrategy owns warmup completion (its warmup phase variant runs + ``aiperf.timing.strategies.graph_ir_replay.rewrite_for_warmup``, a flat + START-rooted graph firing exactly the live chains' boundary turns, then + drains), so this carries no stop-condition counts -- only the concurrency + (inherited from the profiling phase so the warmup primes at the same width + the profiling phase will run). + + Grace: the boundary-priming default is an infinite grace (the barrier holds + until every priming credit returns). When the cache-pressure stage is active + (a set pressure duration), the drain is instead bounded by + :func:`_graph_pressure_grace_sec` so a wedged or lost pressure return cannot + hang the run; ``user_grace`` (a graph run's explicit ``--warmup-grace-period`` + carried through the mode-owned supersede) is honored verbatim there. + + Returns ``None`` when there is no profiling phase to inherit concurrency from + (degenerate config); the caller then simply skips the auto-warmup. + """ + if not profiling_phases: + return None + base = profiling_phases[0] + return CreditPhaseConfig( + phase=CreditPhase.WARMUP, + timing_mode=TimingMode.GRAPH_IR, + concurrency=base.concurrency, + prefill_concurrency=base.prefill_concurrency, + arrival_pattern=ArrivalPattern.CONCURRENCY_BURST, + seamless=False, + grace_period_sec=( + _graph_pressure_grace_sec(user_grace, cache_pressure_duration) + if cache_pressure_duration is not None + else float("inf") + ), + dataset_sampling_strategy=dataset_sampling_strategy, + allow_dataset_wrap=allow_dataset_wrap, + cache_bust=cache_bust, + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + trajectory_start_min_ratio=trajectory_start_min_ratio, + trajectory_start_max_ratio=trajectory_start_max_ratio, + burst_phase_starts=burst_phase_starts, + random_seed=random_seed, + cache_pressure_duration=cache_pressure_duration, ) def _build_profiling_config( - phase: PhaseConfig, *, artifact_dir: Path | None = None + phase: PhaseConfig, + *, + is_graph: bool = False, + artifact_dir: Path | None = None, + dataset_sampling_strategy: DatasetSamplingStrategy | None = None, + allow_dataset_wrap: bool | None = None, + cache_bust: CacheBustTarget | None = None, + num_dataset_entries: int | None = None, + max_context_length: int | None = None, + trajectory_start_min_ratio: float | None = None, + trajectory_start_max_ratio: float | None = None, + burst_phase_starts: bool | None = None, + random_seed: int | None = None, + cache_pressure_duration: float | None = None, ) -> CreditPhaseConfig: """Build a profiling CreditPhaseConfig from a profiling PhaseConfig. Main benchmark phase where all performance metrics are collected. Grace period allows in-flight requests to complete after the stop condition is met, ensuring metrics include requests that were sent before the deadline. + + ``is_graph`` forces ``TimingMode.GRAPH_IR`` so a weka graph workload selects + ``GraphIRReplayStrategy`` regardless of the phase's ``type``. """ + timing_mode = TimingMode.GRAPH_IR if is_graph else _phase_timing_mode(phase) return CreditPhaseConfig( phase=CreditPhase.PROFILING, - timing_mode=_phase_timing_mode(phase), + timing_mode=timing_mode, expected_duration_sec=phase.duration, total_expected_requests=phase.requests, expected_num_sessions=phase.sessions, @@ -444,4 +940,14 @@ def _build_profiling_config( phase, "adaptive_min_completed_requests", 1 ), adaptive_sla_filters=tuple(getattr(phase, "sla", ()) or ()), + dataset_sampling_strategy=dataset_sampling_strategy, + allow_dataset_wrap=allow_dataset_wrap, + cache_bust=cache_bust, + num_dataset_entries=num_dataset_entries, + max_context_length=max_context_length, + trajectory_start_min_ratio=trajectory_start_min_ratio, + trajectory_start_max_ratio=trajectory_start_max_ratio, + burst_phase_starts=burst_phase_starts, + random_seed=random_seed, + cache_pressure_duration=cache_pressure_duration, ) diff --git a/src/aiperf/timing/conversation_source.py b/src/aiperf/timing/conversation_source.py index b301449a35..7e2f00b64a 100644 --- a/src/aiperf/timing/conversation_source.py +++ b/src/aiperf/timing/conversation_source.py @@ -15,6 +15,8 @@ session route to the same worker. """ +from __future__ import annotations + import uuid from dataclasses import dataclass @@ -46,6 +48,10 @@ class SampledSession: parent's accumulated message history and pins to the same worker; SPAWN starts with a fresh context. Ignored when parent_correlation_id is None. + root_correlation_id: The x_correlation_id of the depth-0 root of this + session's tree. None on a depth-0 root (it is its own tree root); + inherited by children/subchildren so per-tree finality accounting + can key on ``effective_root_correlation_id``. """ conversation_id: str @@ -53,6 +59,7 @@ class SampledSession: x_correlation_id: str agent_depth: int = 0 parent_correlation_id: str | None = None + root_correlation_id: str | None = None branch_mode: ConversationBranchMode = ConversationBranchMode.FORK @property @@ -65,6 +72,16 @@ def routing_key(self) -> str: """ return self.parent_correlation_id or self.x_correlation_id + @property + def effective_root_correlation_id(self) -> str: + """Tree root id, defaulting to this session's own x_correlation_id. + + A root session is its own tree root; a child/subchild inherits the + depth-0 root's id, set by the spawning code so per-tree finality + accounting can key on it. + """ + return self.root_correlation_id or self.x_correlation_id + def build_first_turn(self, max_turns: int | None = None) -> TurnToSend: """Build first turn (turn_index=0) from sampled conversation. @@ -80,7 +97,11 @@ def build_first_turn(self, max_turns: int | None = None) -> TurnToSend: num_turns=max_turns or len(self.metadata.turns), agent_depth=self.agent_depth, parent_correlation_id=self.parent_correlation_id, + root_correlation_id=self.root_correlation_id, has_forks=first_meta.has_forks if first_meta is not None else False, + has_branches=bool(first_meta.branch_ids) + if first_meta is not None + else False, branch_mode=self.branch_mode, ) @@ -132,6 +153,7 @@ def start_branch_child( child_conversation_id: str, agent_depth: int, *, + root_correlation_id: str | None = None, branch_mode: ConversationBranchMode = ConversationBranchMode.FORK, ) -> SampledSession: """Build a SampledSession for a DAG child conversation. @@ -143,6 +165,11 @@ def start_branch_child( SPAWN-mode children start with a fresh context, but the sticky pin to the parent's correlation_id is preserved at this layer — routing freedom is enforced upstream by the orchestrator/router. + + ``root_correlation_id`` is the depth-0 root of the spawning parent's + tree; the child inherits it so all descendants of one root share a + single per-tree finality key. Defaults to ``parent_correlation_id`` + when not supplied (a depth-1 child's parent IS the root). """ metadata = self._metadata_lookup[child_conversation_id] return SampledSession( @@ -151,6 +178,7 @@ def start_branch_child( x_correlation_id=str(uuid.uuid4()), agent_depth=agent_depth, parent_correlation_id=parent_correlation_id, + root_correlation_id=root_correlation_id or parent_correlation_id, branch_mode=branch_mode, ) diff --git a/src/aiperf/timing/graph_channel.py b/src/aiperf/timing/graph_channel.py new file mode 100644 index 0000000000..93eb056431 --- /dev/null +++ b/src/aiperf/timing/graph_channel.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""First-class carrier for graph-phase state shared across phase runners.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from aiperf.dataset.graph.models import ParsedGraph + from aiperf.timing.graph_warmup_handoff import GraphWarmupHandoff + + +@dataclass(slots=True) +class GraphPhaseChannel: + """Graph-run state threaded orchestrator -> runner -> strategy. + + Replaces the former attribute smuggling on the generic + ``ConversationSource`` (``parsed_graph`` / ``graph_warmup_handoff``): + graph phases receive this typed channel instead of a conversation-shaped + object they never sample from. + """ + + parsed_graph: ParsedGraph + """The structural conversation DAG the graph strategies plan from.""" + + warmup_handoff: GraphWarmupHandoff | None = None + """Consume-once WARMUP -> PROFILING handoff; the WARMUP strategy stashes + it at teardown and the first PROFILING graph phase pops it.""" diff --git a/src/aiperf/timing/graph_ir_source.py b/src/aiperf/timing/graph_ir_source.py new file mode 100644 index 0000000000..bf084198bc --- /dev/null +++ b/src/aiperf/timing/graph_ir_source.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GraphIRConversationSource — per-trace t* selection. + +The schedule-plane source the ``GraphIRReplayStrategy`` consults to decide, per +weka trace, the wall-clock snapshot instant ``t*`` at/after which nodes are +profiled (pre-``t*`` firings are cache-priming warmup history). + +t* selection mirrors agentx ``TrajectorySource`` (``timing/trajectory_source.py``): +a per-trace ``t*`` is drawn uniformly from +``[start_min_ratio, start_max_ratio] * trace_duration`` with a deterministic, +LANE-salted RNG (``sha256(f"{seed}:{trace_id}:{lane_index}")``), so the same +``(seed, ratios, lane_index)`` reproduce the exact same t* on every run and +across both parse planes. AgentX production ALWAYS uses the lane-salted seed +(``_seed_for_trace_lane``, ``trajectory_source.py:410``); each replayed trace +INSTANCE is a lane, and the single-pass non-recycling case is lane ``0``. The +default ratios ``[0.0, 0.0]`` give ``t*=0`` == full native replay: no warmup +history, every node profiled (the working profiling path is unchanged). +""" + +from __future__ import annotations + +import hashlib +from collections.abc import Iterator +from dataclasses import dataclass + +import numpy as np + +from aiperf.dataset.graph.models import ParsedGraph, TraceRecord +from aiperf.graph.analysis import trace_duration_us + +__all__ = ["GraphIRConversationSource", "GraphTrace"] + + +@dataclass(slots=True, frozen=True) +class GraphTrace: + """One weka trace's snapshot plan: the sampled snapshot instant ``t*``. + + Attributes: + trace_id: The trace's root id. + parsed_graph: The shared ``ParsedGraph`` the trace belongs to (passed + through to the executor; t* reconstruction, when enabled, rewrites it + per trace via the trie snapshot chop). + t_star_us: The sampled snapshot instant in microseconds, kept as a float + (agentx parity: ``trajectory_source.py`` keeps t* a float ms -- no + integer truncation). ``0`` means full native replay (no warmup + history). + """ + + trace_id: str + parsed_graph: ParsedGraph + t_star_us: float + + +class GraphIRConversationSource: + """Yields a :class:`GraphTrace` (t* + node partition) per weka trace. + + Holds the shared ``ParsedGraph`` and the per-build node-ordinal catalog, and + samples a deterministic per-trace ``t*`` from the configured ratio window. + """ + + def __init__( + self, + *, + parsed: ParsedGraph, + start_min_ratio: float = 0.0, + start_max_ratio: float = 0.0, + random_seed: int = 0, + lane_index: int = 0, + ) -> None: + """Initialize the source. + + Args: + parsed: The built ``ParsedGraph`` whose ``traces`` this source plans. + start_min_ratio: Lower bound (fraction of trace duration) of the t* + sampling window. Default ``0.0``. A ratio of ``0.0`` selects + full native replay (t*=0, no warmup). + start_max_ratio: Upper bound of the t* window. Default ``0.0``. + Must be >= start_min_ratio. + random_seed: Base seed for per-trace t* sampling; salted with the + trace id AND ``lane_index`` so traces are uncorrelated yet + reproducible. + lane_index: The replay-lane index this source plans for. AgentX seeds + t* per ``(trace_id, lane_index)`` so the same trace recurring + across lanes decorrelates. The single-pass non-recycling case is + lane ``0`` (the default), matching AgentX's first/only lane. + """ + if start_min_ratio > start_max_ratio: + raise ValueError( + f"start_min_ratio ({start_min_ratio}) must be <= " + f"start_max_ratio ({start_max_ratio})" + ) + self._parsed = parsed + self._start_min_ratio = start_min_ratio + self._start_max_ratio = start_max_ratio + self._random_seed = random_seed + self._lane_index = lane_index + + def iter_traces(self) -> Iterator[GraphTrace]: + """Yield one :class:`GraphTrace` per trace in the parsed graph.""" + for trace in self._parsed.traces: + yield self._plan_trace(trace) + + def _plan_trace(self, trace: TraceRecord) -> GraphTrace: + """Sample ``t*`` for ``trace``.""" + t_star_us = self._sample_t_star(trace) + return GraphTrace( + trace_id=trace.id, + parsed_graph=self._parsed, + t_star_us=max(t_star_us, 0.0), + ) + + def _sample_t_star(self, trace: TraceRecord) -> float: + """Draw a deterministic per-trace ``t*`` in microseconds. + + Agentx parity: ``t* = uniform(lo, hi)`` over + ``[start_min_ratio, start_max_ratio] * trace_duration``, with a + LANE-salted RNG (``_seed_for_trace_lane`` -- AgentX production always + passes a concrete lane, ``trajectory_source.py:410``). ``lo == hi`` + (e.g. the default ``[0, 0]`` window or a zero-duration trace) collapses + to that exact instant with no draw. The draw is kept as a ``float`` + (no integer-microsecond truncation) to match AgentX, which keeps t* a + float ms. + """ + duration_us = trace_duration_us(self._parsed, trace) + if duration_us <= 0: + return 0.0 + lo = self._start_min_ratio * duration_us + hi = self._start_max_ratio * duration_us + if hi <= lo: + return float(lo) + rng = np.random.default_rng( + _seed_for_trace_lane(self._random_seed, trace.id, self._lane_index) + ) + return float(rng.uniform(lo, hi)) + + +def _seed_for_trace_lane(base_seed: int, trace_id: str, lane_index: int) -> int: + """Derive a per-trace-per-lane RNG seed (agentx ``_seed_for_trace_lane``). + + Agentx ``TrajectorySource._seed_for_trace_lane`` parity + (``trajectory_source.py:149``): per-(trace, lane) t* values must be + deterministic given ``base_seed`` yet decorrelated across both traces and + lanes, so we SHA-256 ``f"{base_seed}:{trace_id}:{lane_index}"`` and take the + low 8 bytes. AgentX production ALWAYS uses this lane variant + (``_build_trajectory_for_lane`` passes a concrete lane, TS:410); the plain + no-lane ``_seed_for_trace`` is unreachable in production. + """ + digest = hashlib.sha256(f"{base_seed}:{trace_id}:{lane_index}".encode()).digest() + return int.from_bytes(digest[:8], "big") diff --git a/src/aiperf/timing/graph_ir_trace_view.py b/src/aiperf/timing/graph_ir_trace_view.py new file mode 100644 index 0000000000..073df162d7 --- /dev/null +++ b/src/aiperf/timing/graph_ir_trace_view.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-trace single-graph projection for the weka graph-IR replay strategy. + +Extracted from ``GraphIRReplayStrategy`` (no behavior change) to keep that file +under the file-size ergonomics ceiling. The runtime ``TraceExecutor`` (and the +``rewrite_for_snapshot`` / ``rewrite_for_warmup`` / ``compute_snapshot`` paths it +drives) read ``parsed.graph`` directly, so every trace must be projected onto a +single-graph ``ParsedGraph`` view whose ``.graph`` is the trace's OWN resolved +topology before it reaches the executor. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import msgspec + +from aiperf.dataset.graph.models import resolve_trace_graph + +if TYPE_CHECKING: + from aiperf.dataset.graph.models import ParsedGraph, TraceRecord + +__all__ = ["parsed_for_trace"] + + +def parsed_for_trace(parsed: ParsedGraph, trace: TraceRecord) -> ParsedGraph: + """Project ``parsed`` onto a single-graph view whose ``.graph`` is ``trace``'s. + + Single-graph workloads (``trace.graph_ref is None``) return ``parsed`` + unchanged -- ``resolve_trace_graph`` already yields the shared ``parsed.graph``, + so the projection is a strict identity and the byte-unchanged single-file / + hand-authored path is untouched. + + Multi-graph workloads (a heterogeneous directory / HuggingFace corpus of weka + traces) keep each trace's own topology in ``parsed.graphs`` keyed by + ``trace.graph_ref``; ``parsed.graph`` is only the FIRST file's graph. This + returns a ``ParsedGraph`` whose ``.graph`` is this trace's resolved topology + (``parsed.graphs[trace.graph_ref]``) so the executor and the t* rewrites + operate on the CORRECT per-trace graph rather than the first file's graph. + The ``traces`` list is narrowed to this one trace so the rewrite's + matching-trace lookup resolves it; ``segment_pool`` and ``graphs`` ride along + via ``replace``. The pool rides along as the segment-store-IR marker + (``segment_pool is not None`` gates the executor's dispatch-failure + sentinel writes) and for debug + materialization via ``read_prompt_segment_ids`` -- NOT for inline node + prompts (trie-route ``LlmNode``s carry ``prompt=[]``; the worker materializes + from the mmap segment store). + + Without this projection a non-first trace runs the wrong topology -- the + first file's ``parsed.graph`` -- instead of its own resolved one. + """ + graph = resolve_trace_graph(parsed, trace) + if graph is parsed.graph: + return parsed + return msgspec.structs.replace(parsed, graph=graph, traces=[trace]) diff --git a/src/aiperf/timing/graph_warmup_handoff.py b/src/aiperf/timing/graph_warmup_handoff.py new file mode 100644 index 0000000000..5f8067c963 --- /dev/null +++ b/src/aiperf/timing/graph_warmup_handoff.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Extended-warmup handoff payload -- the WARMUP -> PROFILING re-cut channel. + +The graph-native counterpart of AgentX v1.0's ``finalize_phase`` trajectory +replacement (ported from the retired agentic-replay plane): the WARMUP ``GraphIRReplayStrategy`` +builds one :class:`GraphWarmupHandoff` at ``teardown_phase`` (after every +warmup credit return has landed) and stashes it on the SHARED +``GraphPhaseChannel``; ``PhaseRunner._build_graph_ir_strategy`` pops it +(consume-once) into the PROFILING strategy, which resumes each lane at its +recorded execution frontier via ``chop_trie_at_frontier``. + +Everything deterministic (the per-(trace, lane) t* plan) is NOT carried here +-- both phases re-derive it from the seeded sampler. The handoff carries only +what determinism cannot reproduce: which template each lane was mid-flight on +at drain, which nodes actually executed, and when their returns landed. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +__all__ = ["GraphWarmupHandoff", "LaneHandoff"] + + +@dataclass(slots=True, frozen=True) +class LaneHandoff: + """One lane's live-at-drain execution state.""" + + # Template trace id the lane was mid-flight on at drain (may differ from + # the lane's pass-0 template when the pressure stage recycled the lane). + template_trace_id: str + # The live pressure instance id at drain (e.g. ``t-1#0.p2``). The profiling + # resume reuses it verbatim as the resumed instance's id so the per-instance + # cache-bust marker (digest of ``credit.trace_id``; see + # ``build_trace_instance_marker``) is continuous across the handoff and the + # KV built during pressure transfers instead of cold-prefilling behind a + # fresh ``.0`` marker. + instance_id: str + # The instance's t* (lane-salted plan for the pressure pass-0 instance; + # 0.0 for recycled full-replay instances). Pre-t* nodes are warmup history + # and are dropped by the frontier chop alongside the executed set. + t_star_us: float + # Node ids of the live instance that dispatched AND returned during + # warmup/pressure -- the server holds their KV; profiling must not refire. + executed_node_ids: frozenset[str] + # Monotonic return wall times (microseconds, strategy ledger clock) used + # to compute residual re-root delays; includes the lane's boundary-priming + # returns merged in for pass-0 instances. + return_wall_us: dict[str, float] + + +@dataclass(slots=True, frozen=True) +class GraphWarmupHandoff: + """The full warmup -> profiling handoff, one entry per live lane.""" + + # Lanes live at drain. Lanes absent here (template completed exactly at + # drain, or lane index beyond the pressure lane count) resume the normal + # pass-0 t* path in profiling. + lanes: dict[int, LaneHandoff] + # Drain-end instant on the same monotonic clock as the return walls, + # stamped at warmup teardown (after all returns landed). + drain_end_wall_us: float + # Next corpus draw index after the pressure stage's last recycle draw. + # Profiling's BOUNDED recycle loop continues the wrap from here so freed + # lanes don't re-serve templates the pressure stage just replayed (agentx + # shares ONE sampler across pressure / handoff / profiling draws). Single- + # pass profiling (no stop conditions) deliberately ignores it -- full-corpus + # coverage takes precedence over cursor continuity there. + corpus_cursor: int + # Number of pressure lanes (0..K-1) the warmup ran. A lane below this + # count with NO entry in ``lanes`` completed at drain: profiling must + # fresh-start it (next cursor template, full t*=0 replay) instead of + # re-running a t* resume the pressure stage already executed -- agentx + # ``_build_handoff_trajectories`` empty-lane parity. + pressure_lane_count: int diff --git a/src/aiperf/timing/manager.py b/src/aiperf/timing/manager.py index 8a84d5de61..6985577d10 100644 --- a/src/aiperf/timing/manager.py +++ b/src/aiperf/timing/manager.py @@ -5,8 +5,10 @@ from __future__ import annotations import asyncio +from pathlib import Path from typing import TYPE_CHECKING +from aiperf.common.aiperf_logger import AIPerfLogger from aiperf.common.base_component_service import BaseComponentService from aiperf.common.enums import CommandType, MessageType from aiperf.common.environment import Environment @@ -25,6 +27,7 @@ ProfileConfigureCommand, ) from aiperf.common.models import DatasetMetadata +from aiperf.common.models.dataset_models import GraphSegmentClientMetadata from aiperf.credit.sticky_router import StickyCreditRouter from aiperf.timing.config import TimingConfig from aiperf.timing.phase.publisher import PhasePublisher @@ -32,6 +35,12 @@ if TYPE_CHECKING: from aiperf.config.resolution.plan import BenchmarkRun + from aiperf.dataset.graph.models import ParsedGraph + +# Module logger for run-once configure-time advisories (a stable logger name the +# tests capture via caplog; the service's per-instance ``self.warning`` logger +# name is derived from the runtime service id and is not test-stable). +_logger = AIPerfLogger(__name__) class TimingManager(BaseComponentService): @@ -39,7 +48,8 @@ class TimingManager(BaseComponentService): Central Service for the credit system. Creates a PhaseOrchestrator which internally instantiates the appropriate TimingMode based on mode - (REQUEST_RATE, FIXED_SCHEDULE, or USER_CENTRIC_RATE). + (REQUEST_RATE, FIXED_SCHEDULE, USER_CENTRIC_RATE, ADAPTIVE_SCALE, or + GRAPH_IR). Handles commands: PROFILE_CONFIGURE (create orchestrator), PROFILE_START (begin credit issuance), @@ -69,6 +79,7 @@ def __init__( self._dataset_failed_event = asyncio.Event() self._dataset_failure_reason: str | None = None self._dataset_metadata: DatasetMetadata | None = None + self._graph_client_metadata: GraphSegmentClientMetadata | None = None # StickyCreditRouter handles everything: routing, sending, returns, # worker lifecycle. Created early to handle worker connections @@ -88,12 +99,19 @@ async def _on_dataset_configured_notification( ) -> None: """Store dataset metadata and signal configuration ready.""" self.debug( - lambda: f"Received dataset configured notification: " - f"{len(message.metadata.conversations)} conversations, " - f"{message.metadata.sampling_strategy.value} sampling strategy" + lambda: ( + f"Received dataset configured notification: " + f"{len(message.metadata.conversations)} conversations, " + f"{message.metadata.sampling_strategy.value} sampling strategy" + ) ) self._dataset_metadata = message.metadata + self._graph_client_metadata = ( + message.client_metadata + if isinstance(message.client_metadata, GraphSegmentClientMetadata) + else None + ) self._dataset_configured_event.set() @on_message(MessageType.DATASET_CONFIGURATION_FAILED) @@ -132,15 +150,178 @@ async def _profile_configure_command( self.debug(f"Configuring phase orchestrator for {self.service_id}") + # Graph IR workloads need the structural conversation DAG handed to + # the graph timing strategy. Load the mandatory graph_meta sidecar + # from the path the graph-typed dataset broadcast advertised (hard + # fail if unadvertised or unloadable); runs after + # _wait_for_dataset_or_failure, so the broadcast has been consumed. + parsed_graph = self._load_graph_sidecar() + if parsed_graph is not None: + self._advise_non_streaming_first_token_sources(parsed_graph) + # Create orchestrator that executes phases self._phase_orchestrator = PhaseOrchestrator( config=self.config, phase_publisher=self.phase_publisher, credit_router=self.sticky_router, dataset_metadata=self._dataset_metadata, + parsed_graph=parsed_graph, ) await self._phase_orchestrator.initialize() + def _load_graph_sidecar(self) -> ParsedGraph | None: + """Load the mandatory structural sidecar iff this run is a graph workload. + + The DatasetManager writes ``graph_meta.msgpack`` on EVERY graph build + route and advertises its exact path on the graph-typed + ``DatasetConfiguredNotification.client_metadata``; the schedule plane + ingests the sidecar from that path only. A graph run whose broadcast + is not graph-typed, or whose advertised file is missing, undecodable, + or store-divergent, is a hard configure-time failure. No re-parse + fallback, no env-convention path re-derivation. + """ + from aiperf.dataset.graph.codecs import decode_graph_meta_sidecar + from aiperf.dataset.graph.workload_detect import resolve_graph_workload + + if resolve_graph_workload(self.run) is None: + return None + + meta = self._graph_client_metadata + if meta is None: + raise InvalidStateError( + "graph workload run, but the DatasetConfiguredNotification " + "did not carry GraphSegmentClientMetadata: every graph build " + "must broadcast the graph store and sidecar locations" + ) + sidecar = meta.sidecar_path + if not sidecar.exists(): + raise InvalidStateError( + f"graph_meta sidecar missing at {sidecar} (the path the " + "DatasetManager advertised): if the DatasetManager runs on a " + "different host or filesystem, set " + "AIPERF_DATASET_MMAP_BASE_PATH to a location shared with the " + "TimingManager." + ) + try: + graph, _fp, _version = decode_graph_meta_sidecar(sidecar.read_bytes()) + except Exception as e: + raise InvalidStateError( + f"graph_meta sidecar at {sidecar} is unreadable: {e!r}; " + "rebuild the run so the DatasetManager rewrites it" + ) from e + if not self._sidecar_passes_index_check(graph, sidecar): + raise InvalidStateError( + f"graph_meta sidecar at {sidecar} failed the unified-store " + "index cross-check: the sidecar topology diverged from the " + "stored envelopes the workers read" + ) + self.info(f"Loaded structural graph sidecar for timing: {sidecar}") + return graph + + def _advise_non_streaming_first_token_sources( + self, parsed_graph: ParsedGraph + ) -> None: + """Warn once if a post-TTFT edge's SOURCE node does not stream. + + A first-token-anchored ``StaticEdge`` + (``delay_after_predecessor_first_token_us`` set, post-TTFT anchoring, + validator rule 55) releases its successor at the SOURCE node's OBSERVED + first token + the refined delay. That observation only exists when the + source node itself streams. Each graph node dispatches per its own + recorded ``streaming`` mode (a per-request override), so the global + ``--streaming`` flag does not govern whether a first-token event is + emitted -- a recorded-streaming source streams regardless of the global. + The residual failure is a first-token edge whose source ``LlmNode`` + carries ``streaming=False``: it emits no SSE first-token event, so its + post-TTFT-anchored children SILENTLY fall back to waiting on the + source's COMPLETION latch (dispatch-time delay) instead of its observed + first token -- a replay-fidelity degradation with no other signal. This + is possible only in hand-authored/degenerate graphs; recorded corpora + are consistent by construction (the same recorded ttft drives both the + edge refinement and the source node's streaming mode). Runs once per run + (``_profile_configure_command`` fires once) and is a no-op when every + first-token source streams (or the corpus carries no first-token edges). + """ + from aiperf.dataset.graph.models import LlmNode + from aiperf.timing.strategies.graph_ir_replay import first_token_sources + + graphs = [ + parsed_graph.graph, + *parsed_graph.graphs.values(), + ] + for graph in graphs: + if graph is None: + continue + for source_id in first_token_sources(graph): + source = graph.nodes.get(source_id) + if isinstance(source, LlmNode) and source.streaming is False: + _logger.warning( + "graph corpus contains first-token-anchored edges " + "(post-TTFT anchoring, validator rule 55) whose SOURCE " + f"node ({source_id!r}) has streaming=False: a " + "non-streaming source emits no SSE first-token event, so " + "its post-TTFT-anchored children silently wait on the " + "source's COMPLETION latch (dispatch-time delay) instead " + "of its observed first token -- a replay-fidelity " + "degradation. This is only possible in " + "hand-authored/degenerate graphs; recorded corpora are " + "consistent by construction. Set streaming=True on the " + "source node to restore post-TTFT anchoring." + ) + return + + def _sidecar_passes_index_check(self, graph: ParsedGraph, sidecar: Path) -> bool: + """Return True when the index cross-check passes or is not reachable. + + Cross-checks the sidecar's per-trace ordinals against the unified + store's manifest index (the STORE the worker will actually read). Any + open/read failure -- including a missing store -- is treated as "not + reachable" so the sidecar is accepted as-is (best-effort, additive). + + The store location comes from the broadcast + ``GraphSegmentClientMetadata`` typed fields (``store_base_path`` / + ``benchmark_id``), the same source the workers open, not from + string-parsing the ``sidecar`` path. + """ + from aiperf.dataset.graph.graph_meta_sidecar import sidecar_matches_index + + def _matches(offsets: dict) -> bool: + return sidecar_matches_index(graph, offsets) + + meta = self._graph_client_metadata + if meta is None: # no graph broadcast -> store not reachable + return True + base_path = meta.store_base_path + benchmark_id = meta.benchmark_id + + # The unified store is the sole store shape; absent means not reachable. + try: + from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedClient, + _unified_dir, + ) + + if not _unified_dir(base_path, benchmark_id).exists(): + return True + client = GraphSegmentUnifiedClient( + base_path=base_path, benchmark_id=benchmark_id + ).open() + try: + # inner keys are ":" (see _encode_inner_key); + # sidecar_matches_index wants {(ordinal, variant): _} per trace. + offsets: dict[str, dict[tuple[int, str], object]] = {} + for trace_id, inner in client._node_offsets.items(): + decoded: dict[tuple[int, str], object] = {} + for key, value in inner.items(): + ord_str, _, variant = key.partition(":") + decoded[(int(ord_str), variant)] = value + offsets[trace_id] = decoded + finally: + client.close() + return _matches(offsets) + except Exception: # any failure -> treat as not reachable + return True + async def _wait_for_dataset_or_failure(self) -> None: """Wait for either the dataset-configured or dataset-failed event. diff --git a/src/aiperf/timing/phase/credit_counter.py b/src/aiperf/timing/phase/credit_counter.py index 6df07137a4..0bfc3ad6a7 100644 --- a/src/aiperf/timing/phase/credit_counter.py +++ b/src/aiperf/timing/phase/credit_counter.py @@ -9,6 +9,8 @@ from typing import TYPE_CHECKING +from aiperf.plugin.enums import TimingMode + if TYPE_CHECKING: from aiperf.credit.structs import TurnToSend from aiperf.timing.config import CreditPhaseConfig @@ -25,6 +27,13 @@ class CreditCounter: def __init__(self, config: CreditPhaseConfig) -> None: self._config = config + # Graph-IR phases bypass the linear session arithmetic on BOTH sides: + # ``increment_sent`` detects graph credits per-turn (``trace_id``), but + # ``increment_returned`` receives no turn/credit info, so the bypass is + # keyed off the phase's timing mode -- a GRAPH_IR phase issues ONLY + # trace_id-stamped graph credits (``CreditIssuer.issue_graph_credit``), + # making the two detections equivalent. + self._graph_phase = getattr(config, "timing_mode", None) == TimingMode.GRAPH_IR # Progress counters self._requests_sent: int = 0 @@ -187,13 +196,19 @@ def freeze_completed_counts(self) -> None: def increment_sent(self, turn_to_send: TurnToSend) -> tuple[int, bool]: """Atomically increment sent count and return (credit_index, is_final_credit). + Graph-IR credits (``turn_to_send.trace_id is not None``) bump only + ``_requests_sent`` (the per-node wire/record count) and ALWAYS return + ``is_final_credit=False``: they bypass the linear session arithmetic + entirely because ``GraphIRReplayStrategy`` owns completion and the + session/request stop semantics for a fan-out DAG. + DAG children (``turn_to_send.agent_depth > 0``) count as real HTTP requests and DO bump ``_requests_sent`` — the user-visible "requests sent" metric must reflect actual wire traffic including DAG offspring. They do NOT bump ``_sent_sessions`` or ``_total_session_turns`` because they inherit the parent's - session slot (``CreditIssuer.issue_credit`` skips session-slot - acquisition for them). + session slot (they dispatch via ``CreditIssuer._dispatch_dag_turn``, + which bypasses session-slot acquisition). ``is_final_credit`` flips when the request-count cap is crossed on either a root or a child. This is what drives @@ -209,6 +224,39 @@ def increment_sent(self, turn_to_send: TurnToSend) -> tuple[int, bool]: credit_index = self._requests_sent new_sent_count = self._requests_sent + 1 + if turn_to_send.trace_id is not None: + # Graph-IR credits BYPASS all linear session arithmetic. A weka + # trace is a fan-out DAG, not a linear ``turn_index==0 -> + # num_turns-1`` session: every distinct node fires with + # ``turn_index==0`` on a fresh per-node session key, so the linear + # path would count each NODE as a fresh session and trip + # ``is_final_credit`` after the Nth node (freezing the sent-count + # mid-trace -> lost records). ``GraphIRReplayStrategy`` owns + # completion and the session/request stop semantics, so here we + # only bump the wire-request counter (one per node dispatch, which + # IS the record count) and NEVER flip ``is_final_credit`` from a + # graph credit -- the strategy freezes counts when its executors + # drain. ``--request-count`` is still enforced upstream by the + # ``RequestCountStopCondition`` gate against ``requests_sent``. + self._requests_sent = new_sent_count + return credit_index, False + + return self._increment_linear_sent(turn_to_send, credit_index, new_sent_count) + + def _increment_linear_sent( + self, turn_to_send: TurnToSend, credit_index: int, new_sent_count: int + ) -> tuple[int, bool]: + """Linear (non-graph) sent-count arithmetic for root + DAG-child credits. + + Extracted from ``increment_sent`` (the graph-bypass early-return stays + the leading branch there). Behavior is identical to the prior inline + body: DAG children bump only ``_requests_sent`` and honor the + request-count cap; roots also advance session / root-wire / turn + counters and flip ``is_final_credit`` on the request-count OR + session-completion predicate. + + Lock-free: no async calls. + """ if turn_to_send.agent_depth > 0: # Children: bump the wire-request counter only (slot is # inherited, sampler-plan counters stay root-only). Flip @@ -257,6 +305,18 @@ def increment_returned( ) -> bool: """Atomically increment returned count and check phase completion. + Graph-IR phases (``self._graph_phase``) bypass the session counters + entirely, mirroring the sent-side graph bypass in + :meth:`increment_sent`: every graph credit is minted ``turn_index=0, + num_turns=1`` so ``is_final_turn`` is always True, and counting it here + would report one "completed session" per NODE (with ``sent_sessions`` + pinned at 0 -> negative ``in_flight_sessions`` and a bogus progress %). + Trace-level completion is not visible at this callsite, so graph + session stats stay 0 on both sides (like ``sent_sessions``) and phase + progress is driven purely by the request counters; per-trace + completion is owned and reported by ``GraphIRReplayStrategy`` + (``admitted_traces`` / ``completed_traces``). + Lock-free: no async calls. Args: @@ -271,13 +331,14 @@ def increment_returned( True if ALL sent credits have now been returned or cancelled (phase sending must be complete for this to ever return True). """ + count_sessions = is_final_turn and not self._graph_phase if cancelled: self._requests_cancelled += 1 - if is_final_turn: + if count_sessions: self._cancelled_sessions += 1 else: self._requests_completed += 1 - if is_final_turn: + if count_sessions: self._completed_sessions += 1 if errored: self._request_errors += 1 diff --git a/src/aiperf/timing/phase/runner.py b/src/aiperf/timing/phase/runner.py index 19dd63746f..e02d126161 100644 --- a/src/aiperf/timing/phase/runner.py +++ b/src/aiperf/timing/phase/runner.py @@ -9,6 +9,7 @@ from __future__ import annotations import asyncio +import math from collections.abc import Callable from typing import TYPE_CHECKING @@ -24,7 +25,12 @@ from aiperf.timing.phase.progress_tracker import PhaseProgressTracker from aiperf.timing.phase.stop_conditions import StopConditionChecker from aiperf.timing.ramping import Ramper, RamperConfig, RampType -from aiperf.timing.strategies.core import RateSettableProtocol +from aiperf.timing.session_tree import SessionTreeRegistry +from aiperf.timing.strategies.core import ( + LaneSettableProtocol, + PhaseTeardownStrategyProtocol, + RateSettableProtocol, +) from aiperf.timing.url_samplers import URLSelectionStrategyProtocol if TYPE_CHECKING: @@ -34,6 +40,7 @@ from aiperf.timing.concurrency import ConcurrencyManager from aiperf.timing.config import CreditPhaseConfig from aiperf.timing.conversation_source import ConversationSource + from aiperf.timing.graph_channel import GraphPhaseChannel from aiperf.timing.phase.publisher import PhasePublisher from aiperf.timing.request_cancellation import RequestCancellationSimulator from aiperf.timing.strategies.core import TimingStrategyProtocol @@ -73,7 +80,7 @@ def __init__( self, *, config: CreditPhaseConfig, - conversation_source: ConversationSource, + conversation_source: ConversationSource | None, phase_publisher: PhasePublisher, credit_router: CreditRouterProtocol, concurrency_manager: ConcurrencyManager, @@ -81,13 +88,16 @@ def __init__( callback_handler: CreditCallbackHandler, url_selection_strategy: URLSelectionStrategyProtocol | None = None, branch_orchestrator: BranchOrchestrator | None = None, + graph_channel: GraphPhaseChannel | None = None, **kwargs, ) -> None: """Initialize phase runner. Args: config: Phase configuration (phase enum, stop conditions, concurrency limits). - conversation_source: Source for conversation data (shared across phases). + conversation_source: Source for conversation data (shared across + phases). ``None`` for graph runs, which sample inside the graph + strategy over ``graph_channel.parsed_graph.traces``. phase_publisher: Publishes phase lifecycle events to message bus. credit_router: Routes credits to workers (for cancel_all_credits on timeout). concurrency_manager: Manages session and prefill concurrency slots. @@ -99,16 +109,20 @@ def __init__( ``_is_phase_complete`` consults ``has_pending_branch_work`` so completion blocks while DAG children are still in flight, even after ``--request-count`` is reached. + graph_channel: Graph-run state (parsed graph + consume-once warmup + handoff) threaded from the orchestrator. ``None`` for non-graph + runs; required when ``timing_mode`` is ``GRAPH_IR``. """ super().__init__(**kwargs) self._config = config self._conversation_source = conversation_source self._branch_orchestrator = branch_orchestrator + self._graph_channel = graph_channel # For FIXED_SCHEDULE mode, use actual dataset size instead of config values. # Config values may reflect pre-filtered file size, but dataset_metadata # reflects the actual filtered dataset after start/end offset filtering. - metadata = conversation_source.dataset_metadata + metadata = conversation_source.dataset_metadata if conversation_source else None if config.timing_mode == TimingMode.FIXED_SCHEDULE and metadata: self._config = config.model_copy( update={ @@ -132,6 +146,14 @@ def __init__( lifecycle=self._lifecycle, counter=self._progress.counter, ) + # Per-tree finality ledger, shared by the issuer (opens trees + stamps + # finality) and the branch orchestrator (registers descendants + drives + # terminal transitions). Only built for DAG-shaped datasets; None + # elsewhere so finality stays conservative. Must precede the issuer so + # it can be injected at construction. + self._session_tree_registry = self._build_session_tree_registry( + conversation_source + ) self._credit_issuer = self._build_credit_issuer(url_selection_strategy) self._maybe_construct_branch_orchestrator(conversation_source) @@ -141,6 +163,23 @@ def __init__( self._was_cancelled = False self._rampers: list[Ramper] = [] + def _build_session_tree_registry( + self, conversation_source: ConversationSource | None + ) -> SessionTreeRegistry | None: + """Build the per-tree finality ledger for DAG-shaped datasets only. + + Mirrors the ``_maybe_construct_branch_orchestrator`` gate: a registry is + pointless without DAG fan-out (no descendants to track), and leaving it + None on non-DAG runs keeps ``_finality_for_issue`` conservative. Graph + runs have NO conversation source (they sample inside the graph + strategy) -- the graph adapter stamps its own identity/finality facts. + """ + if conversation_source is None: + return None + if not self._is_dag_dataset(conversation_source.dataset_metadata): + return None + return SessionTreeRegistry() + def _build_credit_issuer( self, url_selection_strategy: URLSelectionStrategyProtocol | None ) -> CreditIssuer: @@ -156,20 +195,25 @@ def _build_credit_issuer( cancellation_policy=self._cancellation_policy, lifecycle=self._lifecycle, url_selection_strategy=url_selection_strategy, + session_tree_registry=self._session_tree_registry, ) def _maybe_construct_branch_orchestrator( - self, conversation_source: ConversationSource + self, conversation_source: ConversationSource | None ) -> None: """Construct ``BranchOrchestrator`` for DAG-shaped datasets. A "DAG-shaped" dataset is one whose metadata declares any branches OR contains any non-root conversations (``agent_depth > 0``). Non-DAG runs leave ``self._branch_orchestrator`` as None and the - callback / strategy paths skip orchestrator hooks. + callback / strategy paths skip orchestrator hooks. Graph runs have no + conversation source (they sample inside the graph strategy) and no + DAG-dataset branches, so the orchestrator is never constructed. """ if self._branch_orchestrator is not None: return + if conversation_source is None: + return if not self._is_dag_dataset(conversation_source.dataset_metadata): return sticky_router = getattr(self._credit_router, "sticky_router", None) @@ -177,6 +221,7 @@ def _maybe_construct_branch_orchestrator( conversation_source=conversation_source, credit_issuer=self._credit_issuer, sticky_router=sticky_router, + session_tree_registry=self._session_tree_registry, ) @property @@ -303,12 +348,15 @@ async def run( raise e finally: self._detach_orchestrator_and_cleanup() + await self._teardown_strategy(strategy) def _build_strategy(self) -> TimingStrategyProtocol: """Construct the timing strategy class for this phase.""" StrategyClass = plugins.get_class( PluginType.TIMING_STRATEGY, self._config.timing_mode ) + if self._config.timing_mode == TimingMode.GRAPH_IR: + return self._build_graph_ir_strategy(StrategyClass) return StrategyClass( config=self._config, conversation_source=self._conversation_source, @@ -321,6 +369,95 @@ def _build_strategy(self) -> TimingStrategyProtocol: progress=self._progress, ) + def _build_graph_ir_strategy( + self, StrategyClass: type[TimingStrategyProtocol] + ) -> TimingStrategyProtocol: + """Construct ``GraphIRReplayStrategy`` with its graph-only injection channel. + + The fixed timing-strategy kwarg set has no slot for the ``ParsedGraph`` + or the graph-return observer, so this branch extends construction + minimally WITHOUT touching the other strategies' path. The ``ParsedGraph`` is read from the typed ``GraphPhaseChannel`` + threaded orchestrator -> runner -> strategy; the single unconditional + graph-return observer is installed via the shared + ``CreditCallbackHandler.set_graph_return_observer``. + + The per-trace t* snapshot window (``start_min_ratio`` / + ``start_max_ratio``) and the phase-start burst mode are per-run config + carried on this phase's ``CreditPhaseConfig`` + (``--trajectory-start-min/max-ratio`` / ``--burst-phase-starts``, + threaded by ``TimingConfig.from_run``); the sampling seed + (``t_star_random_seed``) is the run's resolved ``--random-seed`` -- + the same seed that drives synthesized content, so one seed reproduces + the whole run. Without this the strategy would + fall back to its inert ``[0, 0]`` defaults (t*=0 identity replay) and + the t* machinery would never engage. The default window is ``0.0..0.0`` + (t*=0, full native replay); the ``0.0..1.0`` AgentX + ``TrajectorySource`` window is applied only by + ``--scenario inferencex-agentx-mvp`` or explicit flags, so a bare + graph run replays natively out of the box. + + The extended (cache-pressure) warmup duration + (``--agentic-cache-warmup-duration``, per-run config on this phase's + ``CreditPhaseConfig``) is threaded in as + ``cache_pressure_duration_s``, arming the WARMUP strategy's pressure + stage (None disables it). Any ``GraphWarmupHandoff`` the WARMUP phase + stashed on the shared ``GraphPhaseChannel`` is popped consume-once and + threaded in as ``warmup_handoff`` so the first PROFILING graph phase + resumes each lane at its recorded execution frontier. + """ + channel = self._graph_channel + if channel is None: + raise RuntimeError( + "TimingMode.GRAPH_IR selected but no GraphPhaseChannel was " + "provided; graph runs must thread the parsed graph through " + "the orchestrator's graph channel." + ) + parsed_graph = channel.parsed_graph + # Extended-warmup handoff: pop (consume-once) so a later phase built + # from the same shared channel never sees a stale re-cut. The WARMUP + # phase is built before any handoff exists, so the pop is a no-op + # there; the first PROFILING graph phase after an extended warmup + # consumes it. + warmup_handoff = channel.warmup_handoff + channel.warmup_handoff = None + # getattr: unit harnesses inject duck-typed handlers exposing only the + # set_* API; the production CreditCallbackHandler always has clear_*. + return StrategyClass( + config=self._config, + conversation_source=self._conversation_source, + scheduler=self._scheduler, + stop_checker=self._stop_checker, + credit_issuer=self._credit_issuer, + lifecycle=self._lifecycle, + parsed_graph=parsed_graph, + graph_channel=channel, + register_observer=self._callback_handler.set_graph_return_observer, + register_first_token_observer=( + self._callback_handler.set_graph_first_token_observer + ), + unregister_observer=getattr( + self._callback_handler, "clear_graph_return_observer", None + ), + unregister_first_token_observer=getattr( + self._callback_handler, "clear_graph_first_token_observer", None + ), + start_min_ratio=getattr(self._config, "trajectory_start_min_ratio", None) + or 0.0, + start_max_ratio=getattr(self._config, "trajectory_start_max_ratio", None) + or 0.0, + t_star_random_seed=getattr(self._config, "random_seed", None) or 0, + burst_phase_starts=bool(getattr(self._config, "burst_phase_starts", None)), + cache_pressure_duration_s=getattr( + self._config, "cache_pressure_duration", None + ), + warmup_handoff=warmup_handoff, + dataset_sampling_strategy=getattr( + self._config, "dataset_sampling_strategy", None + ), + allow_dataset_wrap=getattr(self._config, "allow_dataset_wrap", None), + cache_bust=getattr(self._config, "cache_bust", None), + ) + def _register_strategy_with_callback_handler( self, strategy: TimingStrategyProtocol ) -> None: @@ -337,6 +474,30 @@ def _register_strategy_with_callback_handler( if self._branch_orchestrator is not None: self._callback_handler.set_branch_orchestrator(self._branch_orchestrator) + async def _teardown_strategy(self, strategy: TimingStrategyProtocol) -> None: + """Invoke the strategy's optional ``teardown_phase`` hook exactly once. + + Runs in ``run()``'s ``finally`` so per-phase resources (e.g. the graph + strategy's return / first-token observers and retained sticky trace + lifecycles) are released on EVERY exit path -- drain, timeout, error, + cancel. Non-final seamless phases defer the teardown to the background + return-wait completion: their returns are still in flight when + ``run()`` exits and the strategy's return observer must stay installed + until they land. Strategies without the hook (all linear strategies) + are skipped -- the default no-op. + """ + if not isinstance(strategy, PhaseTeardownStrategyProtocol): + return + if self._return_wait_task is not None and not self._return_wait_task.done(): + self._return_wait_task.add_done_callback( + lambda _task: self.execute_async(strategy.teardown_phase()) + ) + return + try: + await strategy.teardown_phase() + except Exception as exc: + self.warning(f"Strategy teardown_phase failed: {exc!r}") + def _detach_orchestrator_and_cleanup(self) -> None: """Final-pass orchestrator teardown for the phase. @@ -414,6 +575,15 @@ async def _run_strategy( await self._wait_for_returning_complete() self._progress_task.cancel() + report = getattr(strategy, "report_warmup_failures", None) + if report is not None: + # A WARMUP phase that recorded terminal failures aborts the + # run here (agentx parity): the raise propagates through + # run()'s except -> _publish_phase_failure_lifecycle so other + # services see the phase end. Linear strategies lack the hook + # (getattr None) and are unaffected. + report() + for ramper in self._rampers: ramper.stop() self._scheduler.cancel_all() @@ -464,10 +634,15 @@ def _create_rampers(self, strategy: TimingStrategyProtocol) -> None: self._rampers = [] config = self._config - # Session concurrency ramper (stepped mode) + # Session concurrency ramper (stepped mode). Graph strategies expose + # LANE admission ramping (LaneSettableProtocol) because graph credits + # bypass session slots -- the same ramper drives both targets, so one + # flag shapes the ramp on either plane. if config.concurrency_ramp_duration_sec and config.concurrency: + lane_settable = isinstance(strategy, LaneSettableProtocol) + what = "lane" if lane_settable else "session" self.info( - f"Starting session concurrency ramp: 1 → {config.concurrency} " + f"Starting {what} concurrency ramp: 1 → {config.concurrency} " f"over {config.concurrency_ramp_duration_sec}s" ) ramp_config = RamperConfig( @@ -478,9 +653,9 @@ def _create_rampers(self, strategy: TimingStrategyProtocol) -> None: ) def setter(limit: float) -> None: - return self._concurrency_manager.set_session_limit( - config.phase, int(limit) - ) + self._concurrency_manager.set_session_limit(config.phase, int(limit)) + if lane_settable: + strategy.set_lane_limit(int(limit)) self._rampers.append(Ramper(setter=setter, config=ramp_config)) @@ -638,6 +813,19 @@ async def _wait_for_returning_complete(self) -> None: return timeout = self._lifecycle.time_left_in_seconds(include_grace_period=True) + if timeout is None and self._config.phase == CreditPhase.WARMUP: + # A duration-less phase never reaches the lifecycle's grace math + # (time_left_in_seconds returns None when expected_duration_sec is + # None), which made a finite warmup grace_period_sec DEAD config: + # the returning wait was unbounded and a lost pressure return hung + # the run. Bound the WARMUP drain by the finite grace directly -- + # agentx's dedicated accelerated-warmup drain wait does exactly + # this (its runner.py:665-686). Infinite grace (the boundary- + # priming default) keeps today's unbounded wait; non-WARMUP phases + # are out of scope. + grace = self._config.grace_period_sec + if grace is not None and math.isfinite(grace): + timeout = grace timed_out = await self._wait_for_event_with_timeout( name=f"{self._config.phase} phase credits returned", event=self._progress.all_credits_returned_event, diff --git a/src/aiperf/timing/phase_orchestrator.py b/src/aiperf/timing/phase_orchestrator.py index d829747c52..03e13dcc2a 100644 --- a/src/aiperf/timing/phase_orchestrator.py +++ b/src/aiperf/timing/phase_orchestrator.py @@ -24,6 +24,7 @@ from aiperf.plugin.enums import PluginType from aiperf.timing.concurrency import ConcurrencyManager from aiperf.timing.conversation_source import ConversationSource +from aiperf.timing.graph_channel import GraphPhaseChannel from aiperf.timing.phase.runner import PhaseRunner from aiperf.timing.request_cancellation import RequestCancellationSimulator from aiperf.timing.url_samplers import URLSelectionStrategyProtocol @@ -31,6 +32,8 @@ if TYPE_CHECKING: from aiperf.common.models import DatasetMetadata from aiperf.credit.sticky_router import CreditRouterProtocol + from aiperf.dataset.graph.models import ParsedGraph + from aiperf.dataset.protocols import DatasetSamplingStrategyProtocol from aiperf.timing.config import TimingConfig from aiperf.timing.phase.publisher import PhasePublisher @@ -86,6 +89,7 @@ def __init__( phase_publisher: PhasePublisher, credit_router: CreditRouterProtocol, dataset_metadata: DatasetMetadata, + parsed_graph: ParsedGraph | None = None, **kwargs, ) -> None: """Initialize timing strategy and orchestration components. @@ -95,6 +99,12 @@ def __init__( phase_publisher: Publishes phase events to message bus credit_router: Routes credits to workers dataset_metadata: Dataset for conversation sampling + parsed_graph: For graph IR workloads, the parsed conversation DAG. + When set, the orchestrator owns a ``GraphPhaseChannel`` carrying + it (and the WARMUP handoff) to every ``PhaseRunner``; the + conversation-shaped sampler and source are NOT built because + graph runs sample inside the graph strategy over + ``parsed_graph.traces``. ``None`` for non-graph workloads. """ super().__init__(**kwargs) self._config = config @@ -102,32 +112,23 @@ def __init__( self._credit_router = credit_router self._dataset_metadata = dataset_metadata - # Create dataset sampler - SamplerClass = plugins.get_class( - PluginType.DATASET_SAMPLER, - self._dataset_metadata.sampling_strategy, - ) - # Only root conversations are sampled by the strategy. DAG - # children belong to their root's session and are dispatched by - # the BranchOrchestrator on credit return — sampling them as - # roots would create duplicate root sessions. Filter on - # ``is_root`` rather than ``agent_depth == 0`` so SPAWN-mode - # children (which keep ``agent_depth == 0`` for fresh-context - # semantics but carry ``is_root=False``) are also excluded. - root_conv_ids = [ - c.conversation_id - for c in self._dataset_metadata.conversations - if getattr(c, "is_root", True) - ] - self._dataset_sampler = SamplerClass( - conversation_ids=root_conv_ids - or [c.conversation_id for c in self._dataset_metadata.conversations] - ) - # Long-lived components (shared across phases) - self._conversation_source = ConversationSource( - self._dataset_metadata, self._dataset_sampler - ) + if parsed_graph is None: + self._dataset_sampler: DatasetSamplingStrategyProtocol | None = ( + self._build_dataset_sampler() + ) + self._conversation_source: ConversationSource | None = ( + self._build_conversation_source() + ) + self._graph_channel: GraphPhaseChannel | None = None + else: + # Graph runs sample inside the graph strategy + # (GraphIRConversationSource over parsed_graph.traces); the + # conversation-shaped sampler and source have nothing to sample + # (conversations is empty) and are not built. + self._dataset_sampler = None + self._conversation_source = None + self._graph_channel = GraphPhaseChannel(parsed_graph=parsed_graph) self._concurrency_manager = ConcurrencyManager() self._cancellation_policy = RequestCancellationSimulator( config.request_cancellation @@ -154,17 +155,47 @@ def __init__( # Active phase runners (for cancellation) - multiple possible with seamless mode self._active_runners: list[PhaseRunner] = [] + def _build_dataset_sampler(self) -> DatasetSamplingStrategyProtocol: + """Construct the root-conversation sampler for this run. + + Only root conversations are sampled by the strategy. DAG children belong + to their root's session and are dispatched by the BranchOrchestrator on + credit return -- sampling them as roots would create duplicate root + sessions. Filter on ``is_root`` rather than ``agent_depth == 0`` so + SPAWN-mode children (which keep ``agent_depth == 0`` for fresh-context + semantics but carry ``is_root=False``) are also excluded. + """ + SamplerClass = plugins.get_class( + PluginType.DATASET_SAMPLER, + self._dataset_metadata.sampling_strategy, + ) + root_conv_ids = [ + c.conversation_id + for c in self._dataset_metadata.conversations + if getattr(c, "is_root", True) + ] + return SamplerClass( + conversation_ids=root_conv_ids + or [c.conversation_id for c in self._dataset_metadata.conversations] + ) + + def _build_conversation_source(self) -> ConversationSource: + """Construct the long-lived conversation source (non-graph runs only).""" + return ConversationSource(self._dataset_metadata, self._dataset_sampler) + @property - def conversation_source(self) -> ConversationSource: - """Conversation source for dataset access.""" + def conversation_source(self) -> ConversationSource | None: + """Conversation source for dataset access (``None`` for graph runs).""" return self._conversation_source @on_init async def _init_orchestrator(self) -> None: """Log configured phases (actual initialization happens per-phase in _execute_phases).""" self.info( - lambda: f"Initialized {len(self._ordered_phase_configs)} phase(s): " - f"{[p.phase.replace('_', ' ').title() for p in self._ordered_phase_configs]}" + lambda: ( + f"Initialized {len(self._ordered_phase_configs)} phase(s): " + f"{[p.phase.replace('_', ' ').title() for p in self._ordered_phase_configs]}" + ) ) @on_start @@ -208,6 +239,7 @@ async def _execute_phases(self) -> None: cancellation_policy=self._cancellation_policy, callback_handler=self._callback_handler, url_selection_strategy=self._url_sampler, + graph_channel=self._graph_channel, ) # For seamless non-final phases, set callback to remove from active runners @@ -229,9 +261,17 @@ async def _execute_phases(self) -> None: await self.cancel() raise e - # Remove from active runners when fully complete - # For seamless phases, this happens after returns complete (background task) - if not is_seamless_non_final: + # Remove from active runners when fully complete. + # For seamless phases, this happens after returns complete (background + # task). A concurrent graceful cancel (PROFILE_CANCEL / Ctrl+C -> + # ``cancel`` -> ``_cancel_active_runners``) clears the whole list while + # this coroutine is parked in ``runner.run`` (e.g. a graph phase parked + # mid idle-gap replay), so the runner may already be gone here. + # Discard rather than ``remove`` so the teardown does not raise + # ``ValueError: list.remove(x): x not in list`` -- which would fail the + # orchestrator and abort the records-manager's in-flight graceful export + # (no metrics JSON, ``was_cancelled`` never set). + if not is_seamless_non_final and runner in self._active_runners: self._active_runners.remove(runner) def _phase_runner_cleanup_callback(self, runner: PhaseRunner) -> Callable[[], None]: @@ -273,7 +313,9 @@ async def _stop_orchestrator(self) -> None: """ if self._active_runners: self.debug( - lambda: f"Stopping orchestrator with {len(self._active_runners)} active runner(s)" + lambda: ( + f"Stopping orchestrator with {len(self._active_runners)} active runner(s)" + ) ) self._cancel_active_runners() diff --git a/src/aiperf/timing/session_tree.py b/src/aiperf/timing/session_tree.py new file mode 100644 index 0000000000..5466423381 --- /dev/null +++ b/src/aiperf/timing/session_tree.py @@ -0,0 +1,299 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-session-TREE liveness bookkeeping for DAG-lineage finality. + +A DAG session is not a single root conversation -- it is a whole TREE: a +depth-0 root plus every child it spawns, recursively (FORK/SPAWN children and +their subchildren). ``SessionTreeRegistry`` tracks the liveness of each such +tree, keyed by the tree's ``root_correlation_id`` (the depth-0 root's +x_correlation_id), so the credit issuer can stamp two conservative facts on +every emitted ``Credit`` at issue time: + + - ``is_parent_final`` -- the credit's parent had already returned its terminal + turn (v1: determinable only when the parent IS the root). + - ``is_tree_final`` -- this credit is provably the last request the whole tree + will ever send. + +This registry is **finality bookkeeping only**. Unlike the agentic-replay +variant it was ported from, it does NOT own session-slot release: on main the +credit callback handler releases the root's session slot on the root's terminal +turn (``_release_slots_for_return``) exactly as before, and DAG children inherit +the root's slot. The registry never touches the concurrency manager and fires no +drain callback -- it only answers the two finality queries below from live tree +state. + +Wiring (DAG datasets only; the registry is None on non-DAG paths, so +``_finality_for_issue`` returns the conservative ``(None, False)``): + - ``CreditIssuer._open_session_tree`` -> :meth:`open_tree` when a root + session's first credit is issued (turn 0 session-slot acquisition). + - ``BranchOrchestrator`` -> :meth:`register_descendants` when a parent turn + spawns children, :meth:`on_descendant_done` at every descendant-terminal + point, and :meth:`on_root_terminal` from ``intercept`` after the root's + final-turn return is processed (so final-turn spawns register first). + - ``BranchOrchestrator.cleanup`` -> :meth:`release_all` at phase teardown to + drop any still-open trees. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from aiperf.common.aiperf_logger import AIPerfLogger +from aiperf.common.enums import CreditPhase + +_logger = AIPerfLogger(__name__) + + +@dataclass(slots=True) +class _TreeState: + """Liveness of one session tree.""" + + phase: CreditPhase + """Phase the tree was opened under; teardown targets the same phase.""" + root_pending: bool + """True until the root's terminal turn returns.""" + outstanding: int = 0 + """Descendants currently live or registered-pending (any depth).""" + released: bool = False + """Set once the tree has been retired; guards against double retire.""" + + @property + def drained(self) -> bool: + """True when the root is done and no descendant work remains.""" + return not self.root_pending and self.outstanding <= 0 + + +class SessionTreeRegistry: + """Tracks per-session-tree liveness for DAG-lineage finality queries. + + See the module docstring for the invariant and wiring. Single-threaded + asyncio: every method runs to completion between awaits, so the counter + mutations are atomic without locking. Finality bookkeeping only -- this + class never acquires or releases a concurrency slot. + """ + + def __init__(self) -> None: + self._trees: dict[str, _TreeState] = {} + # Descendants registered before their tree is opened. On main's flow a + # root's tree is opened at turn-0 issue, before any child can spawn (a + # spawn only happens on a parent turn RETURN), so this buffer is a + # defensive belt that keeps the branch's state transitions identical. + self._pending_descendants: dict[str, int] = {} + # Roots whose tree already RETIRED (drained), mapped to the phase they + # were opened under. A descendant's final-turn SPAWN registers its + # grandchildren AFTER that descendant's own on_descendant_done drained + # the tree (callback order: the child-completion decrement precedes the + # return-intercept that spawns). Without this ledger those grandchildren + # would buffer into a retired root that nothing drains, and is_tree_final + # could never answer True for them. register_descendants consults it to + # RESURRECT the tree root-terminal instead. Cleared at teardown. + self._retired_roots: dict[str, CreditPhase] = {} + # Peak simultaneously-open trees == peak tree concurrency; logged at + # teardown so any overshoot is visible. + self._peak_open: int = 0 + # A non-zero count means a descendant returned for a tree that was + # already retired -- i.e. the tree drained while that work was still in + # flight (premature drain). Surfaces span-overhang. + self._late_events: int = 0 + + def open_tree( + self, root_corr: str, phase: CreditPhase, *, root_pending: bool + ) -> None: + """Record a newly admitted tree when its root's first credit is issued. + + Args: + root_corr: the tree's ``root_correlation_id`` (depth-0 root's + x_correlation_id). + phase: phase the tree was opened under. + root_pending: True when a root credit is in flight that will reach a + terminal turn (always True on main -- open_tree fires only at a + root session start). + """ + existing = self._trees.get(root_corr) + if existing is not None and not existing.released: + # Duplicate open for a still-live tree: keep the original (its + # outstanding count is authoritative). Should not happen with + # correct wiring; log so a regression is visible. + _logger.warning( + lambda: f"open_tree for already-open tree root_corr={root_corr!r}; " + "ignoring duplicate" + ) + return + state = _TreeState(phase=phase, root_pending=root_pending) + state.outstanding += self._pending_descendants.pop(root_corr, 0) + # A freshly-opened tree supersedes any stale retired record for the same + # id (correlation ids are unique, so this is defensive). + self._retired_roots.pop(root_corr, None) + self._trees[root_corr] = state + if len(self._trees) > self._peak_open: + self._peak_open = len(self._trees) + + @property + def peak_open(self) -> int: + """Maximum simultaneously-open trees seen (== peak tree concurrency).""" + return self._peak_open + + @property + def late_events(self) -> int: + """Count of returns for already-retired trees (premature-drain evidence).""" + return self._late_events + + def has_tree(self, root_corr: str) -> bool: + """True when this registry is tracking ``root_corr`` (engagement gate).""" + return root_corr in self._trees + + def root_terminal(self, root_correlation_id: str) -> bool | None: + """True when the tree's root has returned its terminal turn; None if unknown. + + Queried at credit-issue time while the tree is still live (a descendant + being issued keeps it in ``_trees``). A fully drained tree has been + retired and reads as unknown/None. + """ + state = self._trees.get(root_correlation_id) + if state is None: + return None + return not state.root_pending + + def is_last_tree_request( + self, + root_correlation_id: str, + *, + is_final_turn: bool, + is_root_credit: bool, + has_branches: bool, + ) -> bool: + """Conservative: True only when this credit is provably the tree's last request. + + ``has_branches`` means "this turn declares ANY branch (FORK or SPAWN) + and will therefore spawn descendants on its return". It must NOT be the + FORK-only ``has_forks`` flag: SPAWN children register with this registry + only at return-intercept, AFTER the issuer stamped finality, so a + spawning final turn would otherwise read as last while its children are + pending (wrong-True, violating the conservative contract). + """ + if not is_final_turn or has_branches: + return False + state = self._trees.get(root_correlation_id) + if state is None: + return False + if is_root_credit: + return state.outstanding <= 0 + return (not state.root_pending) and state.outstanding == 1 + + def register_descendants(self, root_corr: str, n: int = 1) -> None: + """Add ``n`` descendants (spawned under one root) to a tree. + + Called when a parent turn spawns children, keyed by the tree's root id. + Three cases when the tree is not live: + - RETIRED root (in ``_retired_roots``): resurrect it root-terminal. + A descendant's final-turn SPAWN registers its grandchildren after + that descendant's own ``on_descendant_done`` drained the tree; the + root was terminal at retire (``drained`` requires ``not + root_pending``), so recreate the tree with ``root_pending=False`` + and these descendants outstanding -- keeping ``is_tree_final`` + answerable and the count coherent (the grandchildren re-drain it + when they finish). + - Unopened root: buffer the count so a later ``open_tree`` folds it in + (defensive -- on main the tree is always open first). + """ + if n <= 0: + return + state = self._trees.get(root_corr) + if state is not None: + state.outstanding += n + return + retired_phase = self._retired_roots.pop(root_corr, None) + if retired_phase is not None: + self._trees[root_corr] = _TreeState( + phase=retired_phase, root_pending=False, outstanding=n + ) + if len(self._trees) > self._peak_open: + self._peak_open = len(self._trees) + return + self._pending_descendants[root_corr] = ( + self._pending_descendants.get(root_corr, 0) + n + ) + + def on_descendant_done(self, root_corr: str) -> bool: + """Account one descendant terminally completing (leaf / error / stopped / + dispatch-rollback). Retires the tree iff it is now drained. + + Returns True if the tree was retired by this call. + """ + state = self._trees.get(root_corr) + if state is None: + pending = self._pending_descendants.get(root_corr, 0) + if pending > 0: + self._pending_descendants[root_corr] = pending - 1 + else: + self._late_events += 1 + return False + if state.outstanding > 0: + state.outstanding -= 1 + return self._retire_if_drained(root_corr, state) + + def on_root_terminal(self, root_corr: str) -> bool: + """Account the root's terminal turn returning (or root cancel/truncate). + + Clears ``root_pending``; retires the tree iff it is now drained. Must be + called AFTER the returning root credit's intercept has run, so any + children spawned on the final turn are already registered. + + Returns True if the tree was retired by this call. False both when the + tree is held (descendants remain) AND when ``root_corr`` is untracked. + """ + state = self._trees.get(root_corr) + if state is None: + return False + state.root_pending = False + return self._retire_if_drained(root_corr, state) + + def _retire_if_drained(self, root_corr: str, state: _TreeState) -> bool: + """Pop a drained tree from tracking. Finality-only: releases no slot.""" + if state.released or not state.drained: + return False + state.released = True + self._trees.pop(root_corr, None) + # Remember the retired root (with its phase) so a late final-turn SPAWN + # can RESURRECT it instead of silently buffering. See register_descendants. + self._retired_roots[root_corr] = state.phase + return True + + def release_all(self, phase: CreditPhase | None = None) -> int: + """Retire every still-open tree at teardown (optionally one phase). + + Finality-only: drops the tracking entries; releases no concurrency slot + (the callback handler owns slot release on main). The registry is + created per-phase, so ``phase=None`` (retire all) is the common + teardown call. Returns the number of trees retired. + """ + to_release = [ + root_corr + for root_corr, state in self._trees.items() + if (phase is None or state.phase == phase) and not state.released + ] + for root_corr in to_release: + state = self._trees.pop(root_corr, None) + if state is None or state.released: + continue + state.released = True + # Drop the transient buffers: with the trees above retired, any + # pending/retired descendant accounting refers to no live tree. On a + # full teardown (``phase is None`` -- the common per-phase call, since + # the registry is created per-phase) clear both; a phase-scoped call + # clears only that phase's retired roots (``_pending_descendants`` + # carries no phase and is left untouched then). + if phase is None: + self._pending_descendants.clear() + self._retired_roots.clear() + else: + self._retired_roots = { + r: p for r, p in self._retired_roots.items() if p != phase + } + return len(to_release) + + def open_count(self, phase: CreditPhase | None = None) -> int: + """Number of trees currently tracked (optionally for one phase).""" + if phase is None: + return len(self._trees) + return sum(1 for state in self._trees.values() if state.phase == phase) diff --git a/src/aiperf/timing/snapshot_chop.py b/src/aiperf/timing/snapshot_chop.py new file mode 100644 index 0000000000..4543c272db --- /dev/null +++ b/src/aiperf/timing/snapshot_chop.py @@ -0,0 +1,303 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""t* frontier chop for a segment-trie ``ParsedGraph``. + +Format-agnostic: operates purely on ``ParsedGraph`` topology, ``StaticEdge``s, +per-node ``arrival_offset_us``, and AND-fan-in ``inputs`` -- no recorded-trace +(weka/dynamo) types. Any adapter that emits the segment-trie IR (a graph of +``LlmNode`` + ``StaticEdge`` with ``metadata["trie"]``) can be snapshotted here. +Also hosts the extended-warmup handoff chop (``chop_trie_at_frontier``). +""" + +from __future__ import annotations + +from typing import Any + +import msgspec + +from aiperf.dataset.graph.models import START_NODE_ID, ParsedGraph, StaticEdge + + +def chop_trie_at_tstar(graph: ParsedGraph, t_star_us: int) -> ParsedGraph: + """Chop a segment-trie ``ParsedGraph`` to its live frontier at ``t*``. + + The segment-trie graph is the trivial ``LlmNode`` + ``StaticEdge`` realization + of one recorded trace; every node carries ``arrival_offset_us`` (its recorded + ``t`` * 1e6). The snapshot chop: + + * Drops every node whose ``arrival_offset_us < t_star_us`` -- the pre-``t*`` + turns were already WARMED (primed), not PROFILED. + * Re-roots each SURVIVING node that lost ALL its predecessors to the chop + from ``START`` via a synthetic ``StaticEdge(START -> node)`` whose + ``min_start_delay_us = arrival_offset_us - t_star_us`` -- the node's + ABSOLUTE offset from the instance run-origin ``t*`` (the executor anchors + this via ``absolute_start_offsets=True``). Inter-turn edges between two + SURVIVING nodes are kept verbatim (whichever delay kind they carry -- + ``delay_after_predecessor_us`` end-to-start OR + ``delay_after_predecessor_start_us`` dispatch-to-start -- is unchanged). + * Leaves each surviving node's ``metadata["trie"]["prompt_segment_ids"]`` + UNCHANGED: the full pre-``t*`` prefix stays in the path so the worker still + materializes the EXACT resume prompt (no truncation needed -- the dropped + turns were dispatched during warmup and the server holds their KV). + + ``t_star_us <= 0`` => the graph is returned UNCHANGED (full t*=0 replay). + """ + if t_star_us <= 0: + return graph + + old_graph = graph.graph + survivors = { + nid: node + for nid, node in old_graph.nodes.items() + if (node.arrival_offset_us or 0) >= t_star_us + } + + new_edges = _chop_edges(old_graph.edges, survivors, t_star_us) + + # Recompute each surviving node's AND-fan-in ``inputs`` against the chop: a + # requirement on a DROPPED predecessor's ``{src}_out`` channel would deadlock + # ``await_inputs`` (that channel is never written post-chop). Keep only + # requirements whose source survives; a node re-rooted entirely from START + # ends with empty ``inputs``. + survivor_out_channels = {f"{nid}_out" for nid in survivors} + rescoped = { + nid: _chop_node_inputs(node, survivor_out_channels) + for nid, node in survivors.items() + } + + new_graph = msgspec.structs.replace(old_graph, nodes=rescoped, edges=new_edges) + return msgspec.structs.replace(graph, graph=new_graph) + + +def _chop_edges( + edges: list[Any], survivors: dict[str, Any], t_star_us: int +) -> list[StaticEdge]: + """Recompute the chopped graph's edge set against the surviving frontier. + + An edge survives only when BOTH endpoints survive (or it roots an explicitly + kept node at START). Each surviving node that lost ALL its predecessors to the + chop is re-rooted from START at its t*-relative absolute offset, dropping any + kept START edge for it (the builder rooted pre-t* roots at START with the + recorded absolute offset; the frontier re-root replaces it with the + t*-relative offset). + """ + kept_edges: list[StaticEdge] = [] + has_surviving_pred: set[str] = set() + for edge in edges: + if not isinstance(edge, StaticEdge): + continue + src, tgt = edge.source, edge.target + if tgt not in survivors: + continue + if src == START_NODE_ID or src in survivors: + kept_edges.append(edge) + if src != START_NODE_ID: + has_surviving_pred.add(tgt) + + new_edges = [ + e + for e in kept_edges + if not (e.source == START_NODE_ID and e.target not in has_surviving_pred) + ] + for nid, node in survivors.items(): + if nid not in has_surviving_pred: + new_edges.append( + StaticEdge( + source=START_NODE_ID, + target=nid, + min_start_delay_us=float((node.arrival_offset_us or 0) - t_star_us), + ) + ) + return new_edges + + +def _chop_node_inputs(node: Any, survivor_out_channels: set[str]) -> Any: + """Drop a surviving node's ``inputs`` requirements on dropped predecessors. + + Non-``LlmNode`` (or input-free) nodes pass through untouched. An ``inputs`` + list whose surviving subset equals the original is returned unchanged so the + common no-chop-of-this-node case avoids a struct rebuild. + """ + inputs = getattr(node, "inputs", None) + if not inputs: + return node + kept = [req for req in inputs if req.channel in survivor_out_channels] + if len(kept) == len(inputs): + return node + return msgspec.structs.replace(node, inputs=kept) + + +def chop_trie_at_frontier( + graph: ParsedGraph, + *, + t_star_us: float, + executed: frozenset[str], + return_wall_us: dict[str, float], + drain_end_wall_us: float, + residual_cap_us: float | None = None, +) -> ParsedGraph: + """Chop a segment-trie ``ParsedGraph`` to its extended-warmup handoff frontier. + + The extended (cache-pressure) warmup replays the post-``t*`` remainder with + zero idle delay; at drain, PROFILING must resume each chain at its first + NOT-yet-executed node rather than re-firing from ``t*``. The chop: + + * Drops every node with ``arrival_offset_us < t_star_us`` (pre-``t*`` + history, primed by the boundary warmup) AND every node in ``executed`` + (dispatched-and-returned during warmup/pressure -- the server holds + their KV). + * Keeps inter-survivor edges verbatim (recorded pacing resumes past the + frontier). + * Re-roots each surviving node that lost ALL its real predecessors from + ``START`` with ``min_start_delay_us`` set to its RESIDUAL delay: for each + dropped predecessor edge, ``max(0, recorded_delay - max(0, + drain_end_wall_us - return_wall_us[pred]))`` -- the recorded gap minus + wall time already spent waiting for the drain; + AND-fan-in takes the max across + dropped predecessors, and the result is clamped to ``residual_cap_us`` + when set (graph edge delays + are not bounded by the build-plane idle-gap warp when another stream's + active interval covers the gap, so the cap is load-bearing). The + recorded base uses ONLY end-anchored quantities + (``delay_after_predecessor_us``, edge ``min_start_delay_us``): the + ledger wall is the predecessor's RETURN, so a dispatch-anchored delay + (``delay_after_predecessor_start_us`` / first-token) debited from a + return-anchored elapsed would over-delay by the pred's live service + time -- start-anchored edges contribute 0 (burst). + A dropped predecessor with no recorded wall contributes 0 + (fire immediately -- consistent with compressed pressure pacing, where + recorded leads are considered consumed). + * Rescopes surviving nodes' AND-fan-in ``inputs`` exactly like + :func:`chop_trie_at_tstar` (a requirement on a dropped predecessor's + channel would deadlock ``await_inputs``). A survivor that KEEPS a + surviving-pred edge but LOST a residual-carrying binding edge to the chop + is NOT re-rooted; instead that dropped edge's residual is FOLDED into the + node's ``min_start_delay_us`` (max-combined with any existing node value). + Under ``absolute_start_offsets=True`` the executor anchors that node-level + gate to the instance run-start -- the same anchor the re-root residuals + use -- and max-combines it with the surviving edge gates + (``_compute_firing_gate_us``), so the dropped binding gap survives instead + of vanishing behind a zero-delay join edge. + + ``return_wall_us`` and ``drain_end_wall_us`` share one monotonic clock + (the warmup strategy's ledger); only differences are meaningful. + """ + old_graph = graph.graph + survivors = { + nid: node + for nid, node in old_graph.nodes.items() + if (node.arrival_offset_us or 0) >= t_star_us and nid not in executed + } + + new_edges, kept_pred_residuals = _frontier_edges( + old_graph.edges, + survivors, + return_wall_us=return_wall_us, + drain_end_wall_us=drain_end_wall_us, + residual_cap_us=residual_cap_us, + ) + + survivor_out_channels = {f"{nid}_out" for nid in survivors} + rescoped: dict[str, Any] = {} + for nid, node in survivors.items(): + node = _chop_node_inputs(node, survivor_out_channels) + residual = kept_pred_residuals.get(nid) + if residual is not None: + # A dropped binding edge's residual must not vanish just because a + # zero-delay join edge from a SURVIVING pred remains: fold it into + # the node-level gate. Under the strategy's absolute_start_offsets + # this anchors to the instance run-start -- the same anchor the + # re-root residuals use -- and the executor max-combines it with the + # surviving edge gates (models.py min_start_delay_us contract). + node = msgspec.structs.replace( + node, + min_start_delay_us=max(node.min_start_delay_us or 0.0, residual), + ) + rescoped[nid] = node + + new_graph = msgspec.structs.replace(old_graph, nodes=rescoped, edges=new_edges) + return msgspec.structs.replace(graph, graph=new_graph) + + +def _frontier_edges( + edges: list[Any], + survivors: dict[str, Any], + *, + return_wall_us: dict[str, float], + drain_end_wall_us: float, + residual_cap_us: float | None = None, +) -> tuple[list[StaticEdge], dict[str, float]]: + """Edge set for the handoff chop: keep inter-survivor, re-root at residuals. + + Mirrors :func:`_chop_edges` structurally; the ONLY divergence is the + re-root offset: the t* chop rebases to the recorded absolute offset + (``arrival - t*``) because nothing was replayed yet, while the frontier + chop uses the residual-of-recorded-gap because pressure already consumed + the recorded leads. + + Returns ``(edges, kept_pred_residuals)``. ``kept_pred_residuals`` maps a + survivor id to the leftover residual of a dropped binding edge whose target + ALSO retains a surviving-pred edge (so it is not re-rooted). Those residuals + would otherwise be discarded; the caller folds them node-level so a dropped + binding gap is not lost just because a zero-delay join edge from a surviving + pred remains. + """ + kept_edges: list[StaticEdge] = [] + has_surviving_pred: set[str] = set() + residual_by_target: dict[str, float] = {} + for edge in edges: + if not isinstance(edge, StaticEdge): + continue + src, tgt = edge.source, edge.target + if tgt not in survivors: + continue + if src == START_NODE_ID or src in survivors: + kept_edges.append(edge) + if src != START_NODE_ID: + has_surviving_pred.add(tgt) + continue + # src was dropped (executed / pre-t* history): fold its recorded gap, + # minus the wall time already waited since its return, into the + # target's re-root offset. The executor gate takes the max of a node's + # incoming delays, so AND-fan-in residuals max-combine here too. + # END-anchored quantities only: the ledger wall is src's RETURN, so a + # dispatch-anchored delay (start / first-token) debited from a + # return-anchored elapsed would over-delay by src's live service time + # -- those edges burst (residual 0). + recorded_us = max( + edge.delay_after_predecessor_us or 0.0, + edge.min_start_delay_us or 0.0, + ) + wall = return_wall_us.get(src) + residual = ( + max(0.0, recorded_us - max(0.0, drain_end_wall_us - wall)) + if wall is not None + else 0.0 + ) + if residual_cap_us is not None: + residual = min(residual, residual_cap_us) + residual_by_target[tgt] = max(residual_by_target.get(tgt, 0.0), residual) + + new_edges = [ + e + for e in kept_edges + if not (e.source == START_NODE_ID and e.target not in has_surviving_pred) + ] + kept_pred_residuals: dict[str, float] = {} + for nid in survivors: + if nid not in has_surviving_pred: + new_edges.append( + StaticEdge( + source=START_NODE_ID, + target=nid, + min_start_delay_us=residual_by_target.get(nid, 0.0), + ) + ) + else: + residual = residual_by_target.get(nid, 0.0) + if residual > 0.0: + kept_pred_residuals[nid] = residual + return new_edges, kept_pred_residuals + + +__all__ = ["chop_trie_at_tstar", "chop_trie_at_frontier"] diff --git a/src/aiperf/timing/strategies/cache_bust.py b/src/aiperf/timing/strategies/cache_bust.py new file mode 100644 index 0000000000..506b5e35a8 --- /dev/null +++ b/src/aiperf/timing/strategies/cache_bust.py @@ -0,0 +1,170 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Deterministic per-session cache-bust marker builder for graph snapshot replay. + +Same ``(benchmark_id, recycle_pass, lane_index, trace_id)`` always yields the +same digest - reproducible across reruns. Adding ``trace_id`` to the digest +input ensures every ``(recycle_pass, lane, trace)`` combination is unique by +construction: without it, two different traces landing on the same +``(recycle_pass, lane)`` tuple would collide on one marker. + +Byte-compatible with the agentx ``timing/strategies/cache_bust.py`` digest +(``sha256(f"{benchmark_id}:{recycle_pass}:{trajectory_index}:{trace_id}")[:12]``, +wrapped as ``[rid:]\\n\\n`` for the prefix target). The graph snapshot +path uses only the ``FIRST_TURN_PREFIX`` / ``NONE`` variants. +""" + +from __future__ import annotations + +import hashlib +from typing import Any + +from aiperf.common.enums import CacheBustTarget + +_DIGEST_LEN = 12 # 12 hex chars = 48 bits, ample for in-run uniqueness + + +def build_cache_bust_marker( + benchmark_id: str, + recycle_pass: int, + lane_index: int, + trace_id: str, + *, + target: CacheBustTarget, +) -> str | None: + """Render the marker text for the given inputs and target position. + + The digest tuple is intentionally phase-agnostic. A lane's WARMUP priming + turn and its first PROFILING turn must share the same marker so the warmup + KV-cache work transfers to profiling; adding phase to the digest would + defeat that - keep it out. + + Returns ``None`` when target is ``NONE`` so callers can treat "no marker" + uniformly -- ``inject_marker_at_first_user_message`` no-ops on a falsy + marker (returning ``""`` would introduce a third "no marker" value + distinct from ``None``). + + Args: + benchmark_id: Run-scoped salt (the source's ``random_seed`` stringified) + so two runs with different seeds mint different markers. + recycle_pass: Never-restarting per-lane pass counter; rotates the digest + on every recycle so a recycled lane can never collide with a warmed + digest from an earlier pass. + lane_index: Absolute lane ordinal; decorrelates concurrent lanes that + sampled the same trace. + trace_id: The lane's sampled conversation/trace id; makes the digest + unique per trace within a ``(recycle_pass, lane)`` slot. + target: Injection position. ``NONE`` returns ``None``; + ``FIRST_TURN_PREFIX`` wraps the digest with a trailing blank line. + """ + if target == CacheBustTarget.NONE: + return None + + unique_str = f"{benchmark_id}:{recycle_pass}:{lane_index}:{trace_id}" + digest = hashlib.sha256(unique_str.encode()).hexdigest()[:_DIGEST_LEN] + rid = f"[rid:{digest}]" + return f"{rid}\n\n" + + +def build_trace_instance_marker( + benchmark_id: str, + trace_instance_id: str, + *, + target: CacheBustTarget, +) -> str | None: + """Mint the per-TRACE-INSTANCE cache-bust marker for the graph-IR replay path. + + Matches agentx's per-trajectory-TREE scoping (``cache_bust.py::resolve_tree_marker`` + keys the marker on ``root_correlation_id`` and shares ONE value across every + member of the tree -- main turns, subagents, flat agents). On our graph path + the trajectory-tree analog is the trace INSTANCE id (``credit.trace_id``, e.g. + ``t-1#0``): every dispatch of one instance carries the same ``trace_id`` (the + adapter pins ``credit.trace_id`` to the root instance for nested/subagent + dispatches too), so digesting it yields ONE marker shared across all of the + instance's turns/dispatches. Distinct instances (``t-1#0`` vs ``t-2#0``) get + distinct markers, and a RECYCLED template (a new session slot, e.g. ``t-1#1``) + is a fresh instance id -> a fresh marker, mirroring agentx's per-recycle + ``recycle_pass`` bump. + + The recycle ordinal is already baked into the instance id (the ``#N`` suffix), + so the digest tuple's ``recycle_pass`` / ``lane_index`` slots are pinned to + ``0`` and the instance id fills the ``trace_id`` slot. The rendered marker is + byte-identical to :func:`build_cache_bust_marker` -- a ``[rid:<12hex>]\\n\\n`` + prefix whose digest is ``sha256(f"{benchmark_id}:0:0:{trace_instance_id}")[:12]``. + + Args: + benchmark_id: Run-scoped salt so two runs mint different markers. + trace_instance_id: The trace INSTANCE id (``credit.trace_id``); shared by + every dispatch of the instance, distinct per instance, regenerated on + recycle. + target: Injection position. ``NONE`` returns ``None`` (no marker); + ``FIRST_TURN_PREFIX`` renders the prefix marker. + + Returns: + The rendered marker text, or ``None`` when ``target`` is ``NONE``. + """ + return build_cache_bust_marker( + benchmark_id, + recycle_pass=0, + lane_index=0, + trace_id=trace_instance_id, + target=target, + ) + + +def _content_has_marker_prefix(content: Any, marker: str) -> bool: + """Whether ``content`` already begins with ``marker`` (idempotency guard). + + Mirrors agentx ``worker.py::_content_has_marker_at_edge`` for the prefix + case: a plain string is checked with ``startswith``; an OpenAI multimodal + list-of-parts is checked against a leading ``{"type": "text", "text": ...}`` + part. Any other content shape is treated as "no marker present". + """ + if isinstance(content, str): + return content.startswith(marker) + if isinstance(content, list) and content: + return content[0] == {"type": "text", "text": marker.strip()} + return False + + +def inject_marker_at_first_user_message( + messages: list[dict[str, Any]], marker: str | None +) -> None: + """Prepend ``marker`` to the first ``role == "user"`` message, in place. + + Mirrors agentx ``worker.py::_inject_marker_into_first_user_turn`` for the + ``FIRST_TURN_PREFIX`` target: walk the wire ``messages`` forward, find the + first user-role message, and prepend the marker to its content. Plain-string + content becomes ``marker + content``; OpenAI multimodal list content gets a + leading ``{"type": "text", "text": marker.strip()}`` part. The injection is + idempotent (re-stamping the same marker is a no-op) and stamps ONLY the first + user message. No-op when ``marker`` is falsy (``NONE`` target) or no + user-role message exists. + + Args: + messages: The materialized wire ``messages`` list (mutated in place). + marker: The rendered marker text, or ``None`` to stamp nothing. + """ + if not marker: + return + for idx, msg in enumerate(messages): + if not (isinstance(msg, dict) and msg.get("role") == "user"): + continue + content = msg.get("content", "") + if _content_has_marker_prefix(content, marker): + return + if isinstance(content, str): + messages[idx] = {**msg, "content": marker + content} + return + if isinstance(content, list): + marker_part = {"type": "text", "text": marker.strip()} + messages[idx] = {**msg, "content": [marker_part, *content]} + return + return + + +__all__ = [ + "build_cache_bust_marker", + "build_trace_instance_marker", + "inject_marker_at_first_user_message", +] diff --git a/src/aiperf/timing/strategies/core.py b/src/aiperf/timing/strategies/core.py index d8fe45d7bf..a3526cf46b 100644 --- a/src/aiperf/timing/strategies/core.py +++ b/src/aiperf/timing/strategies/core.py @@ -29,6 +29,11 @@ class TimingStrategyProtocol(Protocol): 3. execute_phase(): Send first turns (main timing loop) 4. handle_credit_return(): Handle credit return, dispatch next turn if needed + Optional hooks live on sibling protocols so strategies that do not need + them are unaffected: ``PhaseTeardownStrategyProtocol.teardown_phase`` (the + PhaseRunner-invoked post-phase cleanup), ``CreditResultAwareStrategyProtocol``, + and ``RateSettableProtocol``. + Fresh strategy instances are created per-phase by PhaseRunner. All dependencies are injected via __init__ for clean, testable design. """ @@ -72,9 +77,10 @@ async def handle_credit_return(self, credit: Credit) -> None: should be sent, and if so, dispatches it via the appropriate path (immediate, scheduled, or queued). - Note: CreditCallbackHandler checks can_send_any_turn() before calling. - Implementations only need to check conversation-specific conditions - (e.g., is_final_turn). + Note: CreditCallbackHandler gates this on can_send_any_turn() for root + credits; DAG child non-final returns are always delivered so pending + joins can drain. Implementations only need to check + conversation-specific conditions (e.g., is_final_turn). Args: credit: Completed credit with conversation/turn info @@ -91,6 +97,26 @@ async def handle_credit_result(self, credit_return: CreditReturn) -> None: ... +@runtime_checkable +class PhaseTeardownStrategyProtocol(Protocol): + """Optional teardown hook for strategies holding phase-scoped resources. + + ``PhaseRunner`` awaits ``teardown_phase()`` in a ``finally`` once the + phase's execute -> sending-complete -> returning-complete pipeline settles + (deferred to the background return-wait completion for non-final seamless + phases, whose returns are still in flight when ``run()`` exits). Strategies + that do not define the method are skipped by the runner -- the default is a + no-op -- so the base ``TimingStrategyProtocol`` surface is unchanged for + the linear strategies. ``GraphIRReplayStrategy`` uses this to detach its + graph-return / first-token observers and close retained sticky trace + lifecycles so a subsequent phase never routes into a torn-down registry. + """ + + async def teardown_phase(self) -> None: + """Release phase-scoped resources after the phase completes.""" + ... + + @runtime_checkable class FirstTokenAwareStrategyProtocol(Protocol): """Optional hook for strategies that need TTFT samples.""" @@ -124,3 +150,19 @@ def set_request_rate(self, rate: float) -> None: rate: New request rate in requests per second (must be > 0). """ ... + + +@runtime_checkable +class LaneSettableProtocol(Protocol): + """Protocol for strategies whose CONCURRENT-LANE admission can be ramped. + + The graph replay strategy runs ``--concurrency`` recycling lanes; a + session-slot ramp cannot throttle them (graph credits bypass session + slots), so the concurrency ramper drives this setter instead: lanes above + the live limit park before their first instance and are admitted as the + limit rises. + """ + + def set_lane_limit(self, limit: int) -> None: + """Update the number of admitted lanes (clamped to [1, concurrency]).""" + ... diff --git a/src/aiperf/timing/strategies/fixed_schedule.py b/src/aiperf/timing/strategies/fixed_schedule.py index a557130555..49b2d8a243 100644 --- a/src/aiperf/timing/strategies/fixed_schedule.py +++ b/src/aiperf/timing/strategies/fixed_schedule.py @@ -100,6 +100,9 @@ async def setup_phase(self) -> None: x_correlation_id=str(uuid.uuid4()), turn_index=0, num_turns=len(conv.turns), + # Finality stamping must see a spawning turn as + # non-final (any-mode branch flag). + has_branches=bool(conv.turns[0].branch_ids), ), ) ) @@ -151,9 +154,11 @@ async def handle_credit_return( if credit.is_final_turn: return - # This contains the delay_ms or timestamp_ms for the next turn + # This contains the delay_ms or timestamp_ms for the next turn. + # Passed through so has_forks / has_branches derive from the real + # metadata (a spawning final turn must never stamp is_tree_final). next_meta = self._conversation_source.get_next_turn_metadata(credit) - turn = TurnToSend.from_previous_credit(credit) + turn = TurnToSend.from_previous_credit(credit, next_meta) if next_meta.timestamp_ms is not None: self._scheduler.schedule_at_perf_sec( diff --git a/src/aiperf/timing/strategies/graph_ir_replay.py b/src/aiperf/timing/strategies/graph_ir_replay.py new file mode 100644 index 0000000000..be04ac44c0 --- /dev/null +++ b/src/aiperf/timing/strategies/graph_ir_replay.py @@ -0,0 +1,2354 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GraphIRReplayStrategy — the weka IR trace runner on the v1 credit pipeline. + +This timing strategy drives the dataflow ``TraceExecutor`` once per +weka trace and OWNS both completion and concurrency, rather than letting the +linear ``CreditCounter`` / session-slot arithmetic decide. A fan-out trace has +no linear acquire/release; graph credits flow through real issuance -> worker -> +records -> returns but BYPASS the session-slot lifecycle via +``CreditIssuer.issue_graph_credit``. + +Ownership model +--------------- +* **Completion.** AgentX-faithful (``aiperf.timing.phase.stop_conditions``): the + phase finalizes on the STOP CONDITION (session count / request count / + duration), NOT on every admitted trace fully draining its faithful idle-gap + replay. ``execute_phase``'s ``finally`` always signals sending-complete, and a + ``--benchmark-duration`` cancels in-flight executors when the budget elapses + (see :meth:`GraphIRReplayStrategy._run_traces_under_duration_budget`). We do + NOT consult ``CreditCounter.is_final_turn`` (it cannot express a DAG). +* **Concurrency.** A fixed pool of ``max_concurrent_traces`` lanes (see + :meth:`GraphIRReplayStrategy._run_lanes` / ``_resolve_lane_count``) admits + traces: each lane runs one trace instance at a time and recycles onto the next + wrap template when the stop-condition gate still admits a new session. Within a + trace the executor self-gates via its dataflow channels -- the lane pool is the + only cross-trace bound (prevents parked-Future explosion). + +Return routing +-------------- +The strategy installs ONE graph-return observer on the ``CreditCallbackHandler`` +(via ``register_observer``) that routes each ``(credit, error, cancelled)`` to +the owning per-trace ``CreditDispatchAdapter`` by ``credit.trace_id``. It fires +UNCONDITIONALLY (before the gated ``handle_credit_return`` path), so an in-flight +dispatch Future is always resolved even after the phase can no longer send. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import time +import uuid +from typing import TYPE_CHECKING, Any + +import msgspec + +from aiperf.common.aiperf_logger import AIPerfLogger +from aiperf.common.constants import MICROS_PER_SECOND +from aiperf.common.environment import Environment +from aiperf.common.mixins import AIPerfLoggerMixin +from aiperf.common.scenario.base import TrajectoryWarmupFailedError +from aiperf.dataset.graph.graph_path_catalog import build_catalog_context +from aiperf.dataset.graph.models import START_NODE_ID, StaticEdge +from aiperf.graph.credit_dispatch_adapter import ( + CreditDispatchAdapter, + CreditIssueRefusedError, +) +from aiperf.graph.executor import TraceExecutor +from aiperf.graph.scheduler import collapse_leading_start_offsets +from aiperf.timing.graph_ir_source import GraphIRConversationSource, GraphTrace +from aiperf.timing.graph_ir_trace_view import parsed_for_trace +from aiperf.timing.graph_warmup_handoff import ( + GraphWarmupHandoff, + LaneHandoff, +) +from aiperf.timing.snapshot_chop import chop_trie_at_frontier, chop_trie_at_tstar + +if TYPE_CHECKING: + from collections.abc import Callable + + from aiperf.common.enums import CacheBustTarget + from aiperf.credit.messages import FirstToken + from aiperf.credit.structs import Credit + from aiperf.dataset.graph.models import ( + GraphRecord, + LlmNode, + ParsedGraph, + TraceRecord, + ) + from aiperf.plugin.enums import DatasetSamplingStrategy + from aiperf.timing.graph_channel import GraphPhaseChannel + +_logger = AIPerfLogger(__name__) + +__all__ = ["GraphIRReplayStrategy", "first_token_sources", "rewrite_for_warmup"] + + +def first_token_sources(graph: GraphRecord) -> frozenset[str]: + """Source node ids of every first-token-anchored ``StaticEdge`` in ``graph``. + + A ``StaticEdge`` carrying ``delay_after_predecessor_first_token_us`` anchors + its successor to the SOURCE node's observed first token (post-TTFT anchoring, + validator rule 55). That source must therefore emit a ``FirstToken`` -- + i.e. its issued turn carries ``first_token_event=True`` -- so the runtime can + release the successor when the token arrives. Returns the (possibly empty) + set of such source node ids; a gap-free / pre-anchoring graph yields the + empty set (no node emits a first-token event). + """ + return frozenset( + edge.source + for edge in getattr(graph, "edges", []) + if isinstance(edge, StaticEdge) + and edge.delay_after_predecessor_first_token_us is not None + ) + + +def _chain_key(node_id: str) -> str: + """Return the per-session chain key a trie node id belongs to. + + Both live trie producers mint node ids as ``{chain_prefix}_{ordinal}`` -- + one linear chain per recorded session: the weka walk emits ``{trace_id}:{k}`` + for the root request list and ``r_1_0``/``r_1_1`` for the subagent spawned + at index 1 (``adapters/weka/trie_build._walk``), and the dynamo lowering + emits ``{session_id}:{k}`` per session + (``adapters/dynamo/trie_lowering``). Stripping the final ``_``-delimited + token therefore recovers the enclosing session chain (root chain + each + subagent chain) -- the same per-session decomposition the trie build's + recorded structure encodes. Chain identity must come from node ids because + the sidecar-loaded timing plane strips ``metadata["trie"]`` contents (hash + ids / segment paths), leaving ids + edges as the only runtime chain signal, + and the interval-order edges are cross-chain ordering edges, not session + boundaries. A node id with no ``_`` forms a defensive singleton chain. + """ + prefix, sep, _ = node_id.rpartition("_") + return prefix if sep else node_id + + +def _warmup_boundary_nodes(graph: GraphRecord, t_star_us: float) -> dict[str, LlmNode]: + """Return ``{node_id: node}`` of each chain-live-at-``t*`` boundary turn. + + Chains are the per-session linear paths the trie node ids encode + (:func:`_chain_key`), ordered by recorded arrival. A chain is LIVE when it + has BOTH a node arriving before ``t*`` and a node arriving at/after ``t*``; + its boundary is the LAST pre-``t*`` node. Chains with no pre-``t*`` node + need no priming (profiling replays them from their own start); chains + entirely pre-``t*`` are not live (nothing of them is profiled). + """ + chains: dict[str, list[tuple[int, str]]] = {} + for nid, node in graph.nodes.items(): + arrival = getattr(node, "arrival_offset_us", None) or 0 + chains.setdefault(_chain_key(nid), []).append((arrival, nid)) + boundary: dict[str, LlmNode] = {} + for members in chains.values(): + members.sort() + pre = [nid for arrival, nid in members if arrival < t_star_us] + if pre and any(arrival >= t_star_us for arrival, _ in members): + boundary[pre[-1]] = graph.nodes[pre[-1]] + return boundary + + +def rewrite_for_warmup(parsed: ParsedGraph, t_star_us: float) -> ParsedGraph: + """Rewrite ``parsed`` into the WARMUP boundary-priming graph at ``t*``. + + AgentX-parity contract (``timing.config._build_graph_auto_warmup_config``): + warmup dispatches exactly ONE priming credit per chain LIVE at ``t*`` -- + the chain's boundary turn, the last node of that per-session chain whose + recorded arrival precedes ``t*`` (:func:`_warmup_boundary_nodes`). Because + trie prompts are cumulative along a chain, priming the boundary turn's + prompt (at the worker-side warmup ``max_tokens`` cap, keyed off the + ``"warmup"`` phase variant) warms the chain's whole prefix. + + The produced graph is FLAT: only the boundary nodes survive, each re-rooted + from ``START`` with NO leading offset (warmup bursts every priming credit + at phase start rather than replaying recorded gaps) and with fan-in + ``inputs`` cleared (their predecessors are gone). Node identity, the trie + envelope, and ``dispatch_overrides`` are preserved so the worker resolves + the unmodified catalog ordinal and materializes the exact recorded prompt. + ``t_star_us <= 0`` (full native replay, or a zero-duration trace) yields an + EMPTY graph so the warmup phase finalizes immediately. + """ + graph = parsed.graph + boundary = _warmup_boundary_nodes(graph, t_star_us) if t_star_us > 0 else {} + new_nodes = { + nid: msgspec.structs.replace(node, inputs=[], min_start_delay_us=None) + for nid, node in boundary.items() + } + new_edges = [StaticEdge(source=START_NODE_ID, target=nid) for nid in new_nodes] + new_graph = msgspec.structs.replace(graph, nodes=new_nodes, edges=new_edges) + return msgspec.structs.replace(parsed, graph=new_graph) + + +def _leaf_credit_refusal(exc: BaseException) -> CreditIssueRefusedError | None: + """Return the refusal iff ``exc`` is (a group of) ONLY issuer refusals. + + The per-trace executor TaskGroup wraps node failures in an + ``ExceptionGroup``, so a healthy stop-gate refusal + (``CreditIssueRefusedError``: request-count / duration cap reached, or run + cancelled) may surface bare OR grouped -- and multiple concurrent node + coroutines may each carry one. A mixed group (a refusal alongside a genuine + error) is NOT a clean stop and returns ``None`` so the caller keeps the + error path. + """ + if isinstance(exc, CreditIssueRefusedError): + return exc + if isinstance(exc, BaseExceptionGroup): + matched, rest = exc.split(CreditIssueRefusedError) + if matched is not None and rest is None: + leaf: BaseException = matched.exceptions[0] + while isinstance(leaf, BaseExceptionGroup): + leaf = leaf.exceptions[0] + return leaf # type: ignore[return-value] + return None + + +def _seed_for_draw_pass(base_seed: int, pass_index: int) -> int: + """Derive a per-pass RNG seed for the shuffle draw (matches the t* salt). + + Mirrors :func:`aiperf.timing.graph_ir_source._seed_for_trace_lane`: SHA-256 + over ``f"{base_seed}:dataset-draw:{pass_index}"`` and take the low 8 bytes, + so each recycle pass re-permutes under a distinct-yet-deterministic seed + derived from the run's ``t_star_random_seed``. Same base seed + pass index + always yields the same permutation (cross-run reproducibility), while + different passes decorrelate -- the same order-independent SHA-256 derivation + the conversation-plane samplers use via ``rng.derive``. + """ + digest = hashlib.sha256(f"{base_seed}:dataset-draw:{pass_index}".encode()).digest() + return int.from_bytes(digest[:8], "big") + + +class GraphIRReplayStrategy(AIPerfLoggerMixin): + """Drive a ``TraceExecutor`` per weka trace over the v1 credit pipeline. + + Constructed per-phase by ``PhaseRunner._build_strategy`` with the standard + timing-strategy kwargs PLUS the graph-only injection channel + (``parsed_graph`` + ``register_observer``); see the module docstring for the + ownership contract. + """ + + def __init__( + self, + *, + config: Any = None, + conversation_source: Any = None, + scheduler: Any = None, + stop_checker: Any = None, + credit_issuer: Any, + lifecycle: Any = None, + parsed_graph: ParsedGraph, + register_observer: Callable[[Any], None], + register_first_token_observer: Callable[[Any], None] | None = None, + unregister_observer: Callable[[Any], None] | None = None, + unregister_first_token_observer: Callable[[Any], None] | None = None, + max_concurrent_traces: int | None = None, + dispatch_timeout_s: float | None = None, + start_min_ratio: float = 0.0, + start_max_ratio: float = 0.0, + t_star_random_seed: int = 0, + burst_phase_starts: bool = False, + cache_pressure_duration_s: float | None = None, + warmup_handoff: GraphWarmupHandoff | None = None, + graph_channel: GraphPhaseChannel | None = None, + dataset_sampling_strategy: DatasetSamplingStrategy | None = None, + allow_dataset_wrap: bool | None = None, + cache_bust: CacheBustTarget | None = None, + **kwargs: Any, + ) -> None: + """Initialize the graph trace runner. + + Args: + config: Per-phase ``CreditPhaseConfig`` (concurrency bound + stop + thresholds: ``expected_num_sessions`` / ``total_expected_requests``). + conversation_source: Accepted for timing-strategy signature parity; + unused by graph runs (``None``), which carry state on + ``graph_channel`` instead. + scheduler: Unused (loop scheduler); protocol parity. + stop_checker: Unused here; the issuer honors the stop gate itself. + credit_issuer: The real ``CreditIssuer``; adapters call its + ``issue_graph_credit`` (session-slot-bypassing) path. + lifecycle: Phase ``PhaseLifecycle``; read for the duration stop + condition (``time_left_in_seconds()``) bounding the dispatch. + parsed_graph: The built ``ParsedGraph`` whose ``traces`` this phase + replays (the R3 injection channel; see ``PhaseRunner``). + register_observer: Callback installing the single graph-return + observer on the shared ``CreditCallbackHandler``. + register_first_token_observer: Callback installing the single + graph-first-token observer on the shared ``CreditCallbackHandler`` + (post-TTFT anchoring). ``None`` (default, e.g. unit harness) + skips first-token routing entirely -- anchored edges then fall + back to their dispatch-time delay. + unregister_observer: Compare-and-clear detach for the graph-return + observer (``CreditCallbackHandler.clear_graph_return_observer``): + teardown passes its OWN observer and the handler clears the + shared slot only if that observer is still installed. Required + for seamless multi-phase runs where a stale phase's deferred + teardown fires after the next phase installed its observer. + ``None`` (unit harness) falls back to ``register_observer(None)``. + unregister_first_token_observer: Same compare-and-clear detach for + the first-token observer slot. ``None`` falls back to + ``register_first_token_observer(None)``. + max_concurrent_traces: Trace-admission bound; defaults to the phase + ``concurrency`` else ``1`` (the plain aiperf default). + dispatch_timeout_s: Per-dispatch deadlock guard (adapter default). + start_min_ratio: Lower bound (fraction of duration) of the t* window. + Default ``0.0``; together with ``start_max_ratio=0.0`` this + selects full native replay (t*=0, no snapshot rewrite). The + AgentX 0.0..1.0 window is scenario-applied, not a bare default. + start_max_ratio: Upper bound of the t* window. Default ``0.0``. + t_star_random_seed: Base seed for per-trace t* sampling (trace-salted). + burst_phase_starts: AgentX ``--burst-phase-starts`` parity. False + (default) = SPREAD: each trace's first firing keeps its recorded + (warped, <=cap) leading offset from t*. True = BURST: the leading + per-firing dispatch offset is zeroed so every trace's earliest + profiling firing (and every warmup priming credit) fires at + phase-time 0; relative inter-turn delays after the first are + still honored. Governs ONLY the phase-start dispatch pattern. + cache_pressure_duration_s: Extended (cache-pressure) warmup stage + budget in seconds (``float | None``). Non-None on a WARMUP phase + arms the pressure stage: after boundary priming drains, the + post-t* graphs replay compressed for this many seconds, then + drain into the profiling handoff. None (default) disables it, + keeping warmup byte-identical. + warmup_handoff: The consume-once ``GraphWarmupHandoff | None`` a prior + extended warmup stashed on the shared channel, threaded into the + first PROFILING graph phase so each lane resumes at its recorded + execution frontier. None (default) means no prior extended warmup. + graph_channel: The typed ``GraphPhaseChannel`` threaded from the + orchestrator; a WARMUP phase stashes its ``warmup_handoff`` here + at teardown for the next PROFILING phase to consume. None in unit + harnesses that construct the strategy directly. + dataset_sampling_strategy: Resolved run-level dataset sampling + strategy (from ``run.resolved.dataset_sampling_strategy`` via the + CreditPhaseConfig). Consumed by the per-lane trace draw + (``_draw_index``): SHUFFLE/RANDOM remap the draw through a seeded + permutation; SEQUENTIAL/None keep the byte-identical cursor. + None for non-graph phases / until resolution derives it. + allow_dataset_wrap: Resolved graph-plane dataset-wrap policy (from + ``run.resolved.allow_dataset_wrap`` via the CreditPhaseConfig). + Consumed by the setup-phase wrap-guard (over-subscription raises + unless wrapping is allowed). None until derived. + cache_bust: Resolved per-trace-instance cache-bust target (from + ``endpoint.cache_bust`` via the CreditPhaseConfig). Consumed by + the dispatch duplication report: when lanes recycle (factor > 1) + with cache-bust OFF (``None`` / ``NONE``), clones collide on + identical prefixes, so the report warns; ON it stays quiet. + """ + super().__init__(logger_name="GraphIRReplayTiming") + self._config = config + self._credit_issuer = credit_issuer + self._stop_checker = stop_checker + self._parsed = parsed_graph + self._register_observer = register_observer + self._register_first_token_observer = register_first_token_observer + self._unregister_observer = unregister_observer + self._unregister_first_token_observer = unregister_first_token_observer + self._dispatch_timeout_s = dispatch_timeout_s + self._burst_phase_starts = bool(burst_phase_starts) + self._graph_channel = graph_channel + self._cache_pressure_duration_s = cache_pressure_duration_s + self._warmup_handoff = warmup_handoff + # Resolved dataset-selection policy threaded from ``run.resolved`` via + # the CreditPhaseConfig. Consumed by the setup-phase wrap-guard and the + # per-lane sampling draw (``_draw_index``). + self._dataset_sampling_strategy = dataset_sampling_strategy + self._allow_dataset_wrap = allow_dataset_wrap + # Resolved cache-bust target threaded from ``endpoint.cache_bust`` via the + # CreditPhaseConfig. Consumed by the dispatch duplication report below. + self._cache_bust = cache_bust + # Duration stop condition (AgentX parity): ``time_left_in_seconds()`` + # bounds the dispatch; ``None`` with no ``--benchmark-duration``. + self._lifecycle = lifecycle + + self._start_min_ratio = start_min_ratio + self._start_max_ratio = start_max_ratio + self._t_star_random_seed = t_star_random_seed + self._init_tstar_source( + parsed_graph, start_min_ratio, start_max_ratio, t_star_random_seed + ) + self._init_accelerated_warmup() + + self._max_concurrent = self._resolve_concurrency(max_concurrent_traces) + # Live lane-admission limit (LaneSettableProtocol): the concurrency + # ramper drives it 1 -> _max_concurrent; without a ramp every lane is + # admitted immediately. The event is swapped fresh on every raise so + # parked lanes re-check without lost wakeups (single-threaded asyncio: + # a waiter's check-then-await pair cannot be preempted by the setter). + self._lane_limit = self._max_concurrent + self._lane_limit_raised = asyncio.Event() + self._completed_traces = 0 + self._errored_traces = 0 + self._admitted_traces = 0 + # Total instances dispatched off the finite pool (lane starts + every + # serial recycle) for the dispatch duplication report. An instance + # attribute so the report fires on BOTH natural completion AND + # duration-cancel (the ``_run_lanes`` local would be lost when the lane + # future is cancelled by the duration budget). + self._instances_started = 0 + # Loop-clock deadline for the ACTIVE duration-budgeted stage (phase + # duration or pressure duration); None outside one. The budget wrappers + # cancel the lane fan-out task on timeout, but cancellation delivery + # through a TaskGroup whose children complete constantly is unreliable + # on Python 3.11 (the cancel can be lost mid-abort), so the lane + # recycle loops ALSO check this deadline cooperatively -- the stage + # then halts on time even when the cancel never lands. + self._duration_deadline: float | None = None + # Per-(template_trace_id, lane_index) t* plan cache. AgentX seeds t* per + # ``(trace_id, lane)`` so the same template recurring across lanes (or + # recycled onto one lane) resumes at a DIFFERENT t*. Built lazily so a + # large concurrency over a small corpus only plans the lanes it runs. + self._lane_plans: dict[tuple[str, int], GraphTrace] = {} + # Per-lane t* source cache (catalog/namespace map are lane-independent), + # so a fan-out reuses one ``GraphIRConversationSource`` per lane instead + # of rebuilding the corpus catalog on every per-trace plan. + self._lane_sources: dict[int, GraphIRConversationSource] = {} + # Per-(total, pass) shuffled trace-index permutation cache for the + # ``--dataset-sampling-strategy`` draw (:meth:`_draw_index`). Built lazily + # once per pass and reused for every draw in that pass so repeated draws + # are cheap and consistent. A single event loop mutates it with no await + # between read and write, so no lock is needed. Empty / unused under the + # default SEQUENTIAL draw. + self._draw_perm_cache: dict[tuple[int, int], list[int]] = {} + + def _init_tstar_source( + self, + parsed_graph: ParsedGraph, + start_min_ratio: float, + start_max_ratio: float, + t_star_random_seed: int, + ) -> None: + """Build the t* snapshot source, per-trace t* plans, and node catalog. + + The t* source samples a per-trace snapshot instant and partitions nodes + into warmup (history) / profiling. Default ratios [0, 0] => t*=0 => full + native replay (identity rewrite), so the profiling path is byte-identical + unless a caller supplies a positive window. + """ + self._source = GraphIRConversationSource( + parsed=parsed_graph, + start_min_ratio=start_min_ratio, + start_max_ratio=start_max_ratio, + random_seed=t_star_random_seed, + ) + # ``{trace_id: GraphTrace}`` -- the lane-0 per-trace t* plan (the default + # single-pass disposition). Lanes > 0 and recycle passes resolve their own + # lane-salted plan lazily via :meth:`_plan_for_lane`. + self._plans = {gt.trace_id: gt for gt in self._source.iter_traces()} + self._catalog = build_catalog_context(parsed_graph) + # ``{instance_id: CreditDispatchAdapter}`` -- the de-mux registry the + # single return observer routes by ``credit.trace_id`` (the per-recycle + # instance id, e.g. ``t-1#0``), so two concurrent instances of one + # template never collide. + self._adapters: dict[str, CreditDispatchAdapter] = {} + # Instance ids whose parent run has finished; the adapter is popped here + # (not earlier) so a mid-run ``on_drained`` callback never reaps it while + # the executor is still firing. Serialized with ``_registry_lock`` so the + # parent finally and a late ``on_drained`` callback never race the pop. + self._parent_done: set[str] = set() + # ``{template_trace_id: frozenset[node_id]}`` -- the sources of each + # trace's first-token-anchored edges (post-TTFT anchoring), computed once + # per template from its OWN projected graph and reused across lanes / + # recycles (the set is lane/t*-independent: a t* chop only drops edges). + self._first_token_sources_cache: dict[str, frozenset[str]] = {} + # ``{template_trace_id: node_id -> (agent_depth, parent_node_id)}`` -- + # the dag_jsonl lowering's ``metadata["dag"]`` legacy-identity stamp, + # read once per template from its OWN projected graph. ``None`` for + # stamp-free producers (weka/dynamo) so the adapter's root-chain + # fallback stays byte-identical. + self._node_identity_cache: dict[ + str, dict[str, tuple[int, str | None]] | None + ] = {} + self._registry_lock = asyncio.Lock() + # Background deferred-release tasks scheduled from the sync ``on_drained`` + # callback; tracked so none is garbage-collected mid-flight. + self._release_tasks: set[asyncio.Task[None]] = set() + # Background GraphTraceEnd sends scheduled from the sync adapter-reap + # path; tracked so none is garbage-collected mid-flight. + self._trace_end_tasks: set[asyncio.Task[None]] = set() + + def _init_accelerated_warmup(self) -> None: + """Decide whether this phase runs the accelerated cache-pressure warmup. + + Active ONLY when this is the WARMUP phase AND a cache-pressure + duration is configured (``--agentic-cache-warmup-duration``): every + WARMUP ``TraceExecutor`` then builds with ``compress_edge_delays=True`` + (zero-idle replay) to drive the server KV cache to pressure. + PROFILING and the warmup output cap are untouched; without a duration + warmup honors every captured edge delay. + """ + from aiperf.common.enums import CreditPhase + + phase = getattr(self._config, "phase", None) + self._accelerated_warmup = bool( + phase == CreditPhase.WARMUP and self._cache_pressure_duration_s is not None + ) + # Extended (cache-pressure) warmup stage state. The ledger records + # every NON-CANCELLED warmup-phase return (priming + pressure) so the + # teardown handoff can anchor residual delays; cancelled returns are + # excluded (not executed -- profiling refires them). Keyed + # {instance_id: {node_id: wall_us}} on the _wall_us() monotonic clock. + self._pressure_enabled = bool( + phase == CreditPhase.WARMUP and self._cache_pressure_duration_s is not None + ) + self._pressure_active = False + self._pressure_live: dict[int, tuple[str, str, float, int]] = {} + # (template_id, lane) -> the warmup boundary-PRIMING instance id run on + # that lane. Instance ids carry fresh nonces (no longer recomputable), + # so the pressure handoff's pass-0 priming-wall merge looks the id up + # here instead of reconstructing it. + self._priming_instance_ids: dict[tuple[str, int], str] = {} + self._return_walls: dict[str, dict[str, float]] = {} + self._ordinal_to_node: dict[str, dict[int, str]] = {} + # Next corpus draw index for the pressure recycle loop; stashed into the + # handoff so profiling's bounded recycle continues from it. A single + # event loop mutates it, so an instance attr has the same atomicity the + # prior ``nonlocal`` closure had, and it survives past the lane loop for + # the teardown stash. + self._pressure_next_index = 0 + # Number of pressure lanes this warmup fanned out (== len(pass0_traces) + # in _run_pressure_lanes). Stashed into the handoff so profiling can + # fresh-start any drained-empty lane below this count instead of + # re-running a t* resume the pressure stage already executed. + self._pressure_lane_count = 0 + # Warmup-failure abort state (agentx parity): a WARMUP phase that + # recorded terminal request failures aborts the run before profiling + # (see :meth:`report_warmup_failures`). Gated on the phase, not on the + # pressure stage, so plain boundary priming is covered too. + self._is_warmup_phase = bool(phase == CreditPhase.WARMUP) + self._warmup_failure_count = 0 + self._warmup_failure_samples: list[str] = [] + + @property + def accelerated_warmup(self) -> bool: + """True when this phase runs the accelerated cache-pressure warmup.""" + return self._accelerated_warmup + + @property + def completed_traces(self) -> int: + """Number of traces whose executor run finished (success OR error).""" + return self._completed_traces + + @property + def errored_traces(self) -> int: + """Number of traces whose executor run unwound with an exception.""" + return self._errored_traces + + @property + def admitted_traces(self) -> int: + """Number of traces admitted (run started) this phase.""" + return self._admitted_traces + + def _resolved_num_sessions(self) -> int | None: + """Effective session TARGET for this phase (explicit, else derived corpus size). + + An explicit ``--num-conversations`` (``expected_num_sessions`` on the + CreditPhaseConfig) always wins. Otherwise, for a BARE graph PROFILING run + -- no explicit stop condition at all (``expected_num_sessions`` / + ``total_expected_requests`` / ``expected_duration_sec`` and the lifecycle + duration all unset; the bare-graph config case validated against the + graph dataset in ``check_phase_dataset_compatibility``) -- the + single-corpus-pass target is derived from the loaded corpus: + ``N = len(self._parsed.traces)`` -- mirroring dag_jsonl's roots->sessions + convention and giving progress reporting a concrete ``N``. + + SCOPE (deliberate): this value is a TARGET used ONLY for + :meth:`_resolve_lane_count` lane clamping and as the reported + expected-session count. It is NOT a recycle-enabling stop condition -- + :meth:`_recycle_has_stop_condition` and :meth:`_can_recycle` read the + EXPLICIT ``expected_num_sessions`` directly, so a bare run stays in + SINGLE-PASS mode (each trace once, no recycle, clean pass-0 plans). + Routing this derived ``N`` into the recycle gate would flip a bare run + into the bounded/recycle path and silently change its dispatch (fresh-start + plan salting, divergent t* instants) -- which is why it is excluded there. + + Returns ``None`` (no derived target) when another explicit stop already + bounds the run without a session count (``--request-count`` / + ``--benchmark-duration``), and for WARMUP phases (their priming / + cache-pressure fan-out keeps its historical behavior). + """ + explicit = getattr(self._config, "expected_num_sessions", None) + if explicit is not None and explicit > 0: + return int(explicit) + if self._is_warmup_phase: + return None + if ( + getattr(self._config, "total_expected_requests", None) + or getattr(self._config, "expected_duration_sec", None) + or ( + self._lifecycle is not None + and self._lifecycle.time_left_in_seconds() is not None + ) + ): + return None + traces = getattr(self._parsed, "traces", None) + return len(traces) if traces else None + + def _resolve_lane_count(self, total: int, recycle_is_bounded: bool) -> int: + """Resolve how many concurrent lanes this phase fans out. + + AgentX builds exactly ``concurrency`` lanes wrapping the corpus + (``trajectory_source.py:212`` ``_target_size = concurrency``), so lane + fan-out can EXCEED the corpus size and is sustained by recycle. We mirror + that: the lane count is the phase ``concurrency`` (``_max_concurrent``), + independent of corpus size, when a stop condition exists for the lanes to + recycle toward. + + Two clamps keep semantics sane: + * ``--num-conversations N`` (explicit ``expected_num_sessions``) caps the + TOTAL distinct roots ever started, so starting more than ``N`` lanes is + pointless -- clamp lanes to ``N``. (The per-lane recycle gate + ``_can_recycle`` enforces the same ``N`` cap across recycles.) A BARE + graph run derives the SAME ``N = len(traces)`` target + (:meth:`_resolved_num_sessions`), so lanes clamp to the corpus size -- + but the bare run stays in SINGLE-PASS mode (``recycle_is_bounded`` is + False; see :meth:`_recycle_has_stop_condition`), so each lane does one + pass with clean pass-0 plans and no recycle. + * The ``recycle_is_bounded=False`` clamp to the corpus size is the bare + run's actual single-pass bound; the derived session target produces the + same lane clamp, so both agree. + """ + lanes = self._max_concurrent + sessions = self._resolved_num_sessions() + if sessions is not None and sessions > 0: + lanes = min(lanes, int(sessions)) + if not recycle_is_bounded: + lanes = min(lanes, total) + return max(1, lanes) + + def set_lane_limit(self, limit: int) -> None: + """Update the live lane-admission limit (``LaneSettableProtocol``). + + Driven by the concurrency ramper (``--concurrency-ramp-duration``): + lanes park in :meth:`_wait_for_lane_admission` before their first + instance and are released as the limit rises. Clamped to + ``[1, _max_concurrent]``; lowering the limit mid-run does not stop + already-admitted lanes (ramp semantics are raise-only, matching the + session-slot ramp). + """ + new_limit = max(1, min(int(limit), self._max_concurrent)) + if new_limit == self._lane_limit: + return + self._lane_limit = new_limit + # Swap-then-set: waiters parked on the OLD event wake and re-check; + # later waiters park on the fresh event. + released, self._lane_limit_raised = ( + self._lane_limit_raised, + asyncio.Event(), + ) + released.set() + + async def _wait_for_lane_admission(self, lane: int) -> None: + """Park lane ``lane`` until the live lane limit admits it (index < limit).""" + while lane >= self._lane_limit: + await self._lane_limit_raised.wait() + + def _resolve_concurrency(self, override: int | None) -> int: + """Resolve the trace-admission bound. + + Precedence: explicit ``override`` > phase ``concurrency`` (when set and + positive) > ``1`` (the plain aiperf default). + """ + if override is not None and override > 0: + return override + cfg_conc = getattr(self._config, "concurrency", None) + if cfg_conc is not None and cfg_conc > 0: + return int(cfg_conc) + return 1 + + def _guard_explicit_oversubscription(self, distinct: int) -> None: + """Fail loud when concurrency EXPLICITLY over-subscribes a non-wrapping corpus. + + Turns the old silent clone-to-fill into a hard configuration error + (#1106-adjacent). ``distinct`` is ``len(self._parsed.traces)`` -- the + per-tree distinct-loaded-traces count. We compare against the RESOLVED + phase ``concurrency`` (``self._max_concurrent``, the raw over-subscription + signal), NOT the clamped ``_resolve_lane_count`` (which already folds the + corpus size in and so could never trip). When concurrency exceeds the + corpus AND wrapping is not allowed, there are too few distinct traces to + fill the lanes without cloning the operator never asked for -- so raise. + EXCEPTION: an explicit session budget (``--num-conversations`` -> + ``expected_num_sessions``) that fits within the distinct corpus bounds + total instances below any cloning need, so concurrency stays a mere + ceiling and the guard stands down. + + An empty corpus (``distinct < 1``) is already fail-loud upstream in the + adapters; we do not double-handle it. The default concurrency ``1`` never + exceeds a non-empty corpus, so a plain run never raises. + + ``wrap_allowed`` mirrors ``GraphDispatchResolver``'s default so a + direct-construction + path (``_allow_dataset_wrap is None``) never spuriously raises: wrapping is + allowed when explicitly set True, or -- unset -- when cache-bust is on. + """ + if distinct < 1: + return + from aiperf.common.enums import CacheBustTarget + + wrap_allowed = ( + self._allow_dataset_wrap + if self._allow_dataset_wrap is not None + else (self._cache_bust != CacheBustTarget.NONE) + ) + if self._max_concurrent <= distinct or wrap_allowed: + return + + # An explicit session budget (--num-conversations) bounds TOTAL + # instances: when it fits within the distinct corpus, no lane can + # ever need a clone -- concurrency is only a ceiling on simultaneous + # lanes, not a demand for that many traces -- so over-provisioned + # concurrency is not over-subscription. + sessions = getattr(self._config, "expected_num_sessions", None) + if sessions is not None and sessions <= distinct: + return + + # CAPPED phrasing when the corpus was bounded by a selection knob: + # ``num_dataset_entries``/``max_context_length`` are threaded onto the + # phase ``CreditPhaseConfig`` in ``TimingConfig.from_run`` (mirroring + # allow_dataset_wrap/cache_bust), so a live run distinguishes a capped + # corpus from one that simply has fewer distinct traces. + capped_by = [ + flag + for attr, flag in ( + ("num_dataset_entries", "--num-dataset-entries"), + ("max_context_length", "--max-context-length"), + ) + if getattr(self._config, attr, None) + ] + if capped_by: + shortfall = f"distinct loaded traces (capped to {distinct} by {'/'.join(capped_by)})" + else: + shortfall = ( + f"distinct loaded traces; only {distinct} distinct traces available" + ) + + # Reaching this raise always means wrapping is disallowed, so the note is + # unconditional and neutral -- it must NOT attribute the disable to the + # user: ``GraphDispatchResolver`` collapses an UNSET --allow-dataset-wrap to False + # when cache-bust is off (the default), so ``self._allow_dataset_wrap is + # False`` cannot distinguish user-explicit-false from resolver-derived. + parts = [ + f"concurrency {self._max_concurrent} exceeds {distinct} {shortfall}.", + "Dataset wrapping is disabled, so the corpus cannot fill the requested lanes.", + f"Reduce --concurrency to <= {distinct}, or pass --allow-dataset-wrap " + "to intentionally clone/recycle the corpus.", + ] + if self._cache_bust is None or self._cache_bust == CacheBustTarget.NONE: + parts.append( + "If you enable wrapping, also enable cache-bust (e.g. --cache-bust " + "first_turn_prefix) so clones do not collide on identical prefixes." + ) + from aiperf.common.exceptions import ConfigurationError + + raise ConfigurationError(" ".join(parts)) + + def _plan_for_lane(self, trace: Any, lane_index: int) -> GraphTrace | None: + """Resolve the t* plan for ``trace`` on ``lane_index`` (AgentX lane salt). + + Lane ``0`` reuses the prebuilt ``_plans`` entry (byte-identical to the + single-pass path). Higher lanes / recycle passes draw a DISTINCT + lane-salted t* (``sha256(seed:trace_id:lane)``) so the same template + recurring across lanes resumes at a different snapshot instant, exactly + as AgentX's ``_build_trajectory_for_lane`` seeds per absolute lane. The + per-lane source is cached by ``(trace_id, lane_index)`` so repeated + recycle passes onto the same lane re-plan only once. With the default + ``[0, 0]`` window every lane collapses to ``t*=0`` (identity), so lane + fan-out adds no t* divergence on the working profiling path. + """ + key = (trace.id, lane_index) + cached = self._lane_plans.get(key) + if cached is not None: + return cached + if lane_index == 0: + plan = self._plans.get(trace.id) + else: + # Plan ONLY this trace on the lane-salted source. ``iter_traces`` would + # plan EVERY corpus trace just to find one (O(lanes * traces) snapshot + # elaborations across a fan-out), so call ``_plan_trace`` directly. The + # per-lane source is cached by ``lane_index`` (catalog / namespace map + # are lane-independent) so a fan-out reuses one source per lane. + plan = self._lane_source(lane_index)._plan_trace(trace) + if plan is not None: + self._lane_plans[key] = plan + return plan + + def _lane_source(self, lane_index: int) -> GraphIRConversationSource: + """Return the (cached) lane-salted t* source for ``lane_index``. + + The source's catalog + namespace map are lane-independent, so one source + per lane is reused across every trace planned on that lane (avoids the + O(lanes * traces) catalog rebuild a per-call construction would incur). + """ + source = self._lane_sources.get(lane_index) + if source is None: + source = GraphIRConversationSource( + parsed=self._parsed, + start_min_ratio=self._start_min_ratio, + start_max_ratio=self._start_max_ratio, + random_seed=self._t_star_random_seed, + lane_index=lane_index, + ) + self._lane_sources[lane_index] = source + return source + + def _draw_index(self, x: int, total: int) -> int: + """Remap a monotonic draw counter ``x`` to a trace index in ``[0, total)``. + + This is the single choke point every cross-trace draw in the lane + fan-out / recycle loop routes through (:meth:`_resolve_pass0_lanes`, the + index-safe pass-0 fallback, the profiling recycle draw, the pressure + fresh-start + recycle draws), so ``--dataset-sampling-strategy`` governs + WHICH template a freed lane serves without changing the draw counters. + + * ``sequential`` (or ``None``): return ``x % total`` -- byte-for-byte the + historical cursor-with-wrap draw. Sequential must be unchanged. + * ``shuffle``: map ``x`` to ``perm[pass][x % total]`` where + ``pass = x // total``, drawing each pass's permutation from a + pass-salted seed (:func:`_seed_for_draw_pass`). Each pass of ``total`` + draws covers every index exactly ONCE (without replacement), then a + fresh seeded permutation begins -- the same music-shuffle contract the + conversation-plane ``ShuffleSampler`` provides. + * ``random``: coerced to ``shuffle`` (without-replacement) semantics. + Each lane recycle here is a single corpus pass, so with-replacement + ``random`` would duplicate/omit templates within a pass; coercing to + shuffle keeps coverage exact. random == shuffle in this context. + """ + if total <= 0: + return 0 + if not self._draw_is_shuffled(): + return x % total + pass_index, offset = divmod(x, total) + return self._draw_permutation(pass_index, total)[offset] + + def _draw_is_shuffled(self) -> bool: + """True iff the resolved sampling strategy permutes (shuffle / random). + + ``None`` and ``sequential`` take the byte-identical ``x % total`` draw; + ``shuffle`` and ``random`` (coerced to without-replacement) permute. + """ + strategy = self._dataset_sampling_strategy + if strategy is None: + return False + from aiperf.plugin.enums import DatasetSamplingStrategy + + return strategy in ( + DatasetSamplingStrategy.SHUFFLE, + DatasetSamplingStrategy.RANDOM, + ) + + def _draw_permutation(self, pass_index: int, total: int) -> list[int]: + """Return the cached seeded permutation of ``range(total)`` for a pass. + + Built once per ``(total, pass_index)`` from a pass-salted numpy RNG + (:func:`_seed_for_draw_pass` -> ``np.random.default_rng`` -> in-place + Fisher-Yates ``shuffle``, matching ``ShuffleSampler``'s numpy shuffle), + then reused for every draw in that pass so draws are cheap + consistent. + """ + key = (total, pass_index) + cached = self._draw_perm_cache.get(key) + if cached is not None: + return cached + import numpy as np + + rng = np.random.default_rng( + _seed_for_draw_pass(self._t_star_random_seed, pass_index) + ) + perm = list(range(total)) + rng.shuffle(perm) + self._draw_perm_cache[key] = perm + return perm + + async def setup_phase(self) -> None: + """Guard against over-subscription, then install the graph observers. + + The wrap-guard runs FIRST and on THIS awaited path (not ``execute_phase``, + which ``PhaseRunner`` launches fire-and-forget so a raise there is + swallowed): ``PhaseRunner`` awaits ``setup_phase`` directly, so a + ``ConfigurationError`` here propagates up through ``_run_strategy`` to the + run's failure handler and fails the run loudly (the #1106-adjacent + contract). It must therefore precede any observer install / dispatch. + + Both observers de-multiplex to the owning per-trace adapter by + ``trace_id``: the return observer resolves the parked dispatch Future, + the first-token observer releases a successor gated on a node's observed + first token (post-TTFT anchoring). The first-token observer is installed + only when a registrar was wired (``None`` in the unit harness). + """ + self._guard_explicit_oversubscription(len(self._parsed.traces)) + self._register_observer(self._on_graph_return) + if self._register_first_token_observer is not None: + self._register_first_token_observer(self._on_graph_first_token) + + @staticmethod + def _wall_us() -> float: + """Monotonic wall instant in microseconds (ledger + handoff clock).""" + return time.perf_counter_ns() / 1_000.0 + + def _record_return_wall(self, credit: Credit) -> None: + """Ledger one warmup-phase return: ``{instance_id: {node_id: wall}}``. + + Runs BEFORE adapter de-mux so late returns after a drain-cancel (the + adapter already reaped) still land in the ledger -- they were real + server-side executions and must count as executed in the handoff. + """ + instance_id = getattr(credit, "trace_id", None) + ordinal = getattr(credit, "node_ordinal", None) + if instance_id is None or ordinal is None: + return + template_id = instance_id.split("::", 1)[0] + inverse = self._ordinal_to_node.get(template_id) + if inverse is None: + inverse = { + o: nid for nid, o in self._catalog.catalog.get(template_id, {}).items() + } + self._ordinal_to_node[template_id] = inverse + node_id = inverse.get(ordinal) + if node_id is None: + return + self._return_walls.setdefault(instance_id, {})[node_id] = self._wall_us() + + def _record_warmup_failure(self, credit: Credit, error: str) -> None: + """Account one TERMINAL warmup-phase request failure (agentx parity). + + Called from :meth:`_on_graph_return` for a WARMUP-phase return carrying + a non-None ``error`` that is NOT cancelled. Keeps the running count plus + up to five human-readable samples for the abort message. + + Cancelled returns are deliberately EXCLUDED here, even when they carry + error text. AgentX counts a cancelled warmup credit as a failure because + its warmup drain PAUSES the issuance gate and lets genuinely in-flight + server requests finish -- a cancellation there means the server never + answered. Our extended-warmup drain instead CANCELS the executor + coroutines on the pressure duration timer (:meth:`_run_pressure_stage`), + so a cancellation surfacing at drain is self-inflicted teardown, not a + server failure; counting it would abort every healthy timed warmup. + """ + self._warmup_failure_count += 1 + if len(self._warmup_failure_samples) < 5: + self._warmup_failure_samples.append( + f"{getattr(credit, 'trace_id', '?')}" + f"[node_ordinal={getattr(credit, 'node_ordinal', '?')}]: {error}" + ) + + def report_warmup_failures(self) -> None: + """Raise if this WARMUP phase recorded terminal request failures. + + AgentX parity (``callback_handler.py`` warmup gate + runner report): a + warmup that could not faithfully prime the server's KV cache must abort + the run BEFORE profiling -- benchmark numbers from a degraded pool look + valid and are not. Called by ``PhaseRunner._run_strategy`` after + returning-complete (never on the cancelled early-return path). No-op for + PROFILING phases and clean warmups. + + Reuses :class:`TrajectoryWarmupFailedError`; the recorded failure + samples (``trace_id[node_ordinal]: error``) are passed as its trace + descriptors so the raised message surfaces both the error text and the + count. The exception derives its count from ``len(failed_trace_ids)`` + and the samples are capped at 5, so past the cap a synthetic trailing + entry carries the TRUE total (the message must never under-report a + degraded pool). + """ + if not self._is_warmup_phase or self._warmup_failure_count == 0: + return + samples = list(self._warmup_failure_samples) + overflow = self._warmup_failure_count - len(samples) + if overflow > 0: + samples.append( + f"... and {overflow} more failure(s) " + f"(total {self._warmup_failure_count})" + ) + raise TrajectoryWarmupFailedError(samples) + + def _on_graph_return( + self, credit: Credit, error: str | None, cancelled: bool + ) -> None: + """Route one graph credit return to its owning per-trace adapter. + + The shared ``CreditCallbackHandler`` fires this UNCONDITIONALLY for every + credit carrying a ``trace_id``. We look the adapter up by the credit's + ``trace_id`` (the root trace id the adapter stamped on issue) and let it + resolve / reject the parked dispatch Future. An unknown trace id is a + graceful no-op (e.g. a late return after the trace already unwound). + """ + if self._pressure_enabled and not cancelled: + # Wire-cancelled turns are NOT executed: the grace-expiry drain + # cancels in-flight requests, and the server may never have + # completed them. Keeping them out of the ledger keeps them out of + # executed_node_ids, so profiling refires them -- which also makes + # a SUCCESSFUL cancel-drain (all credits returned-or-cancelled) + # yield a VALID handoff instead of needing the completeness skip. + self._record_return_wall(credit) + if self._is_warmup_phase and error is not None and not cancelled: + self._record_warmup_failure(credit, error) + trace_id = getattr(credit, "trace_id", None) + if trace_id is None: + return + adapter = self._adapters.get(trace_id) + if adapter is None: + node_ordinal = getattr(credit, "node_ordinal", None) + # Routine at the extended-warmup drain (executors cancelled, wire + # credits still landing -- already ledgered above); exceptional + # everywhere else. + log = _logger.debug if self._pressure_active else _logger.warning + log( + lambda: ( + f"graph return for unknown instance_id={trace_id!r} " + f"node_ordinal={node_ordinal!r} dropped: no live adapter is " + "registered for this instance (a late return arrived after the " + "instance was cancelled/timed-out/errored and its adapter was " + "already reaped). The parked dispatch will not resolve from this " + "return." + ) + ) + return + adapter.resolve(credit, error, cancelled) + + def _on_graph_first_token(self, first_token: FirstToken) -> None: + """Route one graph credit's first-token event to its owning adapter. + + Mirrors :meth:`_on_graph_return`: the shared ``CreditCallbackHandler`` + fires this for every ``FirstToken`` carrying a ``trace_id`` (a graph + credit). We look up the owning adapter by that instance id and hand it + the emitting credit's ``(x_correlation_id, turn_index)`` so it can fire + the successor's parked first-token callback (post-TTFT anchoring). An + unknown / ``None`` trace id -- a late token after the trace unwound or + its instance was reaped, or a non-graph fast-path token -- is a + graceful no-op (the anchor simply falls back to its dispatch delay). + """ + trace_id = getattr(first_token, "trace_id", None) + if trace_id is None: + return + adapter = self._adapters.get(trace_id) + if adapter is None: + _logger.debug( + lambda: ( + f"graph first-token for unknown instance_id={trace_id!r} " + "dropped: no live adapter (successor falls back to dispatch delay)" + ) + ) + return + adapter.on_first_token( + getattr(first_token, "x_correlation_id", None), + getattr(first_token, "turn_index", None), + ) + + async def execute_phase(self) -> None: + """Dispatch every admitted trace, owning completion AgentX-faithfully. + + Stop semantics mirror AgentX's stop-condition model + (``aiperf.timing.phase.stop_conditions``), NOT a per-trace drain wait: + ``--num-conversations N`` (``SessionCountStopCondition``) caps admitted + traces (each weka trace is one conversation/DAG); ``--request-count N`` + (``RequestCountStopCondition``) caps node dispatches via + ``issue_graph_credit``'s gate (a refused issue is a clean per-trace stop); + ``--benchmark-duration D`` (``DurationStopCondition``) caps WALL time via + :meth:`_run_traces_under_duration_budget`. A weka trace faithfully + replays its recorded idle gaps, so a single trace can span its whole + recorded duration -- the duration budget is what bounds that, exactly as + AgentX's ``runner.py::_wait_for_sending_complete`` cancels in-flight + scheduled turns when the duration elapses. + + Completion: a graph credit carries a ``trace_id`` so the linear + ``CreditCounter`` never trips ``is_final_credit``; THIS strategy owns the + signal. The ``finally`` ALWAYS freezes the sent count and sets + ``all_credits_sent_event`` (drain / duration-cancel / error alike), so a + never-returning executor can never wedge the phase with the event unset. + The ``PhaseRunner`` then awaits ``all_credits_returned_event`` (set by the + callback handler once every issued-and-not-cancelled credit returns). + """ + traces = list(self._parsed.traces) + if not traces: + self.debug("no traces to replay; phase complete") + self._credit_issuer.mark_graph_sending_complete() + self._credit_issuer.set_graph_all_returned_event() + return + + self._advise_if_idle_gap_corpus_without_duration() + try: + await self._run_traces_under_duration_budget(traces) + if self._pressure_enabled: + await self._run_pressure_stage(traces) + finally: + # ALWAYS runs (drain / duration-cancel / error): no further graph + # credit will be issued, so freeze the authoritative per-node sent + # count and signal sending-complete. If no credit was ever issued + # (no-op corpus, or every trace cancelled pre-dispatch), also set the + # returned event since no return will fire it. + self._credit_issuer.mark_graph_sending_complete() + if self._credit_issuer.graph_all_returned(): + self._credit_issuer.set_graph_all_returned_event() + + def _has_benchmark_duration(self) -> bool: + """True iff a ``--benchmark-duration`` wall budget bounds this phase. + + The budget lives on the lifecycle (``time_left_in_seconds()`` is non-None + only when ``expected_duration_sec`` is configured) or, equivalently, on the + phase config (``expected_duration_sec``). Either source is authoritative; + we check both so the advisory fires correctly under the unit harness (which + may wire only one). + """ + if ( + self._lifecycle is not None + and self._lifecycle.time_left_in_seconds() is not None + ): + return True + return bool(getattr(self._config, "expected_duration_sec", None)) + + def _max_inter_turn_gap_seconds(self) -> float: + """Largest recorded inter-turn idle gap (seconds) across every graph. + + Scans every node's ``min_start_delay_us`` and every ``StaticEdge``'s + ``delay_after_predecessor_us`` (and the start-anchored + ``delay_after_predecessor_start_us``) over ``ParsedGraph.graph`` and every + entry in ``ParsedGraph.graphs`` (the multi-graph corpus). These are the + per-gap-warped delays the executor replays VERBATIM -- both end-to-start + (``delay_after_predecessor_us``) and dispatch-to-start + (``delay_after_predecessor_start_us``); the max is the worst single + stretch a faithful replay parks on. Returns ``0.0`` for + a gap-free corpus (the AgentX bare-CLI default, end-to-start delays OFF). + """ + max_us = 0.0 + graphs = [self._parsed.graph] + graphs.extend(getattr(self._parsed, "graphs", {}).values()) + for graph in graphs: + if graph is None: + continue + for node in getattr(graph, "nodes", {}).values(): + delay = getattr(node, "min_start_delay_us", None) + if delay is not None: + max_us = max(max_us, float(delay)) + for edge in getattr(graph, "edges", []): + delay = getattr(edge, "delay_after_predecessor_us", None) + if delay is not None: + max_us = max(max_us, float(delay)) + start_delay = getattr(edge, "delay_after_predecessor_start_us", None) + if start_delay is not None: + max_us = max(max_us, float(start_delay)) + # Kind-complete scan: the first-token refinement is <= the start + # delay on the same edge, so this never raises the max in practice. + first_token_delay = getattr( + edge, "delay_after_predecessor_first_token_us", None + ) + if first_token_delay is not None: + max_us = max(max_us, float(first_token_delay)) + return max_us / MICROS_PER_SECOND + + def _advise_if_idle_gap_corpus_without_duration(self) -> None: + """Emit a once-per-phase advisory for an idle-gap corpus with no duration. + + A faithful recorded-trace replay honors every recorded inter-turn idle gap verbatim, + so a count/session/bare run (no ``--benchmark-duration``) spans the slowest + admitted trace's full recorded wall time -- exactly AgentX's count-mode + (``is_final_credit`` fires only once the last turn is SENT, after every + scheduled idle gap elapses; ``runner._wait_for_sending_complete`` then waits + on that with a ``None`` timeout). This is faithful, NOT a hang, but on a + human-pace corpus the wall time is minutes-scale with no console output + during the parked gaps. Advise the operator that ``--benchmark-duration`` + bounds it (cancels the still-parked idle nodes, keeps dispatched records) + and that Ctrl+C exports the partial results gracefully. No-op when a + duration is set, when the corpus is gap-free, or when its largest gap is + below ``IDLE_GAP_NO_DURATION_WARN_SECONDS``. Also a no-op for a + WARMUP phase: warmup bursts only the boundary priming turns (see + :func:`rewrite_for_warmup`) and never replays recorded gaps. + """ + if self._phase_variant() == "warmup": + return + if self._has_benchmark_duration(): + return + max_gap_s = self._max_inter_turn_gap_seconds() + threshold_s = Environment.GRAPH.IDLE_GAP_NO_DURATION_WARN_SECONDS + if max_gap_s < threshold_s: + return + self.notice( + lambda: ( + f"idle-gap corpus has a recorded inter-turn gap of up to " + f"{max_gap_s:.0f}s and no --benchmark-duration is set: this phase " + "replays every gap faithfully (AgentX count-mode parity), so its " + "wall time spans the slowest admitted trace's full recorded " + "duration with no console output during the parked gaps. Pass " + "--benchmark-duration to bound the run (it cancels the " + "still-parked idle nodes and keeps the records dispatched so far), " + "or press Ctrl+C to finalize + export the partial results now." + ) + ) + + async def _run_traces_under_duration_budget(self, traces: list[Any]) -> None: + """Run the lane fan-out + recycle loop under the phase duration budget. + + No ``--benchmark-duration`` (``time_left_in_seconds() is None``) -> await + natural completion; every lane recycles fresh templates until the + stop-condition gate (``--num-conversations`` / ``--request-count``) refuses + a new session, matching AgentX's count/session-mode. Otherwise on timeout + the lane-runner task is cancelled, cancelling the outer ``TaskGroup`` and + every in-flight executor (parked Futures reject, coroutines unwind), so the + phase finalizes on the ``DurationStopCondition`` exactly as AgentX's + ``_wait_for_sending_complete`` duration path (Gap-4). Already-returned + records are kept; only not-yet-fired idle-parked nodes drop. + """ + timeout = ( + self._lifecycle.time_left_in_seconds() + if self._lifecycle is not None + else None + ) + if timeout is not None: + self._duration_deadline = asyncio.get_running_loop().time() + timeout + dispatch = asyncio.ensure_future(self._run_lanes(traces)) + try: + try: + await asyncio.wait_for(dispatch, timeout=timeout) + except TimeoutError: + # ``wait_for`` already requested cancellation of ``dispatch`` + # (AgentX cancel_all_pending parity); await its unwind so no + # orphan survives. + self.notice( + lambda: ( + f"graph phase duration budget ({timeout:.1f}s) elapsed with " + f"{self._admitted_traces - self._completed_traces} instance(s) still " + "replaying recorded idle gaps; cancelling in-flight executors " + "(duration stop condition)" + ) + ) + try: + await dispatch + except asyncio.CancelledError: + # ``wait_for``'s own cancellation of ``dispatch`` unwinds here + # and is absorbed (a healthy duration stop). An EXTERNAL cancel + # of THIS task (runner send-timeout, Ctrl+C) racing that unwind + # must NOT be swallowed: re-raise when the current task itself + # has a pending cancellation. + current = asyncio.current_task() + if current is not None and current.cancelling(): + raise + finally: + self._duration_deadline = None + # Duplication report: fire on BOTH natural completion and the + # duration-cancel path (the MOST recycle-heavy mode -- lanes recycle + # until the timer). This is the SOLE ``_run_lanes`` caller and the + # finally runs exactly once, so the report is emitted once per phase. + self._report_dispatch_duplication(self._instances_started, len(traces)) + + def _past_duration_deadline(self) -> bool: + """True once the active stage's duration budget has elapsed. + + Cooperative twin of the budget wrappers' ``wait_for`` cancel: a cancel + delivered into a TaskGroup whose children complete constantly can be + lost on Python 3.11, leaving the recycle loop running unbounded. Lanes + check this each iteration so the stage halts on time regardless of + whether the cancel lands. + """ + deadline = self._duration_deadline + return deadline is not None and (asyncio.get_running_loop().time() >= deadline) + + async def _run_lanes(self, traces: list[Any]) -> None: + """Drive ``concurrency`` recycling lanes over the wrapped corpus. + + AgentX (``trajectory_source.py:212`` ``_target_size = concurrency``; + the retired agentic-replay plane's recycle-on-root-final-turn): we build + ``_resolve_lane_count`` lanes (``concurrency``, clamped only by an + ``--num-conversations`` cap or a no-stop-condition single pass) -- lane + ``i`` starts on ``traces[i % N]`` -- and each lane LOOPS, drawing the next + wrap template (round-robin over + the corpus) and re-dispatching it onto its freed slot until the + stop-condition gate refuses a new session. So ``concurrency`` is SUSTAINED + even when it exceeds the corpus size (the C2 + C1 fix), instead of + decaying to ``N`` after one corpus pass. + + ``--num-conversations N`` caps the TOTAL number of distinct root sessions + ever started (``SessionCountStopCondition.can_start_new_session`` -- the + same gate AgentX recycle consults); ``--request-count N`` caps total + node/LLM dispatches via ``issue_graph_credit``'s per-dispatch gate (a + refused issue cleanly stops the in-flight instance, and the lane's next + recycle attempt sees the gate closed and stops). With NEITHER cap set the + lanes run a SINGLE corpus pass: each freed lane keeps drawing the next + unclaimed template until every corpus position has been claimed exactly + once (no unbounded recycle without a stop condition), so a bare + ``--graph weka.json`` run covers the whole corpus and still terminates. + + Each lane task runs on ``phase_tg``; each lane swallows its own instance + errors so one failed instance never aborts the phase. + """ + recycle_is_bounded = self._recycle_has_stop_condition() + lanes = self._resolve_lane_count(len(traces), recycle_is_bounded) + # Pass-0 lane assignment mirrors AgentX ``_build_trajectories``: lane ``i`` + # (i == spawnable rank) takes the i-th SPAWNABLE trace in corpus order, + # SKIPPING unspawnable ones, so a trace after an unspawnable one gets the + # SAME shifted lane_index (= spawnable rank) AgentX gives it -- making the + # lane-salted t* byte-identical across both engines. The resolution is + # SEQUENTIAL because spawnability is tested at the lane-salted t* for the + # candidate's TARGET rank (rank advances only on a spawnable hit), exactly + # as AgentX seeds ``lane = len(trajectories)`` per draw. + pass0_traces, corpus_cursor = self._resolve_pass0_lanes(traces, lanes) + lanes = len(pass0_traces) + if self._warmup_handoff is not None: + # AgentX dispatches EVERY drained lane at the handoff; the session + # cap gates only recycles (resumed streams are continuations, not + # new sessions). Lanes beyond len(pass0_traces) get an index-safe + # fallback template below -- their real template comes from the + # handoff entry or the fresh-start cursor draw anyway. Applied AFTER + # the re-shrink so an unspawnable-shrunk resolution can't undo it. + lanes = max(lanes, self._warmup_handoff.pressure_lane_count) + # Recycle draws continue from the corpus position AFTER the last one the + # pass-0 spawnable resolution consumed (AgentX's recycle reuses the SAME + # sampler, so it resumes where ``_build_trajectories`` left off -- past the + # skipped-unspawnable traces too), turn-0 full replay onto the freed slot. + # AgentX shared-sampler parity: bounded profiling recycles continue from + # the pressure stage's last draw so freed lanes don't re-serve templates + # pressure just replayed. Single-pass mode (no stop conditions) keeps its + # own cursor -- its termination check (next_index >= len(traces)) encodes + # cover-the-corpus-once, which a carried cursor would break. + if self._warmup_handoff is not None and recycle_is_bounded: + next_index = self._warmup_handoff.corpus_cursor + else: + next_index = corpus_cursor + traces_by_id = {t.id: t for t in traces} + + # Duplication report: every lane start AND every serial recycle is + # one more instance dispatched off the finite trace pool. Counting each + # ``_run_instance`` call here (recycles included) gives the total the + # duplication factor divides by ``distinct_loaded_traces``. Tracked on an + # instance attribute (reset here, one ``_run_lanes`` call per phase) so + # ``_run_traces_under_duration_budget`` can report it even when this lane + # future is CANCELLED by the duration budget. + self._instances_started = 0 + + async with asyncio.TaskGroup() as phase_tg: + for lane in range(lanes): + + async def _lane(lane: int = lane) -> None: + nonlocal next_index + # Lane-admission gate: under a concurrency ramp + # (--concurrency-ramp-duration) lanes above the live limit + # park here and are admitted as the ramper raises it -- + # 1 -> concurrency over the ramp duration. No ramp => the + # limit already equals _max_concurrent and this returns + # immediately. Applies to handoff-resume lanes too (the + # ramp's purpose is spreading load onto a cold server; + # duration-cancel cancels parked waiters cleanly). + await self._wait_for_lane_admission(lane) + # Index-safe pass-0 template: a handoff can bump ``lanes`` + # past ``len(pass0_traces)`` (every drained pressure lane is + # honored), so lanes beyond the spawnable-resolved set take a + # wrap-around fallback here and get their real template from + # the handoff entry or the fresh-start cursor draw below. + trace: Any = ( + pass0_traces[lane] + if lane < len(pass0_traces) + else traces[self._draw_index(lane, len(traces))] + ) + recycle_pass = 0 + entry = self._handoff_for_lane(lane) + fresh_start = False + if entry is not None: + resumed = traces_by_id.get(entry.template_trace_id) + if resumed is not None: + trace = resumed + else: + self.warning( + f"handoff template {entry.template_trace_id!r} " + f"not in corpus; lane {lane} resumes its " + "normal pass-0 assignment" + ) + elif ( + self._warmup_handoff is not None + and recycle_is_bounded + and lane < self._warmup_handoff.pressure_lane_count + ): + # This pressure lane completed at drain: fresh-start it + # from the shared cursor (agentx empty-lane parity) + # instead of re-running a t* resume the pressure stage + # already executed against the server (which would be + # measured warm). Single-pass mode is excluded: a fresh + # draw would consume a corpus position and hole the + # cover-the-corpus-once contract. + template_index = next_index + next_index += 1 + trace = traces[self._draw_index(template_index, len(traces))] + fresh_start = True + while True: + self._instances_started += 1 + await self._run_instance( + trace, lane, recycle_pass, fresh_start=fresh_start + ) + fresh_start = False + if self._past_duration_deadline(): + return + if recycle_is_bounded: + # Recycle while the stop-condition gate still admits + # a new session. + if not self._can_recycle(): + return + elif next_index >= len(traces): + # No stop condition: single corpus pass -- stop once + # every corpus position has been claimed once. + return + template_index = next_index + next_index += 1 + trace = traces[self._draw_index(template_index, len(traces))] + recycle_pass += 1 + + phase_tg.create_task(_lane(), name=f"graph-lane:{lane}") + + if not recycle_is_bounded: + self.info( + f"no stop condition set: single corpus pass complete, covering " + f"{self._admitted_traces} of {len(traces)} trace(s)" + ) + # NOTE: the duplication report is emitted by the SOLE caller + # (``_run_traces_under_duration_budget``) in a ``finally`` so it fires on + # both natural completion and duration-cancel; do NOT report here. + + def _report_dispatch_duplication( + self, total_instances_started: int, distinct_loaded_traces: int + ) -> None: + """Report the phase's dispatch duplication factor; warn only when unsafe. + + ``factor = total_instances_started / distinct_loaded_traces`` counts every + lane start AND every serial recycle as an instance dispatched off the + finite trace pool, so ``factor > 1`` means the same recorded traces were + replayed more than once (lanes recycled to sustain concurrency / satisfy + ``--request-count`` / ``--num-conversations``). This is a REPORT, never a + failure. A WARNING fires ONLY when the duplication has no cache-bust + antidote (``self._cache_bust`` is ``None`` / ``NONE``): identical prefixes + across clones collide in the server KV cache and inflate cache-hit + metrics. With cache-bust ON every instance mints a distinct marker, so the + duplication is safe and the report stays quiet. Warmup phases are skipped + (their boundary priming is meant to warm the cache; duplication there is + expected), mirroring the idle-gap advisory's warmup carve-out. + """ + if distinct_loaded_traces <= 0 or self._phase_variant() == "warmup": + return + factor = total_instances_started / distinct_loaded_traces + if factor <= 1.0: + return + from aiperf.common.enums import CacheBustTarget + + if self._cache_bust is not None and self._cache_bust != CacheBustTarget.NONE: + return + self.warning( + lambda: ( + f"dispatch duplication factor {factor:.2f}x " + f"({total_instances_started} instances started over " + f"{distinct_loaded_traces} distinct loaded trace(s)): lanes " + "recycled the finite trace pool and cache-bust is OFF, so cloned " + "instances replay identical first-turn prefixes and collide in the " + "inference server's KV cache -- inflating prefix-cache-hit metrics. " + "Set --cache-bust first_turn_prefix to give each instance a " + "distinct marker, or reduce --concurrency / raise " + "--num-dataset-entries so the corpus covers the load without reuse." + ) + ) + + async def _run_pressure_stage(self, traces: list[Any]) -> None: + """Run the extended (cache-pressure) warmup stage for the configured duration. + + AgentX v1.0 ``_start_accelerated_warmup`` parity, dataflow-native: the + post-t* remainder of every lane's template replays compressed (the + WARMUP executors are built with ``compress_edge_delays=True`` via + ``accelerated_warmup``; the worker's 1-token warmup cap keys off the + unchanged ``"warmup"`` phase variant), and freed lanes recycle fresh + templates at t*=0. The stage ends on the duration timer: in-flight + executors are cancelled exactly like the ``--benchmark-duration`` stop + (issued credits still return and land in the ledger; un-fired nodes + drop), which is the drain the teardown handoff is built from. + + The pressure warmup phase is mode-owned: user warmup phases are + superseded at config build (timing/config.py), so the stage always gets + its full duration; the drain is bounded by min(duration, + ``PRESSURE_DRAIN_GRACE_CAP``) with the stash completeness gate as + the safety net. + """ + assert self._cache_pressure_duration_s is not None + duration = self._cache_pressure_duration_s + self._pressure_active = True + self.notice( + f"WARMUP cache pressure: replaying post-t* graphs with zero idle " + f"delay and max_tokens=" + f"{Environment.GRAPH.WARMUP_MAX_OUTPUT_TOKENS} for {duration:.1f}s" + ) + self._duration_deadline = asyncio.get_running_loop().time() + duration + dispatch = asyncio.ensure_future(self._run_pressure_lanes(traces)) + try: + await asyncio.wait_for(dispatch, timeout=duration) + except TimeoutError: + self.notice( + "WARMUP cache pressure duration reached; draining in-flight " + "requests for the profiling handoff" + ) + try: + await dispatch + except asyncio.CancelledError: + # Absorb wait_for's own cancellation (healthy stage end) but + # re-raise when THIS task is being cancelled externally + # (runner send-timeout, Ctrl+C) -- same guard as + # _run_traces_under_duration_budget. + current = asyncio.current_task() + if current is not None and current.cancelling(): + raise + finally: + self._duration_deadline = None + + async def _run_pressure_lanes(self, traces: list[Any]) -> None: + """Drive the pressure lanes: same pass-0 assignment as priming, then recycle. + + Reuses ``_resolve_pass0_lanes`` (with the SAME bounded-recycle shape as + the Stage-A priming fan-out) so lane i continues the SAME (template, + lane-salted t*) chain its boundary priming just warmed -- a lane count + wider than priming would run unprimed pass-0 chops and produce handoff + entries with no priming walls to merge. Lanes recycle until the stage + ends (fresh template at t*=0 per freed slot), matching agentx's + recycle-forever pressure loop (``_handle_accelerated_warmup_return`` + -> ``_spawn_from_recycle_or_id``), with two lane-local gates so a + closed issuer stop gate or a deterministically failing server cannot + hot-spin adapter builds for the whole duration: a clean issuer refusal + ends the lane, and consecutive instance errors back off exponentially + (the sleep also yields the loop under the ``wait_for`` budget). + """ + lane_count = self._resolve_lane_count( + len(traces), self._recycle_has_stop_condition() + ) + pass0_traces, corpus_cursor = self._resolve_pass0_lanes(traces, lane_count) + self._pressure_next_index = corpus_cursor + self._pressure_lane_count = len(pass0_traces) + + async with asyncio.TaskGroup() as pressure_tg: + for lane in range(len(pass0_traces)): + + async def _lane( + lane: int = lane, start_trace: Any = pass0_traces[lane] + ) -> None: + trace = start_trace + pressure_pass = 0 + consecutive_errors = 0 + while True: + plan = ( + self._plan_for_lane(trace, lane) + if pressure_pass == 0 + else None + ) + # {template}::{nonce}: instance identity is template + + # fresh nonce (lane/pass are logged, not encoded). + instance_id = f"{trace.id}::{uuid.uuid4().hex}" + self.debug( + lambda t=trace.id, + i=instance_id, + ln=lane, + p=pressure_pass: ( + f"pressure instance {i} (template={t} lane={ln} pass=p{p})" + ) + ) + # Set-before-await / pop-after: entry present == the + # instance is mid-flight, so a duration-cancel landing + # inside _run_instance leaves it for the handoff. + # Benign race: a cancel delivered exactly at the await + # RESUMPTION of a completed instance stashes it as + # live with all nodes executed -- the profiling chop + # is then empty and the lane recycles (harmless). + self._pressure_live[lane] = ( + trace.id, + instance_id, + plan.t_star_us if plan is not None else 0.0, + pressure_pass, + ) + errors_before = self._errored_traces + refused = await self._run_instance( + trace, + lane, + pressure_pass, + pressure=True, + instance_id=instance_id, + ) + self._pressure_live.pop(lane, None) + if refused: + # Issuer stop gate closed (run cancelled): a fresh + # instance would refuse instantly, forever. + return + if self._past_duration_deadline(): + return + if self._errored_traces > errors_before: + consecutive_errors += 1 + await asyncio.sleep(min(0.25 * 2**consecutive_errors, 5.0)) + else: + consecutive_errors = 0 + template_index = self._pressure_next_index + self._pressure_next_index += 1 + trace = traces[self._draw_index(template_index, len(traces))] + pressure_pass += 1 + + pressure_tg.create_task(_lane(), name=f"graph-pressure-lane:{lane}") + + def _resolve_pass0_lanes( + self, traces: list[Any], lanes: int + ) -> tuple[list[Any], int]: + """Resolve the pass-0 lane->trace map (AgentX spawnable-rank assignment). + + Walks ``traces`` in corpus order assigning lane ``i`` (== spawnable rank) + to the next SPAWNABLE trace, computing each candidate's spawnability at + the lane-salted t* for its TARGET rank (so rank advances only on a hit -- + the sequential dependence AgentX's ``lane = len(trajectories)`` loop has). + Unspawnable candidates are skipped (no lane), exactly mirroring + ``TrajectorySource._build_trajectories`` skipping a ``None`` from + ``_build_trajectory_for_lane`` (``_snapshot_for`` is None). Returns the + per-lane trace list (length ``<= lanes``; shorter only when the corpus has + too few spawnable traces) AND the corpus cursor one past the last consumed + position, so recycle resumes the wrap there. + + The per-candidate corpus draw routes through :meth:`_draw_index`, so + ``--dataset-sampling-strategy`` selects WHICH template each rank serves; + under the default ``sequential`` draw it is ``traces[cursor % n]`` + unchanged. With the default ``[0, 0]`` window (t*=0) every trace is + spawnable (full replay always dispatches), so ``sequential`` collapses to + the prior raw-position assignment (lane ``i`` -> ``traces[i]``) + byte-for-byte. + """ + pass0: list[Any] = [] + cursor = 0 + n = len(traces) + # Bound the skip walk to one full corpus pass past the wrap of ``lanes`` + # so an all-unspawnable corpus can't spin (AgentX caps at + # ``_target_size + 2 * pool_size``); we stop at ``lanes`` hits regardless. + max_cursor = lanes + n + while len(pass0) < lanes and cursor < max_cursor: + trace = traces[self._draw_index(cursor, n)] + cursor += 1 + rank = len(pass0) + if self._is_spawnable(trace, rank): + pass0.append(trace) + return pass0, cursor + + def _is_spawnable(self, trace: Any, lane_index: int) -> bool: + """Whether ``trace`` yields a live dispatch at its lane-salted t* (rank). + + Mirrors AgentX ``TrajectorySource._build_timestamped_trajectory`` -> + ``_snapshot_for`` returning ``None``: a trace is UNSPAWNABLE when, at the + t* sampled for ``(trace, lane_index)``, there is no live, dispatchable + stream -- AgentX's ``not states`` / ``not any(not waiting_on_children)``. + The native equivalent over :func:`compute_snapshot` is "the snapshot has + at least one PROFILED LLM firing": a profiled LLM firing is a turn at/after + t* that survived the spawn-completion gate (AgentX's live, non-future + stream), and an LLM turn is the only thing AgentX would dispatch (markers + carry no request) -- so a snapshot with zero profiled LLM firings is one + AgentX builds no sendable session for and skips. + + At t*=0 (the default ``[0, 0]`` window, or a zero-duration trace) every + trace's full timeline is profiled, so this is always True -- the prior + no-skip behavior, byte-for-byte. + + The lane-salted t* is computed directly (``_lane_salted_t_star``) rather + than through ``_plan_for_lane`` so the spawnability scan never builds the + per-lane catalog / node partition (only the t* + one snapshot are needed); + the value is byte-identical to the plan's ``t_star_us`` because both call + the same ``_sample_t_star`` math, so spawnability and the eventual + dispatch plan agree. + """ + t_star_us = self._lane_salted_t_star(trace, lane_index) + if t_star_us <= 0: + return True + from aiperf.dataset.graph.models import NodeKind + from aiperf.graph.analysis import compute_snapshot + + # ``compute_snapshot`` resolves the trace's own graph internally + # (``elaborate_trace`` -> ``resolve_trace_graph``), so passing the raw + # ``self._parsed`` is correct here. + snapshot = compute_snapshot(self._parsed, trace, t_star_us=int(t_star_us)) + return any(sf.firing.kind == NodeKind.LLM for sf in snapshot.profiled) + + def _lane_salted_t_star(self, trace: Any, lane_index: int) -> float: + """Compute ``trace``'s lane-salted t* (us) WITHOUT building a t* source. + + Reuses the prebuilt lane-0 plan when ``lane_index == 0`` (the common + single-pass case). For higher lanes it inlines ``_sample_t_star``'s seed + + duration math (``_seed_for_trace_lane`` -> ``rng.uniform(lo, hi)`` over the + ratio window) so the spawnability scan never constructs a + ``GraphIRConversationSource`` (whose ``__init__`` rebuilds the whole-corpus + node catalog -- an O(lanes * traces) cost across a fan-out the scan does + not need). The value is byte-identical to ``_plan_for_lane``'s + ``t_star_us`` because it is the SAME computation on the SAME seed. + """ + if lane_index == 0: + plan = self._plans.get(trace.id) + return plan.t_star_us if plan is not None else 0.0 + import numpy as np + + from aiperf.graph.analysis import trace_duration_us + from aiperf.timing.graph_ir_source import _seed_for_trace_lane + + # ``trace_duration_us`` resolves the trace's OWN graph internally + # (``elaborate_trace`` -> ``resolve_trace_graph``), so passing the full + # multi-graph parse matches ``GraphIRConversationSource._sample_t_star`` + # byte-for-byte (it too uses the source's raw ``self._parsed``). + duration_us = trace_duration_us(self._parsed, trace) + if duration_us <= 0: + return 0.0 + lo = self._start_min_ratio * duration_us + hi = self._start_max_ratio * duration_us + if hi <= lo: + return float(lo) + rng = np.random.default_rng( + _seed_for_trace_lane(self._t_star_random_seed, trace.id, lane_index) + ) + return float(rng.uniform(lo, hi)) + + def _recycle_has_stop_condition(self) -> bool: + """True iff an EXPLICIT stop condition exists that recycle can saturate toward. + + Recycle runs ONLY when the user set ``--num-conversations`` + (``expected_num_sessions``) / ``--request-count`` + (``total_expected_requests``) / ``--benchmark-duration``. With none, the + corpus is replayed once and the lanes stop (a bare ``--graph weka.json`` + must NOT recycle forever). + + The DERIVED single-pass session target + (:meth:`_resolved_num_sessions` = ``len(traces)`` for a bare run) is + DELIBERATELY EXCLUDED here: it is a lane-clamp + reported-target only, not + a recycle-enabling stop condition. Routing it in would flip a bare run out + of SINGLE-PASS mode into the bounded/recycle path -- firing the + pressure-lane fresh-start gate and lane-salted recycle plans (``.f``/``.1`` + ids, divergent t* instants) instead of the clean pass-0 plans a + cover-the-corpus-once pass must keep. So a bare run stays single-pass: + each lane does one pass, no recycle, no fresh-start. + """ + cfg = self._config + return bool( + getattr(cfg, "expected_num_sessions", None) + or getattr(cfg, "total_expected_requests", None) + or getattr(cfg, "expected_duration_sec", None) + or ( + self._lifecycle is not None + and self._lifecycle.time_left_in_seconds() is not None + ) + ) + + def _can_recycle(self) -> bool: + """Whether a freed lane may start a fresh root (AgentX recycle gate). + + Mirrors AgentX ``_dispatch_recycled_on_lane`` consulting + ``stop_checker.can_start_new_session()``, but the v1 ``CreditCounter`` + does NOT bump ``sent_sessions`` for graph credits (they bypass the linear + session arithmetic -- ``credit_counter.increment_sent``), so the + ``--num-conversations`` cap must be counted by the STRATEGY: the number of + distinct root sessions ever started is ``_admitted_traces`` (one per + instance run), and the cap is ``expected_num_sessions``. The request-count + and cancellation caps ARE expressed in the counter, so we also honor + ``can_send_dag_child_turn`` (which includes ``RequestCountStopCondition`` + and the DAG cancellation gate but excludes the session gate). The duration + cap is enforced by the outer ``wait_for`` (Gap-4 cancel-in-flight), not + here, so an in-flight instance is never abandoned mid-replay by this gate. + + Reads the EXPLICIT ``expected_num_sessions`` only -- never the derived + single-pass target -- so this gate is consulted solely on the bounded + (explicit-stop) recycle path. A bare run never reaches here (it takes the + single-pass branch in :meth:`_run_lanes`), so the derived ``N`` is not a + recycle cap. + """ + sessions = getattr(self._config, "expected_num_sessions", None) + if sessions is not None and sessions > 0 and self._admitted_traces >= sessions: + return False + checker = self._stop_checker + if checker is None: + return True + can_send = getattr(checker, "can_send_dag_child_turn", None) + if callable(can_send): + return bool(can_send()) + return True + + async def _run_instance( + self, + trace: Any, + lane_index: int, + recycle_pass: int, + *, + pressure: bool = False, + fresh_start: bool = False, + instance_id: str | None = None, + ) -> bool: + """Run ONE trace instance on ``lane_index`` (recycle pass ``recycle_pass``). + + Returns True iff the instance stopped cleanly at the issuer's stop gate -- + the pressure lane loop's stop signal. + + ``pressure=True`` marks an extended-warmup instance: ``.p{pass}`` id and + the profiling-style graph (see ``_graph_at_t_star``). + + ``fresh_start=True`` marks a profiling lane that WAS a pressure lane but + drained empty: it suppresses the pass-0 t* plan (a full identity t*=0 + replay of the cursor-drawn template, NOT a resume of a t* the pressure + stage already executed against the server) and mints a ``.f{pass}`` id. + The dedicated ``.f`` namespace exists because a plain ``.0`` id can + COLLIDE with the same lane's warmup boundary-priming id when the cursor + wraps a small corpus back onto the lane's own template; the cache-bust + marker digests the full instance id, so a collision would silently warm + a lane that exists to measure cold (agentx mints a fresh uuid + marker + for its empty-lane fresh conversations). The worker strips everything + after the first ``#``, so catalog / mmap keying is unaffected -- the same + argument as ``.p{k}``. + + The instance id ``f"{trace.id}#{lane_index}.{recycle_pass}"`` is stamped + on every credit's ``trace_id`` so the cache-bust marker ROTATES per recycle + AND decorrelates concurrent lanes that wrapped onto the SAME template (C3 + + AgentX's per-lane ``trajectory_index`` in the digest), while the build-time + catalog + graph store mmap stay keyed by the base template id (the + worker strips everything after the first ``#``). The id is GLOBALLY UNIQUE + per (lane, pass), so two lanes wrapping the same template never collide on + the return-routing registry. Pass 0 uses the lane-salted t* plan (C2); + recycle passes (>0) run the full t*=0 replay -- AgentX recycle restarts a + session from turn 0 with no snapshot/t* slice + (the retired agentic-replay plane's recycled-lane dispatch). + + A per-instance adapter is registered under the instance id for return + routing, bound to a ``TraceExecutor``, and awaited. An instance error is + counted (``errored_traces``) and contained -- it does not propagate out of + the lane's ``TaskGroup`` (which would cancel sibling lanes). + """ + refused = False + handoff_entry = ( + self._handoff_for_lane(lane_index) + if recycle_pass == 0 and not pressure + else None + ) + if handoff_entry is not None and trace.id != handoff_entry.template_trace_id: + # _run_lanes declined the template swap (handoff template not in + # this corpus): fall back to the normal pass-0 path COHERENTLY -- + # applying another template's executed set here would chop the + # wrong graph. + handoff_entry = None + if handoff_entry is not None: + # Marker continuity: reuse the live pressure instance's id so the + # per-instance cache-bust marker (digest of credit.trace_id) is + # unchanged across the handoff and the KV the pressure stage built + # at this id transfers. The pressure instance's adapter was reaped at + # warmup teardown, so re-registering the id in this profiling + # strategy's fresh registry cannot collide; profiling instances mint + # fresh nonces, so no collision there either. + instance_id = handoff_entry.instance_id + elif instance_id is None: + # {template}::{nonce}: instance identity is template + fresh nonce + # (uuid4 -- collision-proof across concurrent lanes, recycles, + # phases, and runs). Lane / pass / flavor are diagnostics, logged + # below instead of encoded in the id. Callers that pre-mint (the + # pressure lane loop stashes the id in _pressure_live BEFORE the + # await) pass theirs in so credits and bookkeeping share ONE id. + instance_id = f"{trace.id}::{uuid.uuid4().hex}" + flavor = "p" if pressure else ("f" if fresh_start else "") + self.debug( + lambda t=trace.id, + i=instance_id, + ln=lane_index, + k=recycle_pass, + f=flavor: (f"instance {i} (template={t} lane={ln} pass={f}{k})") + ) + plan = ( + self._plan_for_lane(trace, lane_index) + if recycle_pass == 0 and handoff_entry is None and not fresh_start + else None + ) + adapter = self._build_adapter( + trace.id, + instance_id, + first_token_sources=self._first_token_sources_for(trace), + node_identity=self._node_identity_for(trace), + ) + self._adapters[instance_id] = adapter + self._admitted_traces += 1 + if not pressure and not fresh_start and recycle_pass == 0: + # Boundary-priming instance for (template, lane): the pressure + # handoff's pass-0 merge needs this id (no longer reconstructable). + self._priming_instance_ids[(trace.id, lane_index)] = instance_id + parsed: ParsedGraph | None = None + try: + if handoff_entry is not None: + parsed, run_trace = self._graph_at_handoff(trace, handoff_entry) + else: + parsed, run_trace = self._graph_at_t_star( + trace, plan, pressure=pressure + ) + executor = TraceExecutor( + parsed, + credit_issuer=adapter, + compress_edge_delays=self._accelerated_warmup, + # Every live stream's frontier turn must fire at its ABSOLUTE + # offset from t*, not relative to when + # its inputs arrive. Anchor node-level leading offsets to the + # shared instance run-start so co-scoped subagent/worker streams + # interleave in recorded-time order instead of + # drifting behind their spawn-parent. Correct for every phase of + # this strategy: WARMUP primes boundary turns at min_start_delay + # 0 (anchor is a no-op) and t*=0 full replay anchors to the trace + # start, which IS the run-start. + absolute_start_offsets=True, + ) + await executor.run(run_trace) + except Exception as exc: + refusal = _leaf_credit_refusal(exc) + if refusal is not None: + # Issuer refusal is a HEALTHY stop (request-count / duration cap + # reached, or run cancelled), not a trace error: log quietly and + # keep ``errored_traces`` untouched. Signal the pressure lane loop + # to end this lane -- a fresh instance would refuse identically. + refused = True + self.debug( + lambda exc=refusal, iid=instance_id: ( + f"graph instance {iid!r} " + f"stopped cleanly at the issuer's stop gate: {exc!r}" + ) + ) + else: + self._errored_traces += 1 + self.warning( + lambda exc=exc, iid=instance_id: ( + f"graph instance {iid!r} unwound with error: {exc!r}" + ) + ) + finally: + self._completed_traces += 1 + async with self._registry_lock: + self._parent_done.add(instance_id) + self._release_adapter_if_idle(instance_id) + return refused + + def _release_adapter_if_idle(self, instance_id: str) -> None: + """Pop ``instance_id``'s adapter once its parent finished and it is idle. + + A mid-run drain of a still-running parent's dispatches must not reap + the adapter (``instance_id`` not yet in ``_parent_done``); the parent + finally and the last return's ``on_drained`` both funnel here so the + pop happens exactly once, whichever lands second. + + Idempotent: a second call after the pop is a no-op. Callers MUST hold + ``_registry_lock``. + """ + if instance_id not in self._parent_done: + return + adapter = self._adapters.get(instance_id) + if adapter is not None and adapter.inflight_count > 0: + return + self._adapters.pop(instance_id, None) + self._parent_done.discard(instance_id) + if adapter is not None: + self._send_graph_trace_end(adapter) + + def _on_adapter_drained(self, adapter: CreditDispatchAdapter) -> None: + """Adapter callback: its last in-flight dispatch just returned. + + Fires synchronously from ``adapter.resolve`` on the event loop. We cannot + ``await`` the registry lock here, so schedule a lock-guarded retry of the + deferred pop; the adapter's ``instance_id`` is its de-mux registry key. + Drives the idle-pop for an instance whose parent finally already ran + with returns still in flight. A spurious schedule (adapter already + popped) is harmless. + """ + + async def _release(instance_id: str = adapter.instance_id) -> None: + async with self._registry_lock: + self._release_adapter_if_idle(instance_id) + + task = asyncio.get_running_loop().create_task(_release()) + self._release_tasks.add(task) + task.add_done_callback(self._release_tasks.discard) + + @staticmethod + def _is_trie_graph(parsed: ParsedGraph) -> bool: + """True iff ``parsed`` is a trie graph (the flat LlmNode IR the weka and + dynamo adapters emit). + + The trie builder stamps ``metadata["trie"]`` on every emitted ``LlmNode``. + Detecting the trie marker on any top-level node confirms the trie path; a + non-trie parse reaching ``t*>0`` is a lowering bug (raises in + ``_graph_at_t_star``). + """ + for node in parsed.graph.nodes.values(): + if "trie" in getattr(node, "metadata", {}): + return True + return False + + def _handoff_for_lane(self, lane_index: int) -> LaneHandoff | None: + """This profiling lane's extended-warmup resume entry, if any. + + Only pass-0 profiling instances consume the handoff (the gate lives at + the call sites); recycle passes always run fresh t*=0 templates, + matching agentx's post-handoff recycle draws. ``self._warmup_handoff`` + is only ever non-None on a PROFILING strategy (the runner pops it there), + so this returns None for every warmup/pressure phase with no extra check. + """ + if self._warmup_handoff is None: + return None + return self._warmup_handoff.lanes.get(lane_index) + + def _graph_at_handoff( + self, trace: TraceRecord, entry: LaneHandoff + ) -> tuple[ParsedGraph, TraceRecord]: + """Reconstruct the profiling resume graph at the warmup drain frontier. + + The frontier chop drops pre-t* history AND every node the pressure + stage executed, re-rooting each chain's first not-yet-executed node + with its residual delay -- see ``chop_trie_at_frontier``. The trie + envelope keeps each node's FULL prompt prefix, so the worker + materializes the exact resume prompt with no back-seeding. + + A fully-executed handoff template yields an EMPTY chop: the executor + finalizes instantly and the lane proceeds to its normal recycle draw -- + no special case needed. + """ + parsed = parsed_for_trace(self._parsed, trace) + if (entry.executed_node_ids or entry.t_star_us > 0) and not self._is_trie_graph( + parsed + ): + raise RuntimeError( + f"trace {trace.id!r}: the extended-warmup handoff requires a " + "trie-stamped graph (metadata['trie'] on every LlmNode); a " + "non-trie parse reaching the handoff is a lowering bug" + ) + assert self._warmup_handoff is not None + cap_s = Environment.GRAPH.HANDOFF_RESIDUAL_CAP + rewritten = chop_trie_at_frontier( + parsed, + t_star_us=entry.t_star_us, + executed=entry.executed_node_ids, + return_wall_us=entry.return_wall_us, + drain_end_wall_us=self._warmup_handoff.drain_end_wall_us, + residual_cap_us=cap_s * MICROS_PER_SECOND if cap_s is not None else None, + ) + # --burst-phase-starts collapses each resumed lane's leading offset + # to 0 at the profiling start, residuals included. + if self._burst_phase_starts: + rewritten = self._burst_collapse_leading_offsets(rewritten) + return rewritten, trace + + def _graph_at_t_star( + self, trace: TraceRecord, plan: GraphTrace | None, *, pressure: bool = False + ) -> tuple[ParsedGraph, TraceRecord]: + """Reconstruct the per-trace graph + trace at this instance's t* disposition. + + ``plan`` is the lane-resolved t* plan for this instance (lane-salted on + pass 0, ``None`` == full t*=0 replay for recycle passes). Returns + ``(parsed_to_run, trace_to_run)`` for the ``TraceExecutor``. + + PROFILING: ``t*==0`` (default full-replay window, or any recycle pass) + => IDENTITY (byte-identical to the original). ``t*>0`` => + ``chop_trie_at_tstar`` (a frontier chop re-rooting each live chain at + the ``t*`` frontier) for a trie graph. Surviving nodes keep their ids so + the adapter resolves the unmodified catalog ordinal. + + WARMUP: :func:`rewrite_for_warmup` -- the flat boundary-priming graph + (one boundary turn per chain live at t*, START-rooted, zero leading + offsets). ``t*<=0`` yields an EMPTY warmup graph so the instance + finalizes immediately (the ``timing.config`` auto-warmup contract). + + A non-trie graph at ``t*>0`` is a lowering bug (raises). Multi-graph + workloads project onto each trace's OWN graph via ``parsed_for_trace`` + first (else a non-first trace runs the first file's topology). + + PRESSURE (extended warmup): ``pressure=True`` routes a WARMUP-phase + instance through the PROFILING branches (t*>0 chop / t*=0 identity) so + the post-t* remainder replays compressed under the warmup token cap -- + never ``rewrite_for_warmup``. + """ + parsed = parsed_for_trace(self._parsed, trace) + t_star_us = plan.t_star_us if plan is not None else 0 + is_warmup = self._phase_variant() == "warmup" and not pressure + if t_star_us <= 0: + if is_warmup: + return rewrite_for_warmup(parsed, 0), trace + return parsed, trace + + # Trie graphs (the flat LlmNode + StaticEdge recorded-trace IR) snapshot via a + # simple frontier chop -- there are no reducers / spawn / await / subgraph + # primitives to re-root, and each node's full pre-t* prompt prefix is + # preserved (the worker materializes the exact resume prompt). The + # executor anchors the re-rooted frontier offsets via the + # ``absolute_start_offsets=True`` already set on every instance. + if not self._is_trie_graph(parsed): + raise RuntimeError( + f"trace {trace.id!r}: t*={t_star_us} requires a trie-stamped graph " + "(metadata['trie'] on every LlmNode); every live producer lowers " + "onto the trie, so a non-trie parse reaching t*>0 is a lowering bug" + ) + if is_warmup: + return rewrite_for_warmup(parsed, t_star_us), trace + rewritten = chop_trie_at_tstar(parsed, t_star_us) + if self._burst_phase_starts: + rewritten = self._burst_collapse_leading_offsets(rewritten) + return rewritten, trace + + def _burst_collapse_leading_offsets(self, rewritten: ParsedGraph) -> ParsedGraph: + """Collapse leading phase-start offsets (AgentX ``--burst-phase-starts``). + + AgentX burst collapses the phase START into a synchronized burst: every + trace's earliest profiling resume fires at once, IGNORING the per-stream + leading offset from t*. On a chopped trie graph that leading offset + lives on each re-rooted node's START in-edge ``min_start_delay_us`` + (stamped by ``snapshot_chop._chop_edges``; the node-level field is never + stamped by the trie producers but is collapsed too for hand-authored + graphs) -- delegate to the pure + :func:`aiperf.graph.scheduler.collapse_leading_start_offsets`. The + inter-turn ``StaticEdge.delay_after_predecessor_us`` end-to-start gaps + are UNTOUCHED -- burst governs only the phase start, not the faithful + inter-turn pacing. Warmup builds its boundary graph offset-free + (:func:`rewrite_for_warmup`), keeping spread/burst warmup identical. + """ + return msgspec.structs.replace( + rewritten, graph=collapse_leading_start_offsets(rewritten.graph) + ) + + def _first_token_sources_for(self, trace: Any) -> frozenset[str]: + """Source node ids of ``trace``'s first-token-anchored edges (cached). + + Computed from the trace's OWN projected graph (``parsed_for_trace``) -- + the same per-trace topology the executor dispatches -- and cached per + template trace id (lane/t*-independent). This matches the executed graph + exactly for the default t*=0 replay; under a t*>0 snapshot chop the set + is a safe superset (a chop only DROPS edges, so a source whose anchor + edge was chopped no longer dispatches -- its stray ``first_token_event`` + flag is inert -- and no chopped graph can contain an anchor edge absent + from the projection). + """ + cached = self._first_token_sources_cache.get(trace.id) + if cached is None: + graph = parsed_for_trace(self._parsed, trace).graph + cached = first_token_sources(graph) + self._first_token_sources_cache[trace.id] = cached + return cached + + def _node_identity_for( + self, trace: Any + ) -> dict[str, tuple[int, str | None]] | None: + """``trace``'s ``node_id -> (agent_depth, parent_node_id)`` map (cached). + + Reads the dag_jsonl lowering's ``metadata["dag"]`` legacy-identity stamp + off the trace's OWN projected graph (``parsed_for_trace``), cached per + template trace id -- the identity is lane/t*-independent (a chop / + warmup rewrite keeps surviving node ids, so a superset map is safe). + Returns ``None`` when no node carries the stamp (weka/dynamo) so the + adapter's root-chain fallback stays byte-identical. + """ + if trace.id in self._node_identity_cache: + return self._node_identity_cache[trace.id] + graph = parsed_for_trace(self._parsed, trace).graph + identity = { + nid: (m["agent_depth"], m.get("parent_node")) + for nid, node in graph.nodes.items() + if (m := (node.metadata or {}).get("dag")) + } + result = identity or None + self._node_identity_cache[trace.id] = result + return result + + def _build_adapter( + self, + trace_id: str, + instance_id: str, + *, + first_token_sources: frozenset[str] = frozenset(), + node_identity: dict[str, tuple[int, str | None]] | None = None, + ) -> CreditDispatchAdapter: + """Construct a fresh per-instance ``CreditDispatchAdapter``. + + ``trace_id`` is the BASE template id keying the catalog + graph store mmap; + ``instance_id`` (``{trace_id}::{nonce}``) is stamped on every + credit's ``trace_id`` for marker rotation + return de-mux. Routing keys + are per TRAJECTORY: the adapter lazily mints one + ``{conversation_id}::{nonce}`` x_correlation_id per scope, so + concurrent instances never share a key. ``first_token_sources`` names + the nodes whose dispatch emits a ``FirstToken`` (post-TTFT anchoring). + ``node_identity`` is the trace's dag legacy-identity map (agent_depth / + parent corr stamping); None for stamp-free producers. + """ + return CreditDispatchAdapter( + credit_issuer=self._credit_issuer, + catalog_context=self._catalog, + trace_id=trace_id, + instance_id=instance_id, + phase_variant=self._phase_variant(), + dispatch_timeout_s=self._dispatch_timeout_s, + on_drained=self._on_adapter_drained, + first_token_sources=first_token_sources, + node_identity=node_identity, + ) + + def _phase_variant(self) -> str: + """Graph phase-variant label for this phase. + + ``"warmup"`` for a WARMUP phase, ``"profiling"`` otherwise. Falls back to + ``"profiling"`` when no config / phase is wired (e.g. unit harness). + """ + from aiperf.common.enums import CreditPhase + + phase = getattr(self._config, "phase", None) + if phase == CreditPhase.WARMUP: + return "warmup" + return "profiling" + + async def handle_credit_return(self, credit: Credit) -> None: + """No-op for graph credits. + + Graph returns are resolved by the UNCONDITIONAL observer + (``_on_graph_return`` -> ``adapter.resolve``) BEFORE this gated path is + reached, so there is nothing to do here. We do NOT route graph returns + through this method: it is skipped once ``can_send_any_turn()`` is False, + which would strand already-issued dispatch Futures. + """ + return + + def _send_graph_trace_end(self, adapter: CreditDispatchAdapter) -> None: + """Fire-and-forget the trace's sticky-lifecycle close (GraphTraceEnd). + + Called from the adapter-reap points -- the successful idle-pop (all + in-flight dispatches drained) and phase teardown for retained adapters + -- NOT the per-instance finally, which can run while credits are still + in flight (duration-end local cancel, dispatch timeout, watchdog, error + unwind). ONE call per instance: graph sessions key on the instance + trace_id, so the router closes the whole instance synchronously before + forwarding to the worker; the chain is idempotent, and the worker pool + has an LRU backstop for lost sends. + """ + end = getattr(self._credit_issuer, "end_graph_trace", None) + if end is None: + return + + task = asyncio.create_task(end(adapter.instance_id, adapter.phase_variant)) + self._trace_end_tasks.add(task) + task.add_done_callback(self._trace_end_tasks.discard) + + def _detach_observer( + self, + unregister: Callable[[Any], None] | None, + register: Callable[[Any], None] | None, + own_observer: Callable[..., None], + label: str, + ) -> None: + """Best-effort detach of one shared observer slot at teardown. + + Prefers the compare-and-clear channel (``unregister(own_observer)``: + the handler clears the slot only if OUR observer is still installed -- + see ``CreditCallbackHandler.clear_graph_return_observer``). Falls back + to the unconditional ``register(None)`` when no unregister + channel was wired (unit harness). Exceptions are swallowed at debug: + teardown must never mask the phase's own exit path. + """ + try: + if unregister is not None: + unregister(own_observer) + elif register is not None: + register(None) + except Exception as exc: + self.debug(lambda exc=exc: f"{label} observer detach failed: {exc!r}") + + def _stash_pressure_handoff(self) -> None: + """Convert the drained pressure state into the profiling handoff. + + Runs at teardown -- for a non-seamless WARMUP phase that is AFTER every + issued credit returned (``PhaseRunner._run_strategy`` awaits + returning-complete before ``run()``'s finally), so the ledger is + complete and ``drain_end`` stamped here credits the whole drain wait + toward each recorded gap (agentx ``finalize_phase``'s + ``finalized_at_ns`` parity). Pass-0 pressure instances merge their + lane's Stage-A boundary-priming walls so chains the pressure never + advanced still anchor their residual on the priming return. + + Completeness gate: two runner paths reach teardown BEFORE every return + landed -- an external cancel (``PhaseRunner._run_strategy``'s cancelled + early-return, right after sending-complete) and a USER warmup phase with + a finite + ``--warmup-grace-period`` (grace-timeout force-complete). A handoff + built from an incomplete ledger would mark server-executed nodes as + not-executed and profiling would REFIRE them, so skip the stash + entirely and let profiling start from the plain t* plans. + """ + channel = self._graph_channel + if channel is None: + self.debug("pressure handoff skipped: no graph channel wired") + return + all_returned = getattr(self._credit_issuer, "graph_all_returned", None) + if all_returned is not None and not all_returned(): + self.notice( + "WARMUP cache pressure handoff skipped: not every warmup " + "credit return landed (cancelled run or finite grace period); " + "profiling will start from the plain t* plans" + ) + return + drain_end = self._wall_us() + lanes: dict[int, LaneHandoff] = {} + for lane, ( + template_id, + instance_id, + t_star_us, + pressure_pass, + ) in self._pressure_live.items(): + live_walls = self._return_walls.get(instance_id, {}) + anchor_walls = dict(live_walls) + if pressure_pass == 0: + # Merge the lane's boundary-PRIMING walls under the pressure + # pass-0 anchor set. Ids carry nonces, so the priming id is + # looked up from the record made at its _run_instance. + priming_id = self._priming_instance_ids.get((template_id, lane)) + priming = self._return_walls.get(priming_id, {}) if priming_id else {} + for node_id, wall in priming.items(): + anchor_walls.setdefault(node_id, wall) + lanes[lane] = LaneHandoff( + template_trace_id=template_id, + instance_id=instance_id, + t_star_us=t_star_us, + executed_node_ids=frozenset(live_walls), + return_wall_us=anchor_walls, + ) + channel.warmup_handoff = GraphWarmupHandoff( + lanes=lanes, + drain_end_wall_us=drain_end, + corpus_cursor=self._pressure_next_index, + pressure_lane_count=self._pressure_lane_count, + ) + self.notice( + f"WARMUP cache pressure handoff: persisted {len(lanes)} live " + "lane(s) for the profiling resume" + ) + + async def teardown_phase(self) -> None: + """Detach the graph-return observer after the phase finalizes. + + Invoked by ``PhaseRunner`` in its ``run()`` ``finally`` (see + ``PhaseTeardownStrategyProtocol``); tests may also call it directly. + Best-effort: a subsequent phase / cleanup must not dispatch into this + torn-down strategy's adapter registry. Also reaps any de-mux entry retained + for an instance not popped at its finally -- the phase is over, so no further + return will arrive and the registry must not leak into the next phase -- + and closes the sticky lifecycle for every retained adapter. + + When the extended-warmup pressure stage ran, the FIRST action here (a + sync prefix, before any awaited sticky close that a cancel could + interrupt) is to build and stash the WARMUP -> PROFILING handoff from + the drained pressure state. + + The SYNC parts (observer detach + registry clear) run FIRST so they + complete even when the phase is being cancelled and the awaited sticky + closes below get interrupted at their first suspension point. + + Detach is compare-and-clear, NOT unconditional: for seamless non-final + phases this teardown is deferred to the background return-wait + completion and can fire AFTER the next phase's ``setup_phase`` installed + ITS observer on the same shared ``CreditCallbackHandler`` slot. Clearing + unconditionally here would drop every subsequent graph return of the + live phase, so the slot is cleared only if OUR observer is still the + one installed (unit harnesses without the unregister channel keep the + unconditional ``register(None)`` detach). + """ + if self._pressure_active: + self._stash_pressure_handoff() + retained = list(self._adapters.values()) + self._adapters.clear() + self._parent_done.clear() + self._detach_observer( + self._unregister_observer, + self._register_observer, + self._on_graph_return, + "graph-return", + ) + self._detach_observer( + self._unregister_first_token_observer, + self._register_first_token_observer, + self._on_graph_first_token, + "first-token", + ) + end = getattr(self._credit_issuer, "end_graph_trace", None) + if end is not None: + for adapter in retained: + try: + await end(adapter.instance_id, adapter.phase_variant) + except Exception as exc: + self.debug(lambda exc=exc: f"teardown trace-end failed: {exc!r}") diff --git a/src/aiperf/timing/strategies/user_centric_rate.py b/src/aiperf/timing/strategies/user_centric_rate.py index e3a5ca408e..3a316496bd 100644 --- a/src/aiperf/timing/strategies/user_centric_rate.py +++ b/src/aiperf/timing/strategies/user_centric_rate.py @@ -407,7 +407,11 @@ async def handle_credit_return( raise ValueError( f"User not found for x_correlation_id: {credit.x_correlation_id}" ) - turn = TurnToSend.from_previous_credit(credit) + # Pass the next turn's metadata through so has_forks / has_branches + # derive from the real metadata (a spawning turn must never stamp + # is_tree_final; a branching turn must stay conservative). + next_meta = self._conversation_source.get_next_turn_metadata(credit) + turn = TurnToSend.from_previous_credit(credit, next_meta) # If the next turn time already passed, the max() will # re-align their schedule to account for the delay. diff --git a/src/aiperf/transports/aiohttp_transport.py b/src/aiperf/transports/aiohttp_transport.py index bfebab23c7..9c3c70817d 100644 --- a/src/aiperf/transports/aiohttp_transport.py +++ b/src/aiperf/transports/aiohttp_transport.py @@ -38,6 +38,7 @@ BaseTransport, FirstTokenCallback, TransportMetadata, + effective_streaming, ) @@ -186,7 +187,7 @@ def get_transport_headers(self, request_info: RequestInfo) -> dict[str, str]: """ accept = ( "text/event-stream" - if request_info.model_endpoint.endpoint.streaming + if effective_streaming(request_info) else "application/json" ) headers: dict[str, str] = {"Accept": accept} @@ -238,7 +239,7 @@ def get_url(self, request_info: RequestInfo) -> str: endpoint_metadata = plugins.get_endpoint_metadata(endpoint_info.type) endpoint_path = endpoint_metadata.endpoint_path if ( - self.model_endpoint.endpoint.streaming + effective_streaming(request_info) and endpoint_metadata.streaming_path is not None ): endpoint_path = endpoint_metadata.streaming_path @@ -278,7 +279,7 @@ def _dedup_path_overlap(base_path: str, sub_path: str) -> str: async def send_request( self, request_info: RequestInfo, - payload: dict[str, Any], + payload: dict[str, Any] | bytes, *, first_token_callback: FirstTokenCallback | None = None, ) -> RequestRecord: @@ -324,7 +325,9 @@ async def send_request( == RequestContentType.MULTIPART_FORM_DATA ) body: bytes | aiohttp.FormData = ( - self._build_form_data(payload) + payload + if isinstance(payload, (bytes, bytearray)) + else self._build_form_data(payload) if use_form_data else orjson.dumps(payload) ) @@ -495,7 +498,7 @@ def _build_form_data(payload: dict[str, Any]) -> aiohttp.FormData: async def _submit_video_job( self, url: str, - payload: dict[str, Any], + payload: dict[str, Any] | bytes, headers: dict[str, str], *, use_form_data: bool = False, @@ -507,7 +510,11 @@ async def _submit_video_job( if self.aiohttp_client is None: raise NotInitializedError("AioHttpClient not initialized") body: bytes | aiohttp.FormData = ( - self._build_form_data(payload) if use_form_data else orjson.dumps(payload) + payload + if isinstance(payload, (bytes, bytearray)) + else self._build_form_data(payload) + if use_form_data + else orjson.dumps(payload) ) record = await self.aiohttp_client.post_request(url, body, headers) result = self._parse_video_response(record, "submit") @@ -619,7 +626,7 @@ async def _download_video_content( async def _send_video_request_with_polling( self, request_info: RequestInfo, - payload: dict[str, Any], + payload: dict[str, Any] | bytes, ) -> RequestRecord: """Send video generation request and poll until complete.""" if self.aiohttp_client is None: diff --git a/src/aiperf/transports/base_transports.py b/src/aiperf/transports/base_transports.py index fb6d1091f5..40433f0b10 100644 --- a/src/aiperf/transports/base_transports.py +++ b/src/aiperf/transports/base_transports.py @@ -35,6 +35,21 @@ """ +def effective_streaming(request_info: RequestInfo) -> bool: + """Per-request wire mode: graph credits may carry a recorded override. + + A graph credit stamps the recorded per-node mode onto + ``RequestInfo.stream_override``; every other request leaves it ``None`` and + follows the global ``endpoint.streaming``. The response READ side is already + per-request (SSE parse keys on the server's ``text/event-stream`` content + type), so this only governs the request-side sites: the Accept header and + the streaming URL path. + """ + if request_info.stream_override is not None: + return request_info.stream_override + return bool(request_info.model_endpoint.endpoint.streaming) + + @runtime_checkable class TransportProtocol(AIPerfLifecycleProtocol, Protocol): """Protocol for a transport that sends requests to an inference server.""" diff --git a/src/aiperf/ui/dashboard/realtime_metrics_dashboard.py b/src/aiperf/ui/dashboard/realtime_metrics_dashboard.py index 2119536107..8f0b740bb2 100644 --- a/src/aiperf/ui/dashboard/realtime_metrics_dashboard.py +++ b/src/aiperf/ui/dashboard/realtime_metrics_dashboard.py @@ -68,7 +68,11 @@ def _should_skip(self, metric: MetricResult) -> bool: INTERNAL and EXPERIMENTAL metrics are already filtered upstream by summarize(), so only ERROR_ONLY and console_group=NONE need filtering here. """ - metric_class = MetricRegistry.get_class(metric.tag) + metric_class = MetricRegistry.get_class_or_none(metric.tag) + if metric_class is None: + # Unregistered tags (e.g. externally injected results) render with + # their own header/unit rather than killing the whole table. + return False if metric_class.has_flags(MetricFlags.ERROR_ONLY): return True return ( @@ -96,11 +100,7 @@ def update(self, metrics: list[MetricResult]) -> None: metrics = [ metric - for metric in sorted( - metrics, - key=lambda m: MetricRegistry.get_class(m.tag).display_order - or sys.maxsize, - ) + for metric in sorted(metrics, key=lambda m: self._display_order(m.tag)) if not self._should_skip(metric) ] _logger.debug(lambda: f"Updating metrics table with {len(metrics)} metrics") @@ -120,6 +120,14 @@ def update(self, metrics: list[MetricResult]) -> None: row_key = self.data_table.add_row(*row_cells) self._metric_row_keys[metric.tag] = row_key + @staticmethod + def _display_order(tag: str) -> int: + """Sort key for a metric tag; unregistered or unordered tags sort last.""" + metric_class = MetricRegistry.get_class_or_none(tag) + if metric_class is None: + return sys.maxsize + return metric_class.display_order or sys.maxsize + def _update_single_row(self, row_cells: list[Text], row_key: RowKey) -> None: """Update a single row's cells.""" for col_name, cell_value in zip(self.COLUMNS, row_cells, strict=True): @@ -138,10 +146,17 @@ def _format_metric_row(self, metric: MetricResult) -> list[Text]: Note: Metrics are pre-converted to display units by summarize(), so values can be used directly without conversion. """ - metric_class = MetricRegistry.get_class(metric.tag) - short_header = metric_class.short_header or metric_class.header + metric_class = MetricRegistry.get_class_or_none(metric.tag) + if metric_class is not None: + short_header = metric_class.short_header or metric_class.header + hide_unit = metric_class.short_header_hide_unit + else: + # Unregistered tag: fall back to the display metadata carried on the + # MetricResult itself so one unknown tag cannot kill the table. + short_header = metric.header or metric.tag + hide_unit = False # Use the metric's unit directly (already converted to display unit) - if not metric_class.short_header_hide_unit and metric.unit: + if not hide_unit and metric.unit: short_header = f"{short_header} ({metric.unit})" return [ Text( diff --git a/src/aiperf/workers/inference_client.py b/src/aiperf/workers/inference_client.py index 27d2098edb..67837d6c64 100644 --- a/src/aiperf/workers/inference_client.py +++ b/src/aiperf/workers/inference_client.py @@ -20,6 +20,8 @@ from aiperf.common.redact import redact_headers from aiperf.plugin import plugins from aiperf.plugin.enums import PluginType, TransportType +from aiperf.transports.base_transports import effective_streaming +from aiperf.workers.session_routing import RoutingContext, SessionRoutingBase if TYPE_CHECKING: from aiperf.transports.base_transports import FirstTokenCallback @@ -72,6 +74,22 @@ def __init__( # Resolved by the worker via record payload-retention auto-detection. self.strip_record_payload_bytes = strip_record_payload_bytes + # Session-routing plugin (selected via --session-routing): one instance + # per worker, invoked at the request-serialization chokepoint to stamp + # per-session identity (headers and/or body). None when routing is off. + self._routing: SessionRoutingBase | None = None + self._routing_mode: str | None = None + self._warned_bytes_routing = False + endpoint_info = model_endpoint.endpoint + if endpoint_info.session_routing is not None: + routing_cls = plugins.get_class( + PluginType.SESSION_ROUTING, endpoint_info.session_routing + ) + self._routing = routing_cls( + routing_cls.Options(**endpoint_info.session_routing_opts) + ) + self._routing_mode = endpoint_info.session_routing + # Detect and set transport type if not explicitly set if not model_endpoint.transport: model_endpoint.transport = TransportType( @@ -89,6 +107,34 @@ def __init__( self.transport = TransportClass(model_endpoint=self.model_endpoint) self.attach_child_lifecycle(self.transport) + @property + def session_routing_active(self) -> bool: + """True when a --session-routing plugin owns per-session identity.""" + return self._routing is not None + + def notify_session_end(self, x_correlation_id: str) -> None: + """Post-session pass-through to the routing plugin (idempotent hook). + + Called by the worker terminal-eviction paths on ANY terminal outcome of + a session. On this codebase those are: a successful final turn, a + cancellation, and a cancel-before-start (the done-callback path whose + finally block never runs). Idempotency is the plugin's responsibility -- + this hook does not dedupe. No-op when session routing is unset. + + A plugin exception is logged (naming the plugin and session) and + swallowed: this cleanup hook must never break the worker's core + session-eviction lifecycle. + """ + if self._routing is None: + return + try: + self._routing.on_session_end(x_correlation_id) + except Exception as e: + self.warning( + f"session-routing plugin {self._routing_mode!r} on_session_end " + f"failed for session {x_correlation_id!r}; continuing eviction: {e!r}" + ) + async def _send_request_to_transport( self, request_info: RequestInfo, @@ -113,13 +159,75 @@ async def _send_request_to_transport( """ request_info.endpoint_headers = self.endpoint.get_endpoint_headers(request_info) request_info.endpoint_params = self.endpoint.get_endpoint_params(request_info) + # Session-routing chokepoint: build the per-request routing context once + # and let the plugin stamp its headers now (merged onto the endpoint + # headers). The same context feeds the structured body transform below. + routing_ctx: RoutingContext | None = None + if self._routing is not None: + routing_ctx = RoutingContext( + x_correlation_id=request_info.x_correlation_id, + parent_correlation_id=request_info.parent_correlation_id, + root_correlation_id=request_info.root_correlation_id, + is_final_turn=request_info.is_final_turn, + is_parent_final=request_info.is_parent_final, + is_tree_final=request_info.is_tree_final, + ) + # Attribute a plugin fault to the routing plugin (not the server): + # this raise is caught by _send_request_internal and becomes an error + # record whose message names the plugin instead of the endpoint. + try: + routing_headers = self._routing.headers(routing_ctx) + except Exception as e: + raise RuntimeError( + f"session-routing plugin {self._routing_mode!r} failed in headers(): {e!r}" + ) from e + request_info.endpoint_headers.update(routing_headers) + + raw_payload_bytes = request_info.turns[-1].raw_payload_bytes raw_payload = request_info.turns[-1].raw_payload - payload = ( - raw_payload - if raw_payload is not None - else self.endpoint.format_payload(request_info) - ) - request_info.payload_bytes = orjson.dumps(payload) + if raw_payload_bytes is not None: + # Pre-serialized body (weka graph-IR bytes path): the bytes ARE valid + # JSON, so send and record them verbatim. orjson.dumps() would + # corrupt payload_bytes into a JSON string and break ISL/raw-export. + # Body-based routing transforms cannot apply to a verbatim-bytes + # payload (the header stamp above still does) -- warn ONCE so a + # body-mutating mode never silently loses its bind/close writes. + if ( + routing_ctx is not None + and self._routing.mutates_body + and not self._warned_bytes_routing + ): + self._warned_bytes_routing = True + self.warning( + f"session-routing mode {self._routing_mode!r} mutates the " + "request BODY, but this workload sends pre-serialized " + "verbatim bytes (graph-IR replay); the body transform is " + "skipped for those requests (headers still apply). Use a " + "header-based mode (e.g. dynamo_headers) for byte-exact " + "graph replay." + ) + payload = raw_payload_bytes + request_info.payload_bytes = raw_payload_bytes + else: + payload = ( + raw_payload + if raw_payload is not None + else self.endpoint.format_payload(request_info) + ) + # Body-based session routing (e.g. Dynamo nvext.session_control): + # overlay onto the structured body, endpoint-agnostic, after the + # payload dict is in hand. transform_body returns a copy, so this + # never mutates a cached Turn.raw_payload dict (the copy-on-write + # contract is load-bearing here). + if routing_ctx is not None and isinstance(payload, dict): + try: + payload = self._routing.transform_body(payload, routing_ctx) + except Exception as e: + raise RuntimeError( + f"session-routing plugin {self._routing_mode!r} failed " + f"in transform_body(): {e!r}" + ) from e + request_info.payload_bytes = orjson.dumps(payload) return await self.transport.send_request( request_info, payload=payload, @@ -190,6 +298,10 @@ async def send_request( if self.is_trace_enabled: self.trace(f"Calling inference API for turn: {request_info.turns[-1]}") record = await self._send_request_internal(request_info, first_token_callback) + # Stamp the per-request effective wire mode as ground truth before the + # downcast in _finalize_request_record. This holds even for error + # records: a mid-stream failure was still a streamed send. + record.streamed = effective_streaming(request_info) # Redact sensitive headers on the request_info now that the transport has # consumed them. This prevents raw credentials from flowing back through # ZMQ messages (which are TRACE-logged as serialised JSON / repr). diff --git a/src/aiperf/workers/session_manager.py b/src/aiperf/workers/session_manager.py index bb1b4db01d..7c8b8f5dc6 100644 --- a/src/aiperf/workers/session_manager.py +++ b/src/aiperf/workers/session_manager.py @@ -10,15 +10,23 @@ from aiperf.common.models.dataset_models import Conversation, Turn -def _compute_is_fork_parent(conversation: Conversation) -> bool: - """True if this conversation declares any FORK-mode branch. - - Stamped onto ``UserSession`` at creation rather than recomputed on - every read because ``conversation.branches`` is dropped on the - PAYLOAD_BYTES context-mode wire round-trip; a lazy read after that - would silently flip the flag to ``False``. +def _count_fork_children(conversation: Conversation) -> int: + """Number of FORK-mode CHILD conversations this conversation declares. + + Counts children, not branch entries: the dag_jsonl loader packs every + fork declared on one turn into a SINGLE ``ConversationBranchInfo`` + whose ``child_conversation_ids`` lists all of them, and each child + pins the parent independently. Stamped onto ``UserSession`` at + creation rather than recomputed on every read because + ``conversation.branches`` is dropped on the PAYLOAD_BYTES context-mode + wire round-trip; a lazy read after that would silently flip the count + to zero. """ - return any(b.mode == ConversationBranchMode.FORK for b in conversation.branches) + return sum( + len(b.child_conversation_ids) + for b in conversation.branches + if b.mode == ConversationBranchMode.FORK + ) class UserSession(AIPerfBaseModel): @@ -71,6 +79,24 @@ class UserSession(AIPerfBaseModel): "``release_fork_child``. Eviction (``evict_if_unpinned``) is a " "no-op while this is non-zero.", ) + fork_children_expected: int = Field( + default=0, + ge=0, + description="Total FORK-mode branches this conversation declares " + "(stamped at ``create_and_store`` alongside ``is_fork_parent``). " + "``release_fork_child`` may collect a pending-eviction parent only " + "after this many children have pinned: a refcount of 0 alone cannot " + "distinguish 'all children joined' from 'a sibling child has not " + "arrived yet' when an earlier child completes its whole conversation " + "before a later child's credit reaches the worker.", + ) + fork_children_pinned: int = Field( + default=0, + ge=0, + description="How many FORK children have pinned (and therefore " + "seeded from) this session so far. Monotonic; compared against " + "``fork_children_expected`` by ``release_fork_child``.", + ) pending_fork_eviction: bool = Field( default=False, description="When True, the parent's terminal turn has already " @@ -166,7 +192,8 @@ def create_and_store( or self._default_context_mode or ConversationContextMode.DELTAS_WITHOUT_RESPONSES ) - is_fork_parent = _compute_is_fork_parent(conversation) + fork_children_expected = _count_fork_children(conversation) + is_fork_parent = fork_children_expected > 0 # FORK seeding hands the parent's accumulated ``turn_list`` to # the child. ``MESSAGE_ARRAY_WITH_RESPONSES`` replaces ``turn_list`` # on every ``advance_turn`` (see below), which would discard the @@ -208,6 +235,7 @@ def create_and_store( turn_list=[], context_mode=context_mode, is_fork_parent=is_fork_parent, + fork_children_expected=fork_children_expected, ) self.store(x_correlation_id, user_session) return user_session @@ -255,6 +283,7 @@ def pin_for_fork_child(self, x_correlation_id: str) -> None: f"{x_correlation_id!r} (parent already evicted before FORK child arrived)" ) session.fork_refcount += 1 + session.fork_children_pinned += 1 def seed_from_parent( self, child_x_correlation_id: str, parent_x_correlation_id: str @@ -289,15 +318,24 @@ def release_fork_child(self, x_correlation_id: str) -> None: practice and must not raise. When ``pending_fork_eviction`` is set (parent's terminal turn - has already fired but was waiting for children to land) and - the refcount drops to 0, the session is evicted in the same - call — there is no other code path that will collect it. + has already fired but was waiting for children to land), the + refcount drops to 0, AND every declared FORK child has pinned + (``fork_children_pinned >= fork_children_expected``), the session + is evicted in the same call — there is no other code path that + will collect it. The pinned-count gate exists because an early + child can complete its WHOLE conversation before a sibling's + credit reaches the worker: refcount alone would hit 0 and evict + the parent out from under the late sibling's seed. """ session = self._cache.get(x_correlation_id) if session is None: return session.fork_refcount = max(0, session.fork_refcount - 1) - if session.fork_refcount == 0 and session.pending_fork_eviction: + if ( + session.fork_refcount == 0 + and session.pending_fork_eviction + and session.fork_children_pinned >= session.fork_children_expected + ): self._cache.pop(x_correlation_id, None) def evict_if_unpinned(self, x_correlation_id: str) -> None: diff --git a/src/aiperf/workers/session_routing.py b/src/aiperf/workers/session_routing.py new file mode 100644 index 0000000000..ca6752620d --- /dev/null +++ b/src/aiperf/workers/session_routing.py @@ -0,0 +1,240 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Session-routing transforms: per-session identity on the wire. + +One plugin category unifies every mechanism that tells an external router +which session a request belongs to, whether header-based (SGLang Model +Gateway routing keys, Dynamo session headers, generic session-ID headers) or +body-based (Dynamo ``nvext.session_control``). The selected plugin is +instantiated once per worker by ``InferenceClient`` and invoked at the +request-serialization chokepoint. + +Contracts: +- Options instances are canonicalized at config resolution; ``self.options`` + is the plugin's own typed model. +- ``transform_body`` must NEVER mutate its input (the structured path + includes cached ``Turn.raw_payload`` dicts shared with the dataset). +- ``on_session_end`` fires strictly AFTER the session's last worker-side + activity, on every terminal path (final turn, cancellation, and + cancel-before-start via the done-callback path). It MUST be idempotent. +- Stateful plugins key instance state on ``ctx.x_correlation_id`` ONLY: + a session tree deliberately spans workers, so tree-keyed worker state + fragments. Tree-scoped behavior uses the stateless per-request facts + (``root_correlation_id`` + ``is_tree_final``) instead. +""" + +from __future__ import annotations + +from abc import ABC +from dataclasses import dataclass +from typing import Any, ClassVar, Generic, Literal, TypeVar + +from pydantic import ConfigDict, Field + +from aiperf.common.models import AIPerfBaseModel + + +class EmptyRoutingOptions(AIPerfBaseModel): + """Options model for parameterless plugins; rejects every opt key.""" + + model_config = ConfigDict(extra="forbid") + + +OptionsT = TypeVar("OptionsT", bound=AIPerfBaseModel) + + +@dataclass(slots=True, frozen=True) +class RoutingContext: + """Per-request identity facts handed to a routing plugin. + + Field naming mirrors ``RequestInfo`` verbatim. ``is_parent_final`` / + ``is_tree_final`` are stamped issuer-side from ``SessionTreeRegistry`` + state and are conservative: ``is_tree_final`` is False whenever + indeterminate, ``is_parent_final`` is None for roots or when unknown. + """ + + x_correlation_id: str + """This session's stable key (same on every turn).""" + parent_correlation_id: str | None + """Immediate parent session's key; None for root sessions.""" + root_correlation_id: str | None + """Session-tree root key, verbatim from RequestInfo (never None on the + dispatch path: the worker passes ``credit.effective_root_correlation_id``).""" + is_final_turn: bool + """True when this is the current session's last request.""" + is_parent_final: bool | None + """True when the parent session had already returned its final turn + at credit-issue time; None for roots or when not determinable.""" + is_tree_final: bool + """Best-effort: True only when this is provably the last request the + whole session tree will send.""" + + +class SessionRoutingBase(ABC, Generic[OptionsT]): # noqa: B024 # ABC marks the plugin protocol; every method has a working default so passthrough subclasses instantiate directly. + """Base for session-routing plugins (``session_routing`` category).""" + + mutates_body: ClassVar[bool] = False + """Protocol metadata: True when ``transform_body`` changes the payload, + marking the mode as body-mutating / incompatible with any verbatim-bytes + request path. The graph-IR bytes path (pre-serialized weka/dynamo node + bodies) skips ``transform_body`` and warns ONCE per worker when the + active mode sets this flag (``InferenceClient``); header-based stamping + still applies to every request, bytes path included.""" + + # ClassVar cannot reference the OptionsT type parameter, so the base type + # is kept here; subclasses narrow it via their Generic parameterization. + Options: ClassVar[type[AIPerfBaseModel]] = EmptyRoutingOptions + """Per-plugin options model, populated from --session-routing-opt + key=value pairs. Every Options model must set extra='forbid'.""" + + def __init__(self, options: OptionsT) -> None: + self.options: OptionsT = options + + def headers(self, ctx: RoutingContext) -> dict[str, str]: + """Extra HTTP headers for this request (merged into endpoint headers).""" + return {} + + def transform_body( + self, payload: dict[str, Any], ctx: RoutingContext + ) -> dict[str, Any]: + """Return a (possibly new) payload dict; never mutate the input.""" + return payload + + def on_session_end(self, x_correlation_id: str) -> None: + """Post-session cleanup: no further requests will be sent for this + session by this worker. Idempotent; default no-op.""" + return None + + +class DynamoHeadersRouting(SessionRoutingBase[EmptyRoutingOptions]): + """Dynamo session affinity via X-Dynamo-Session-ID / X-Dynamo-Parent-Session-ID. + + Pair with a Dynamo frontend running ``--router-session-affinity-ttl-secs``. + """ + + def headers(self, ctx: RoutingContext) -> dict[str, str]: + headers = {"X-Dynamo-Session-ID": ctx.x_correlation_id} + if ctx.parent_correlation_id: + headers["X-Dynamo-Parent-Session-ID"] = ctx.parent_correlation_id + return headers + + +class DynamoNvextOptions(AIPerfBaseModel): + """Options for the deprecated-upstream nvext.session_control transport.""" + + model_config = ConfigDict(extra="forbid") + + timeout_seconds: int = Field( + default=300, + ge=1, + description="Dynamo session_control inactivity timeout carried on " + "every bind (contract=bind) or on the one open (contract=open).", + ) + contract: Literal["bind", "open"] = Field( + default="bind", + description="session_control wire contract. 'bind' (Dynamo >= " + "v1.3.0-dev): re-bind every non-final turn -- idempotent on the " + "router and refreshes the inactivity TTL. 'open' (released Dynamo " + "v1.2.x, whose SessionAction enum accepts only open/close and " + "rejects bind with HTTP 400): open ONCE on the first request this " + "worker sends for the session, then ride bare session_id.", + ) + + +class DynamoNvextRouting(SessionRoutingBase[DynamoNvextOptions]): + """Dynamo session affinity via nvext.session_control request-body metadata. + + Two wire contracts, selected by the ``contract`` opt: + + - ``bind`` (default; Dynamo >= v1.3.0-dev): 'bind' on every non-final + turn (idempotent on the router, refreshes the TTL), 'close' on the + final turn. + - ``open`` (released Dynamo v1.2.x): 'open' (+timeout) on the FIRST + request this worker sends for a session -- 'open' is NOT idempotent, + so opened sessions are tracked per worker; the trigger is + first-request-per-worker, NOT ``turn_index == 0`` (agentic replay's + first request can be a warmup turn). Subsequent turns carry bare + ``session_id``; the final turn sends 'close'. + + Targets Dynamo builds that implement session_control; current upstream + Dynamo main does not (use dynamo_headers there). + """ + + mutates_body: ClassVar[bool] = True + Options: ClassVar[type[AIPerfBaseModel]] = DynamoNvextOptions + + def __init__(self, options: DynamoNvextOptions) -> None: + super().__init__(options) + # Sessions this worker has already opened (contract=open only): + # keyed on ctx.x_correlation_id per the base-class state contract, + # evicted via on_session_end so high-cardinality runs cannot grow it + # unboundedly. + self._opened: set[str] = set() + + def transform_body( + self, payload: dict[str, Any], ctx: RoutingContext + ) -> dict[str, Any]: + if ctx.is_final_turn: + session_control: dict[str, Any] = { + "session_id": ctx.x_correlation_id, + "action": "close", + } + self._opened.discard(ctx.x_correlation_id) + elif self.options.contract == "open": + if ctx.x_correlation_id in self._opened: + # Affinity already bound; bare session_id keeps it sticky. + session_control = {"session_id": ctx.x_correlation_id} + else: + self._opened.add(ctx.x_correlation_id) + session_control = { + "session_id": ctx.x_correlation_id, + "action": "open", + "timeout": self.options.timeout_seconds, + } + else: + session_control = { + "session_id": ctx.x_correlation_id, + "action": "bind", + "timeout": self.options.timeout_seconds, + } + merged = dict(payload) + raw_nvext = merged.get("nvext") + nvext = dict(raw_nvext) if isinstance(raw_nvext, dict) else {} + raw_sc = nvext.get("session_control") + merged_sc = dict(raw_sc) if isinstance(raw_sc, dict) else {} + merged_sc.update(session_control) + nvext["session_control"] = merged_sc + merged["nvext"] = nvext + return merged + + def on_session_end(self, x_correlation_id: str) -> None: + """Release the per-worker opened marker (idempotent).""" + self._opened.discard(x_correlation_id) + + +class SmgRoutingKeyRouting(SessionRoutingBase[EmptyRoutingOptions]): + """SGLang Model Gateway manual-policy stickiness via X-SMG-Routing-Key.""" + + def headers(self, ctx: RoutingContext) -> dict[str, str]: + return {"X-SMG-Routing-Key": ctx.x_correlation_id} + + +class SessionIdHeaderOptions(AIPerfBaseModel): + """Options for the generic additive session-ID header.""" + + model_config = ConfigDict(extra="forbid") + + header_name: str = Field( + default="X-Session-ID", + min_length=1, + description="Header name carrying the per-session correlation ID.", + ) + + +class SessionIdHeaderRouting(SessionRoutingBase[SessionIdHeaderOptions]): + """Generic additive session header for routers expecting a custom name.""" + + Options: ClassVar[type[AIPerfBaseModel]] = SessionIdHeaderOptions + + def headers(self, ctx: RoutingContext) -> dict[str, str]: + return {self.options.header_name: ctx.x_correlation_id} diff --git a/src/aiperf/workers/worker.py b/src/aiperf/workers/worker.py index 7ebc82ef92..5a79ef58aa 100644 --- a/src/aiperf/workers/worker.py +++ b/src/aiperf/workers/worker.py @@ -5,11 +5,14 @@ import asyncio import time import uuid -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any + +import orjson from aiperf.common.base_component_service import BaseComponentService from aiperf.common.constants import BYTES_PER_MIB, WARMUP_SYSTEM_MESSAGE_PREFIX from aiperf.common.enums import ( + CacheBustTarget, CommAddress, CommandType, ConversationBranchMode, @@ -47,8 +50,10 @@ RequestInfo, RequestRecord, SSEMessage, + Turn, WorkerTaskStats, ) +from aiperf.common.models.dataset_models import GraphSegmentClientMetadata from aiperf.common.models.record_models import find_last_non_empty_usage from aiperf.common.protocols import ( PushClientProtocol, @@ -63,11 +68,13 @@ CancelCredits, CreditReturn, FirstToken, + GraphTraceEnd, RouterToWorkerMessage, WorkerReady, WorkerShutdown, ) from aiperf.credit.structs import Credit, CreditContext +from aiperf.dataset.graph_segment_unified_store import GraphSegmentUnifiedClient from aiperf.dataset.protocols import DatasetClientStoreProtocol from aiperf.plugin import plugins from aiperf.plugin.enums import PluginType @@ -76,7 +83,30 @@ from aiperf.workers.session_manager import UserSession, UserSessionManager if TYPE_CHECKING: + from collections.abc import Callable + from aiperf.config.resolution.plan import BenchmarkRun + from aiperf.graph.dynamic_pool import GraphCapturedReply, GraphPoolSentinel + + +def _mint_x_request_id(credit: Any) -> str: + """Fresh per-dispatch request id (``X-Request-ID``), minted per credit. + + Linear credits keep the legacy opaque ``uuid4``. Graph credits mint the + recoverable-plus-nonce form ``{node_id}::{nonce}``: the node id + (``{scope}:{turn}``, the trajectory coordinate) is derived from the + credit's own ``conversation_id`` (``{trace}[::{scope}]``) and + ``turn_index``, and the nonce keeps the id fresh per dispatch. Safe to + enrich because ``x_request_id`` is write-only (never read back, matched, + or deduped) -- it exists so an export row identifies its template node + without any side lookup. + """ + if getattr(credit, "node_ordinal", None) is None or not credit.conversation_id: + return str(uuid.uuid4()) + conversation = credit.conversation_id + _, sep, child_scope = conversation.partition("::") + scope = child_scope if sep else conversation + return f"{scope}:{credit.turn_index}::{uuid.uuid4().hex}" def _phase_needs_first_token_callback(phase) -> bool: @@ -170,6 +200,14 @@ def __init__( self.task_stats: WorkerTaskStats = WorkerTaskStats() self.credit_tasks: dict[int, asyncio.Task] = {} + # Worker-local dynamic-content pool for graph traces (captured + # assistant responses). Engaged only when a + # node envelope carries `capture` -- inert for recorded corpora. + from aiperf.graph.dynamic_pool import GraphDynamicPool + + self._graph_dynamic_pool = GraphDynamicPool( + max_bytes=Environment.GRAPH.DYNAMIC_POOL_MAX_BYTES + ) self.inference_results_push_client: PushClientProtocol = ( self.comms.create_push_client( @@ -227,8 +265,36 @@ def __init__( self._dataset_client: DatasetClientStoreProtocol | None = None self._dataset_configured_event = asyncio.Event() + # Caches a FAILED store open (missing/corrupt files) so the worker does + # not re-attempt the doomed open on every graph credit (F4): a single + # fatal, actionable error is surfaced per credit instead of a silent + # retry-loop that looks like a hang. + self._graph_store_open_error: ErrorDetails | None = None + # Worker-side reader for the graph unified segment store (content pool + + # per-node manifests) -- the SOLE graph store shape; every graph build + # writes one. Opened lazily on the first graph credit from the SAME + # (base_path, benchmark_id) the build-time GraphSegmentUnifiedBackingStore + # wrote to; ``_graph_unified_open_attempted`` guards a single open. + self._graph_unified_client: GraphSegmentUnifiedClient | None = None + self._graph_unified_open_attempted = False + # Why the unified open failed, when the failure is actionable (an + # on-disk store REJECTED as pre-v3 by the A2-strict reader, vs simply + # absent). Folded into the fatal GraphStoreUnavailable message so + # the operator sees "re-parse required" instead of a misleading + # "neither store exists". + self._graph_unified_open_failure: str | None = None + # Graph store location from the dataset broadcast (the graph-typed + # DatasetConfiguredNotification). The worker opens the unified store + # from THIS, never from env conventions -- absence is a recorded failure + # feeding the GraphStoreUnavailable fatal, not a temp-dir fallback. + self._graph_client_metadata: GraphSegmentClientMetadata | None = None + # Detecting first token requires parsing each SSE chunk, so only enable - # FirstToken messages when a downstream consumer needs them. + # FirstToken messages when a downstream consumer needs them. Graph credits + # additionally opt in per-credit via ``credit.first_token_event`` (checked + # at dispatch) for post-TTFT first-token anchoring even when no phase-level + # consumer is active. ``_phase_needs_first_token_callback`` already covers + # prefill-concurrency limiting. self._first_token_observation_enabled: bool = any( _phase_needs_first_token_callback(phase) for phase in self.run.cfg.phases ) @@ -262,6 +328,14 @@ async def _on_dataset_configured(self, msg: DatasetConfiguredNotification) -> No ) self._dataset_client = ClientStoreClass(client_metadata=msg.client_metadata) await self._dataset_client.initialize() + # Graph runs carry the unified-store location on the broadcast; capture + # it so ``_graph_unified_reader`` opens exactly what the DatasetManager + # built (no env re-derivation). Non-graph broadcasts leave this None. + self._graph_client_metadata = ( + msg.client_metadata + if isinstance(msg.client_metadata, GraphSegmentClientMetadata) + else None + ) self.session_manager.set_default_context_mode(msg.metadata.default_context_mode) self._dataset_configured_event.set() self.debug( @@ -310,6 +384,14 @@ async def _on_credit_message(self, message: RouterToWorkerMessage) -> None: self._schedule_credit_drop_task(message) case CancelCredits(): await self._on_cancel_credits_message(message) + case GraphTraceEnd(): + # Sticky-lifecycle close: evict the trace's dynamic pool entry + # (deferred while its credits are still in flight -- their + # capture writes may land after this message on cancelled + # paths). + self._graph_dynamic_pool.trace_end( + message.trace_id, message.phase_variant + ) case _: self.warning( f"Unknown credit message type: {message.__class__.__name__}" @@ -382,6 +464,14 @@ def _on_credit_drop_message_task_done( self.execute_async(self.credit_return_push_client.send(credit_return)) credit_context.returned = True + # Post-session hook for routing plugins: this cancel-before-start path is a + # terminal disposition the finally block never sees (the credit task was + # cancelled before it started), so notify here too. Idempotent by contract. + if credit_context.credit is not None: + self.inference_client.notify_session_end( + credit_context.credit.x_correlation_id + ) + # Explicitly clear references to help refcounting (GC is disabled on workers) credit_context.credit = None credit_context.error = None @@ -464,12 +554,17 @@ async def _process_credit(self, credit_context: CreditContext) -> None: - Subsequent turns: Session retrieved from cache (sticky routing ensures same worker) - Final turn: Session evicted from cache to free memory """ - x_request_id = str(uuid.uuid4()) - x_correlation_id = credit_context.credit.x_correlation_id credit = credit_context.credit - + x_request_id = _mint_x_request_id(credit) + + # First token callback - needed when a phase-level consumer is active + # (prefill concurrency limiting, adaptive-scale SLA) OR when this credit + # requests a per-credit first-token event (post-TTFT first-token + # anchoring). Sends FirstToken to router for prefill slot release and for + # the graph first-token observer. + # Returns True when meaningful content is found to stop looking for first token first_token_callback = None - if self._first_token_observation_enabled: + if self._first_token_observation_enabled or credit.first_token_event: async def first_token_callback(ttft_ns: int, message: SSEMessage) -> bool: parsed = self.inference_client.endpoint.parse_response(message) @@ -481,11 +576,43 @@ async def first_token_callback(ttft_ns: int, message: SSEMessage) -> bool: credit_id=credit.id, phase=credit.phase, ttft_ns=ttft_ns, + trace_id=credit.trace_id, + x_correlation_id=credit.x_correlation_id, + turn_index=credit.turn_index, ) ) credit_context.first_token_sent = True return True + # Graph-IR credits carry a trace_id/node_ordinal: the worker rebuilds the + # node's request from the shared graph store mmap (D1) instead of taking + # the linear session-cache path. Any worker can serve any node. + if credit.trace_id is not None: + await self._process_graph_credit( + credit_context, x_request_id, first_token_callback + ) + return + + await self._process_session_credit( + credit_context, x_request_id, first_token_callback + ) + + async def _process_session_credit( + self, + credit_context: CreditContext, + x_request_id: str, + first_token_callback, + ) -> None: + """Process a linear (non-graph) credit via the sticky-routed session cache. + + The original ``_process_credit`` body: look up / create the session by + ``x_correlation_id`` (seeding fork children from their parent), advance to + the credit's turn, send the request, store the response, and evict on the + final / cancelled turn. Behavior is unchanged -- this is a pure extraction + so ``_process_credit`` stays a thin graph-vs-session dispatcher. + """ + x_correlation_id = credit_context.credit.x_correlation_id + credit = credit_context.credit try: session = self.session_manager.get(x_correlation_id) if session is None: @@ -553,7 +680,7 @@ async def first_token_callback(ttft_ns: int, message: SSEMessage) -> bool: # Mark cancelled before re-raising so finally can evict session credit_context.cancelled = True raise - except Exception as e: + except Exception as e: # capture any request error into the credit context credit_context.error = ErrorDetails.from_exception(e) self.exception(f"Error processing credit: {e!r}") finally: @@ -561,6 +688,565 @@ async def first_token_callback(ttft_ns: int, message: SSEMessage) -> bool: if credit_context.credit.is_final_turn or credit_context.cancelled: self._release_and_evict_for_terminal(credit, x_correlation_id) + def _resolve_graph_session_headers( + self, envelope: dict[str, Any], credit: Credit + ) -> dict[str, str] | None: + """Resolve recorded dynamo session-identity headers for one dispatch. + + Recorded dynamo session ids are stamped verbatim at build time. With a + live --session-routing plugin active, the plugin OWNS session identity + (headers stamped at the chokepoint from the credit's live corr), so the + recorded identity headers are stripped -- forwarding both would put two + conflicting identities on the wire. Otherwise, suffix them per replay + instance so concurrent instances of one trace never share -- or + session-final-evict -- a server session. + """ + from aiperf.graph.worker_materialize import ( + strip_dynamo_session_headers, + uniquify_dynamo_session_headers, + ) + + if self.inference_client.session_routing_active: + return strip_dynamo_session_headers(envelope.get("extra_headers")) + return uniquify_dynamo_session_headers( + envelope.get("extra_headers"), + trace_instance_id=credit.trace_id, + phase_variant=credit.phase_variant, + ) + + async def _process_graph_credit( + self, + credit_context: CreditContext, + x_request_id: str, + first_token_callback, + ) -> None: + """Process a graph-IR credit by materializing its request from the mmap. + + Reads the credit's ``(trace_id, node_ordinal, phase_variant)`` and + rebuilds the node's request payload from the unified interned segment + store (:func:`materialize_graph_request_unified` / + :func:`materialize_graph_request_unified_bytes`) -- the sole graph + store shape. It layers the + run-level endpoint options the verbatim path would otherwise drop + (``endpoint.extra`` / ``stream_options.include_usage``) while keeping + the per-node ``dispatch_overrides`` / ``stream`` winning, wraps it as a + single-turn ``raw_payload``, sends it, and forwards the result. + + A missing/corrupt store sets a fatal ``credit_context.error`` (not a + silent swallow + per-credit retry); a missing node ordinal sets a clean + ``GraphEnvelopeMissing`` error. Every pre-dispatch failure ALSO emits a + synthetic error ``RequestRecord`` (:meth:`_send_graph_error_record`): + the RecordsManager completion barrier counts success+error RECORDS + against the credit-side ``final_requests_completed``, so an errored + credit that produced no record would starve the barrier and hang the + run at "please wait for the results". A catch-all around the body + extends the same guarantee to UNANTICIPATED raisers (corrupt envelope + bytes, materialize bugs, transport raises): the error is attributed on + the context (mirroring ``_process_session_credit``) and a record is + emitted unless the dispatch path already sent one. + """ + # Imported lazily: ``aiperf.graph.__init__`` pulls the full executor + # adapter chain, which is import-order fragile; deferring to the first + # graph credit keeps the worker module importable in any order. + from aiperf.graph.worker_materialize import ( + apply_run_level_payload_options, + materialize_graph_request_unified, + materialize_graph_request_unified_bytes, + read_node_envelope, + stamp_cache_bust_marker, + ) + + credit = credit_context.credit + segment_store = self._graph_store_reader(credit_context) + if segment_store is None: + # Store absent/corrupt: error already set on the context. Do NOT + # send a garbage request; the run surfaces the fatal config error. + await self._send_graph_error_record(credit_context, x_request_id) + return + # ``credit.trace_id`` is the per-recycle INSTANCE id (e.g. ``t-1#0``): + # it rotates the cache-bust marker per recycle pass. The unified store + # is keyed by the BASE/template id, so strip the ``::{nonce}`` instance + # suffix before reading -- every recycle instance of one template reads + # the same build-time manifests. + base_trace_id = credit.trace_id.split("::", 1)[0] + endpoint = self.model_endpoint.endpoint + + # Catch-all boundary: any exception escaping the body must still + # attribute the error on the context (else the CreditReturn reports + # success) AND emit a record (else the RecordsManager barrier starves + # and the run hangs). ``record_emitted`` tracks every emission inside + # the region so a failure AFTER a record landed never double-emits. + record_emitted = False + try: + # Pre-read the node envelope ONCE: routes bytes-vs-dict, carries the + # `capture` flag (and assembly `items`), and is handed to the + # materialize functions so the manifest is never decoded twice. + envelope = read_node_envelope( + segment_store, base_trace_id, credit.node_ordinal, credit.phase_variant + ) + if envelope is None: + self._set_graph_envelope_missing(credit_context, base_trace_id) + await self._send_graph_error_record(credit_context, x_request_id) + record_emitted = True + return + capture = bool(envelope.get("capture")) + has_items = envelope.get("items") is not None + extra_headers = self._resolve_graph_session_headers(envelope, credit) + + # In-flight bracket for the dynamic pool: a GraphTraceEnd arriving + # while this credit is mid-processing defers eviction until the + # bracket closes (a cancelled dispatch's FAILED write must land in a + # live entry, not re-create an evicted one). + pool = self._graph_dynamic_pool + pool.credit_started(credit.trace_id, credit.phase_variant) + try: + # Unified (interned A2) store: every node carries int ``handles``, + # so a ``None`` from the unified fns IS a genuine miss. When no + # content mutation is needed (cache-bust prepends a marker to the + # first user message, which a pre-serialized body cannot do), build + # the request body once from content-pool slices and send it + # verbatim; otherwise take the unified dict path. Slot-carrying + # nodes (``items``) always take the dict path: their messages are + # composed per request from the dynamic pool. + if endpoint.cache_bust == CacheBustTarget.NONE and not has_items: + built = materialize_graph_request_unified_bytes( + segment_store, + base_trace_id, + credit.node_ordinal, + credit.phase_variant, + use_legacy_max_tokens=endpoint.use_legacy_max_tokens, + endpoint=endpoint, + envelope=envelope, + default_model=self.model_endpoint.primary_model_name, + ) + if built is None: + self._set_graph_envelope_missing(credit_context, base_trace_id) + await self._send_graph_error_record( + credit_context, x_request_id + ) + record_emitted = True + return + body, _model, stream = built + request_info = self._build_graph_request_info( + credit_context, + None, + x_request_id, + raw_payload_bytes=body, + stream_override=stream, + extra_headers=extra_headers, + ) + else: + from aiperf.graph.dynamic_pool import GraphPoolMissingError + + try: + payload = materialize_graph_request_unified( + segment_store, + base_trace_id, + credit.node_ordinal, + credit.phase_variant, + use_legacy_max_tokens=endpoint.use_legacy_max_tokens, + envelope=envelope, + default_model=self.model_endpoint.primary_model_name, + slot_resolver=( + ( + lambda ordinal: pool.get( + credit.trace_id, credit.phase_variant, ordinal + ) + ) + if has_items + else None + ), + ) + except GraphPoolMissingError as e: + # Broken stickiness (worker death re-route) or backstop + # eviction: a loud trace error, never a silent omission. + # The dispatch adapter sniffs this prefix. + credit_context.error = ( + f"aiperf.graph.pool_missing: " + f"{credit.trace_id}/{e.src_ordinal}" + ) + await self._send_graph_error_record( + credit_context, x_request_id + ) + record_emitted = True + return + if payload is None: + self._set_graph_envelope_missing(credit_context, base_trace_id) + await self._send_graph_error_record( + credit_context, x_request_id + ) + record_emitted = True + return + env_stream = payload.get("stream") + stream_override = ( + bool(env_stream) if env_stream is not None else None + ) + apply_run_level_payload_options( + payload, + endpoint, + stream_override=stream_override, + skip_endpoint_extra=bool( + envelope.get("endpoint_extra_applied") + ), + ) + if credit.trace_id is not None: + stamp_cache_bust_marker( + payload, + benchmark_id=self.run.benchmark_id, + trace_instance_id=credit.trace_id, + target=endpoint.cache_bust, + ) + request_info = self._build_graph_request_info( + credit_context, + payload, + x_request_id, + stream_override=payload.get("stream"), + extra_headers=extra_headers, + ) + await self._dispatch_graph_request( + request_info, + credit_context, + first_token_callback, + capture=capture, + ) + record_emitted = True + finally: + pool.credit_finished(credit.trace_id, credit.phase_variant) + except Exception as e: + credit_context.error = ErrorDetails.from_exception(e) + self.exception( + f"Error processing graph credit {credit.id} " + f"(trace={credit.trace_id!r} node_ordinal={credit.node_ordinal}): {e!r}" + ) + if not record_emitted: + await self._send_graph_error_record(credit_context, x_request_id) + return + + async def _dispatch_graph_request( + self, + request_info: RequestInfo, + credit_context: CreditContext, + first_token_callback: Callable | None, + *, + capture: bool, + ) -> None: + """Send one materialized graph request; capture the response if flagged. + + The pool write happens strictly BEFORE the CreditReturn (the caller's + task finally), so a successor credit -- issued only after this credit + resolves and sticky-routed to this worker -- always observes the + entry. Every exit stores a value for captured nodes: a structured + :class:`GraphCapturedReply` on success (text plus, for tool_calls + replies, the verbatim assistant message JSON), ``EMPTY`` on a + successful response with no replayable content, ``FAILED`` on a + dispatch error, cancellation, or capture-extraction failure. + """ + from aiperf.graph.dynamic_pool import GraphPoolSentinel + + credit = credit_context.credit + pool_key = (credit.trace_id, credit.phase_variant, credit.node_ordinal) + self.task_stats.total += 1 + try: + record: RequestRecord = await self.inference_client.send_request( + request_info, first_token_callback=first_token_callback + ) + except BaseException: + if capture: + self._graph_dynamic_pool.put(*pool_key, GraphPoolSentinel.FAILED) + raise + if capture: + self._graph_dynamic_pool.put( + *pool_key, self._graph_capture_value(credit_context, record) + ) + await self._send_inference_result_message(record) + if record.error is not None: + credit_context.error = record.error + return + + def _graph_capture_value( + self, credit_context: CreditContext, record: RequestRecord + ) -> GraphCapturedReply | GraphPoolSentinel: + """Extract the pool value for a captured node's response record. + + Plain-text replies capture as a text-only :class:`GraphCapturedReply`; + replies whose endpoint returned ``raw_messages`` (chat ``tool_calls`` / + structured content) also carry the verbatim orjson-serialized + assistant message so downstream splices reproduce the legacy + child-seed rendering byte-for-byte. + """ + from aiperf.graph.dynamic_pool import GraphCapturedReply, GraphPoolSentinel + + if record.error is not None: + return GraphPoolSentinel.FAILED + try: + turn = self.inference_client.endpoint.build_assistant_turn(record) + if turn is None: + return GraphPoolSentinel.EMPTY + if turn.raw_messages: + # A single-message capture cannot faithfully carry more than one + # assistant message (openai_responses can legitimately produce + # several); truncating to raw_messages[0] would silently drop the + # rest, so fail loudly instead. + if len(turn.raw_messages) > 1: + credit_context.error = ( + "aiperf.graph.capture_failed: multi-entry raw_messages " + f"({len(turn.raw_messages)} entries) is not representable " + "in a single-message capture; single-message endpoints only" + ) + return GraphPoolSentinel.FAILED + message = turn.raw_messages[0] + content = message.get("content") + return GraphCapturedReply( + text=content if isinstance(content, str) else "", + message_json=orjson.dumps(message).decode(), + ) + except Exception as e: + credit_context.error = f"aiperf.graph.capture_failed: {e!r}" + return GraphPoolSentinel.FAILED + text = "".join( + content for t in turn.texts or [] for content in (t.contents or []) + ) + if not text: + return GraphPoolSentinel.EMPTY + return GraphCapturedReply(text=text) + + def _build_graph_request_info( + self, + credit_context: CreditContext, + payload: dict[str, Any] | None, + x_request_id: str, + *, + raw_payload_bytes: bytes | None = None, + stream_override: bool | None = None, + extra_headers: dict[str, str] | None = None, + ) -> RequestInfo: + """Wrap a materialized graph payload as a single-turn ``RequestInfo``. + + Two shapes, selected by which arg is set: + + - **dict path** (``payload``): a ``raw_payload`` carrying ``messages``, + per-node ``dispatch_overrides``, ``stream``, and -- when ``--cache-bust`` + is set -- the stamped first-user marker. + - **bytes path** (``raw_payload_bytes``): a pre-serialized body built once + from mmap slices; ``payload`` is ``None``. + + ``Turn.model`` is deliberately left unset on both shapes: the recorded + per-node model rides only the wire body (sent verbatim), while + ``record.model_name`` falls back to the run ``--model`` in + ``_finalize_request_record`` so tokenizer selection behaves like plain + aiperf -- recorded deployment ids (e.g. ``dynamo/org/model-fp8``) are + usually not resolvable tokenizer repos. + + ``stream_override`` carries the recorded per-node wire mode (the FINAL + stamped ``payload["stream"]`` / bytes-path effective stream) onto the + ``RequestInfo`` so the transport picks the matching wire mode per-request + (``effective_streaming``); non-graph paths leave it ``None`` (follow the + global ``endpoint.streaming``). + + Either way the chat endpoint's ``format_payload`` is bypassed and the body + is sent verbatim. + """ + credit = credit_context.credit + turn = Turn( + role="user", + raw_payload=payload, + raw_payload_bytes=raw_payload_bytes, + # Per-node HTTP headers from the envelope (dynamo session + # identity: x-dynamo-session-id / -parent-session-id / + # -session-final). The transport merges the last turn's + # extra_headers into the request headers; body is untouched. + extra_headers=extra_headers, + ) + return RequestInfo( + model_endpoint=self.model_endpoint, + credit_num=credit.id, + credit_phase=credit.phase, + cancel_after_ns=credit.cancel_after_ns, + x_request_id=x_request_id, + x_correlation_id=credit.x_correlation_id, + conversation_id=credit.conversation_id, + turn_index=credit.turn_index, + turns=[turn], + drop_perf_ns=credit_context.drop_perf_ns, + credit_issued_ns=credit.issued_at_ns, + is_final_turn=credit.is_final_turn, + agent_depth=credit.agent_depth, + parent_correlation_id=credit.parent_correlation_id, + # Session-routing identity facts: the graph adapter mints real + # per-trajectory num_turns (is_final_turn = the recorded + # session-final fact) and stamps the instance's root trajectory + # corr; finality stays conservative (no SessionTreeRegistry on + # the graph plane). + root_correlation_id=credit.effective_root_correlation_id, + is_parent_final=credit.is_parent_final, + is_tree_final=credit.is_tree_final, + url_index=credit.url_index, + stream_override=stream_override, + ) + + def _graph_unified_reader(self) -> GraphSegmentUnifiedClient | None: + """Return the lazily-opened unified store client, or ``None`` on failure. + + Cached across credits: the doomed/successful open is attempted exactly + once (``_graph_unified_open_attempted``) and reused. The unified client + carries both the addressing face (``get_node_envelope``) and the content + face (``materialize_handles`` / ``build_request_body_handles``); + ``_graph_store_reader`` returns this ONE + client when the store exists on disk. A failed open is cached as + ``None``; the A2-strict ``ValueError`` (an on-disk pre-v3 store rejected + with "re-parse required") additionally remembers its reason so the + fatal error path reports it instead of claiming no store exists. + """ + if self._graph_unified_open_attempted: + return self._graph_unified_client + self._graph_unified_open_attempted = True + meta = self._graph_client_metadata + if meta is None: + # A graph credit arrived but the dataset broadcast was not + # graph-typed: the store location is unknown by contract (no env + # re-derivation). Feeds the existing GraphStoreUnavailable fatal. + self._graph_unified_open_failure = ( + "dataset broadcast did not carry GraphSegmentClientMetadata " + "(or no dataset-configured broadcast was received); the " + "DatasetManager must build the graph store and broadcast its " + "location before workers can serve graph credits" + ) + return None + try: + self._graph_unified_client = GraphSegmentUnifiedClient( + base_path=meta.store_base_path, benchmark_id=meta.benchmark_id + ).open() + except ValueError as e: + # A2-strict rejection: the store EXISTS on disk but is a legacy + # pre-v3 shape. Unlike a missing store this is actionable, so keep + # the reason for the GraphStoreUnavailable fatal. + self._graph_unified_open_failure = str(e) + self.warning(lambda e=e: f"unified store rejected: {e}") + self._graph_unified_client = None + except Exception as e: + self.debug(lambda e=e: f"unified store not opened: {e!r}") + self._graph_unified_client = None + return self._graph_unified_client + + def _graph_store_reader( + self, credit_context: CreditContext + ) -> GraphSegmentUnifiedClient | None: + """Return the worker's lazily-opened unified store client, or ``None``. + + Resolves the store from the graph-typed dataset broadcast + (:class:`GraphSegmentClientMetadata.store_base_path` + + ``benchmark_id``), so it reads exactly what the build-time + :class:`GraphSegmentUnifiedBackingStore` wrote. There is NO env + fallback: a graph credit whose broadcast carried no graph store + location is a recorded failure, not a temp-dir guess. + + When the store is absent or corrupt, rather than letting the failure + propagate to the broad ``_on_credit_drop_message_task`` ``except`` + (which logs but does NOT attribute the error -- F4), this caches a + fatal, actionable + :class:`ErrorDetails`, sets it on the credit context, and returns + ``None`` so the caller skips the send. Attributing the error here is + NOT sufficient on its own: the caller must still emit a synthetic + error record (:meth:`_send_graph_error_record`), or the RecordsManager + completion barrier starves and the run hangs at end of phase. The + doomed open is attempted exactly ONCE; every later credit reuses the + cached error. + """ + unified = self._graph_unified_reader() + if unified is not None: + return unified + if self._graph_store_open_error is None: + meta = self._graph_client_metadata + if meta is None: + # No graph store location arrived on the broadcast; the failure + # string already explains the missing GraphSegmentClientMetadata. + message = f"No graph store could be opened: {self._graph_unified_open_failure}" + else: + unified_failure = ( + f" The store WAS found but rejected: " + f"{self._graph_unified_open_failure}." + if self._graph_unified_open_failure is not None + else "" + ) + message = ( + f"No graph store could be opened under {meta.store_base_path}: " + f"the unified segment store " + f"(aiperf_graph_segments_{meta.benchmark_id}) could not be " + f"opened.{unified_failure} The store location arrives on the " + f"dataset broadcast (GraphSegmentClientMetadata.store_base_path); " + f"ensure the DatasetManager built the graph store and that its " + f"base path is on a shared filesystem visible to every worker." + ) + self._graph_store_open_error = ErrorDetails( + type="GraphStoreUnavailable", + message=message, + ) + self.error(lambda c=self._graph_store_open_error: c.message) + credit_context.error = self._graph_store_open_error + return None + + def _set_graph_envelope_missing( + self, credit_context: CreditContext, base_trace_id: str + ) -> None: + """Set a clean ``GraphEnvelopeMissing`` error for an unaddressable node.""" + credit = credit_context.credit + credit_context.error = ErrorDetails( + type="GraphEnvelopeMissing", + message=( + f"No graph-store envelope for trace={base_trace_id!r} " + f"(instance={credit.trace_id!r}) " + f"node_ordinal={credit.node_ordinal} " + f"phase_variant={credit.phase_variant!r}" + ), + ) + self.warning(lambda c=credit_context.error: c.message) + + async def _send_graph_error_record( + self, credit_context: CreditContext, x_request_id: str + ) -> None: + """Emit a synthetic error ``RequestRecord`` for a pre-dispatch graph failure. + + Pre-dispatch failures (missing store, missing envelope, missing pool + entry) never reach the inference client, so no ``RequestRecord`` would + otherwise flow to the RecordProcessor. The RecordsManager's completion + barrier counts success+error RECORDS against the credit-side + ``final_requests_completed`` (``records_tracker``), so a credit that + returns with an error but no record starves the barrier and the run + hangs at "please wait for the results". Mirror of the session path's + synthetic error record for a failed conversation fetch + (:meth:`_request_conversation_from_dataset_manager`). The + ``CreditReturn`` semantics are untouched -- the caller's finally still + returns the credit with ``credit_context.error`` attached. + """ + credit = credit_context.credit + error = credit_context.error + if not isinstance(error, ErrorDetails): + error = ErrorDetails( + type="GraphPreDispatchError", + message=str(error) if error else "pre-dispatch graph credit failure", + ) + now_perf_ns = time.perf_counter_ns() + await self._send_inference_result_message( + RequestRecord( + request_info=RequestInfo( + model_endpoint=self.model_endpoint, + conversation_id=credit.conversation_id, + turn_index=credit.turn_index, + turns=[], + credit_num=credit.id, + credit_phase=credit.phase, + x_request_id=x_request_id, + x_correlation_id=credit.x_correlation_id, + drop_perf_ns=credit_context.drop_perf_ns, + ), + model_name=self.model_endpoint.primary_model_name, + timestamp_ns=time.time_ns(), + start_perf_ns=now_perf_ns, + end_perf_ns=now_perf_ns, + error=error, + ) + ) + def _parsed_responses_for_record( self, record: RequestRecord ) -> list[ParsedResponse]: @@ -686,6 +1372,11 @@ def _release_and_evict_for_terminal( Non-FORK and non-parent sessions evict immediately. """ + # Fire the routing plugin's post-session hook on ANY terminal outcome + # (final turn or cancel), not just a successful final turn, so stateful + # plugins release sessions abandoned mid-conversation. Idempotent by + # contract; no-op when session routing is unset. + self.inference_client.notify_session_end(x_correlation_id) if ( credit.parent_correlation_id is not None and credit.branch_mode == ConversationBranchMode.FORK @@ -744,6 +1435,9 @@ def _create_request_info( is_final_turn=credit.is_final_turn, agent_depth=credit.agent_depth, parent_correlation_id=credit.parent_correlation_id, + root_correlation_id=credit.effective_root_correlation_id, + is_parent_final=credit.is_parent_final, + is_tree_final=credit.is_tree_final, # Use session's url_index to ensure all turns hit the same backend url_index=session.url_index, ) @@ -879,6 +1573,11 @@ async def _worker_stop(self) -> None: await dataset_client.stop() self.debug("Dataset client stopped") + if self._graph_unified_client is not None: + self._graph_unified_client.close() + self._graph_unified_client = None + self.debug("Graph unified store client closed") + self.event_loop_monitor.stop() diff --git a/src/aiperf/workers/worker_manager.py b/src/aiperf/workers/worker_manager.py index 74e41bca7e..166ac7fbf0 100644 --- a/src/aiperf/workers/worker_manager.py +++ b/src/aiperf/workers/worker_manager.py @@ -73,7 +73,9 @@ def __init__( ), ) self.debug( - lambda: f"Auto-setting max workers to {self.max_workers} due to no max workers specified." + lambda: ( + f"Auto-setting max workers to {self.max_workers} due to no max workers specified." + ) ) # Cap the worker count to the max concurrency, but only if the user is in concurrency mode. diff --git a/src/aiperf/zmq/dealer_request_client.py b/src/aiperf/zmq/dealer_request_client.py index a7c6b7023e..b1298765df 100644 --- a/src/aiperf/zmq/dealer_request_client.py +++ b/src/aiperf/zmq/dealer_request_client.py @@ -156,7 +156,9 @@ async def callback(response_message: Message) -> None: future.set_result(response_message) else: self.debug( - lambda: f"Received response for request {message.request_id} after it was already completed. Ignoring." + lambda: ( + f"Received response for request {message.request_id} after it was already completed. Ignoring." + ) ) await self.request_async(message, callback) diff --git a/src/aiperf/zmq/router_reply_client.py b/src/aiperf/zmq/router_reply_client.py index 790dbbd21e..1e5da1ea9a 100644 --- a/src/aiperf/zmq/router_reply_client.py +++ b/src/aiperf/zmq/router_reply_client.py @@ -98,8 +98,9 @@ def register_request_handler( ) self.debug( - lambda service_id=service_id, - type=message_type: f"Registering request handler for {service_id} with message type {type}" + lambda service_id=service_id, type=message_type: ( + f"Registering request handler for {service_id} with message type {type}" + ) ) self._request_handlers[message_type] = (service_id, handler) @@ -145,7 +146,9 @@ async def _wait_for_response( if response is None: self.warning( - lambda req_id=request_id: f"Got None as response for request {req_id}" + lambda req_id=request_id: ( + f"Got None as response for request {req_id}" + ) ) response = ErrorMessage( request_id=request_id, diff --git a/src/aiperf/zmq/streaming_dealer_client.py b/src/aiperf/zmq/streaming_dealer_client.py index 8d36a1382f..d86c141f7d 100644 --- a/src/aiperf/zmq/streaming_dealer_client.py +++ b/src/aiperf/zmq/streaming_dealer_client.py @@ -50,9 +50,10 @@ class ZMQStreamingDealerClient(BaseZMQClient): Example: ```python - from aiperf.common.structs import ( - Credit, CancelCredits, WorkerReady, WorkerShutdown, CreditReturn + from aiperf.credit.messages import ( + CancelCredits, CreditReturn, WorkerReady, WorkerShutdown ) + from aiperf.credit.structs import Credit # Create via comms (recommended - handles lifecycle management) dealer = comms.create_streaming_dealer_client( @@ -60,11 +61,11 @@ class ZMQStreamingDealerClient(BaseZMQClient): identity="worker-1", ) - async def handle_message(message: Credit | CancelCredits) -> None: + async def handle_message(message: RouterToWorkerMessage) -> None: match message: case Credit() as credit: do_some_work(credit) - await dealer.send(CreditReturn(credit_id=credit.id)) + await dealer.send(CreditReturn(credit=credit)) case CancelCredits(credit_ids=ids): cancel_credits(ids) @@ -117,10 +118,10 @@ def register_receiver(self, handler: RouterToWorkerHandler) -> None: """ Register handler for incoming messages from ROUTER. - The handler will be called for each message received (Credit or CancelCredits). + The handler will be called for each message received (any RouterToWorkerMessage). Args: - handler: Async function that takes a RouterToWorkerMessage (Credit | CancelCredits) + handler: Async function that takes a RouterToWorkerMessage """ if self._receiver_handler is not None: raise ValueError("Receiver handler already registered") @@ -175,7 +176,7 @@ async def _streaming_dealer_receiver(self) -> None: Background task for receiving messages from ROUTER. Runs continuously until stop is requested. Decodes messages as - RouterToWorkerMessage (Credit | CancelCredits) using msgpack. + RouterToWorkerMessage using msgpack. """ self.debug( lambda: f"Streaming DEALER receiver task started for {self.identity}" diff --git a/src/aiperf/zmq/streaming_router_client.py b/src/aiperf/zmq/streaming_router_client.py index b82a31d8b5..f0b29831a5 100644 --- a/src/aiperf/zmq/streaming_router_client.py +++ b/src/aiperf/zmq/streaming_router_client.py @@ -61,8 +61,8 @@ class ZMQStreamingRouterClient(BaseZMQClient): Example: ```python - from aiperf.common.structs import ( - Credit, WorkerReady, WorkerShutdown, CreditReturn + from aiperf.credit.messages import ( + CreditReturn, WorkerReady, WorkerShutdown ) # Create via comms (recommended - handles lifecycle management) @@ -77,8 +77,8 @@ async def handle_message(identity: str, message: WorkerToRouterMessage) -> None: await register_worker(identity) case WorkerShutdown(): await unregister_worker(identity) - case CreditReturn(credit_id=id, cancelled=c, error=e): - await handle_credit_return(identity, id, c, e) + case CreditReturn(credit=credit, cancelled=c, error=e): + await handle_credit_return(identity, credit, c, e) router.register_receiver(handle_message) @@ -130,7 +130,7 @@ def register_receiver(self, handler: WorkerToRouterHandler) -> None: Register handler for incoming messages from DEALER clients. The handler will be called for each message received, with the DEALER's - identity and the decoded message (WorkerReady | WorkerShutdown | CreditReturn). + identity and the decoded message (any WorkerToRouterMessage). Args: handler: Async function that takes (identity: str, message: WorkerToRouterMessage) @@ -197,7 +197,7 @@ async def send_to(self, identity: str, struct: Struct) -> None: Args: identity: The DEALER client's identity (routing key) - struct: The msgspec Struct to send (Credit or CancelCredits) + struct: The msgspec Struct to send (any RouterToWorkerMessage) Raises: NotInitializedError: If socket not initialized @@ -222,7 +222,7 @@ async def _streaming_router_receiver(self) -> None: Background task for receiving messages from DEALER clients. Runs continuously until stop is requested. Decodes messages as - WorkerToRouterMessage (WorkerReady | WorkerShutdown | CreditReturn) using msgpack. + WorkerToRouterMessage using msgpack. """ self.debug("Streaming ROUTER receiver task started") diff --git a/src/aiperf/zmq/sub_client.py b/src/aiperf/zmq/sub_client.py index 0eef5fed5e..388e99d5fc 100644 --- a/src/aiperf/zmq/sub_client.py +++ b/src/aiperf/zmq/sub_client.py @@ -135,7 +135,9 @@ async def _subscribe_internal( ) else: self.debug( - lambda: f"Adding callback to existing subscription for topic: {topic}" + lambda: ( + f"Adding callback to existing subscription for topic: {topic}" + ) ) self._subscribers.setdefault(topic, []).append(callback) @@ -182,7 +184,9 @@ async def _handle_message(self, topic_bytes: bytes, message_bytes: bytes) -> Non message = Message.from_json(message_bytes) self.trace( - lambda: f"Calling callbacks for message: {message}, {self._subscribers.get(topic)}" + lambda: ( + f"Calling callbacks for message: {message}, {self._subscribers.get(topic)}" + ) ) # Call callbacks with the parsed message object diff --git a/src/aiperf/zmq/zmq_base_client.py b/src/aiperf/zmq/zmq_base_client.py index 5b52c409c7..4eb6f7bb8f 100644 --- a/src/aiperf/zmq/zmq_base_client.py +++ b/src/aiperf/zmq/zmq_base_client.py @@ -96,7 +96,9 @@ async def _initialize_socket(self) -> None: self.scheduler = LoopScheduler() self.socket = self.context.socket(self.socket_type) self.debug( - lambda: f"ZMQ {self.socket_type_name} socket initialized, try {'BIND' if self.bind else 'CONNECT'} to {self.address} ({self.client_id})" + lambda: ( + f"ZMQ {self.socket_type_name} socket initialized, try {'BIND' if self.bind else 'CONNECT'} to {self.address} ({self.client_id})" + ) ) if zmq.IDENTITY in self.socket_ops: @@ -104,7 +106,9 @@ async def _initialize_socket(self) -> None: # otherwise it will not be properly set when the socket is bound/connected self.socket.setsockopt(zmq.IDENTITY, self.socket_ops[zmq.IDENTITY]) self.debug( - lambda: f"Set IDENTITY socket option: {self.socket_ops[zmq.IDENTITY]}" + lambda: ( + f"Set IDENTITY socket option: {self.socket_ops[zmq.IDENTITY]}" + ) ) del self.socket_ops[zmq.IDENTITY] @@ -140,7 +144,9 @@ async def _initialize_socket(self) -> None: self.socket.setsockopt(key, val) self.debug( - lambda: f"ZMQ {self.socket_type_name} socket {'BOUND' if self.bind else 'CONNECTED'} to {self.address} ({self.client_id})" + lambda: ( + f"ZMQ {self.socket_type_name} socket {'BOUND' if self.bind else 'CONNECTED'} to {self.address} ({self.client_id})" + ) ) except Exception as e: @@ -171,7 +177,9 @@ async def _shutdown_socket(self) -> None: self.socket.close() except zmq.ContextTerminated: self.debug( - lambda: f"ZMQ context already terminated, skipping socket close ({self.client_id})" + lambda: ( + f"ZMQ context already terminated, skipping socket close ({self.client_id})" + ) ) return except Exception as e: diff --git a/src/aiperf/zmq/zmq_proxy_base.py b/src/aiperf/zmq/zmq_proxy_base.py index ec2fbdb804..525a5845f9 100644 --- a/src/aiperf/zmq/zmq_proxy_base.py +++ b/src/aiperf/zmq/zmq_proxy_base.py @@ -103,7 +103,9 @@ def __init__( self.config = zmq_proxy_config self.debug( - lambda: f"Proxy Initializing - Frontend: {self.config.frontend_address}, Backend: {self.config.backend_address}" + lambda: ( + f"Proxy Initializing - Frontend: {self.config.frontend_address}, Backend: {self.config.backend_address}" + ) ) self.backend_socket = backend_socket_class( @@ -159,10 +161,14 @@ async def _initialize(self) -> None: """Initialize and start the BaseZMQProxy.""" self.debug("Proxy Initializing Sockets...") self.debug( - lambda: f"Frontend {self.frontend_socket.socket_type.name} socket binding to: {self.config.frontend_address} (for {self.backend_socket.socket_type.name} clients)" + lambda: ( + f"Frontend {self.frontend_socket.socket_type.name} socket binding to: {self.config.frontend_address} (for {self.backend_socket.socket_type.name} clients)" + ) ) self.debug( - lambda: f"Backend {self.backend_socket.socket_type.name} socket binding to: {self.config.backend_address} (for {self.frontend_socket.socket_type.name} services)" + lambda: ( + f"Backend {self.backend_socket.socket_type.name} socket binding to: {self.config.backend_address} (for {self.frontend_socket.socket_type.name} services)" + ) ) try: @@ -234,13 +240,17 @@ async def _monitor_messages(self) -> None: return self.debug( - lambda: f"Proxy Monitor Starting - Capture Address: {self.config.capture_address}" + lambda: ( + f"Proxy Monitor Starting - Capture Address: {self.config.capture_address}" + ) ) capture_socket = self.context.socket(SocketType.SUB) capture_socket.connect(self.config.capture_address) self.debug( - lambda: f"Proxy Monitor Connected to Capture Address: {self.config.capture_address}" + lambda: ( + f"Proxy Monitor Connected to Capture Address: {self.config.capture_address}" + ) ) capture_socket.setsockopt(zmq.SUBSCRIBE, b"") # Subscribe to all messages self.debug("Proxy Monitor Subscribed to all messages") diff --git a/src/aiperf/zmq/zmq_proxy_sockets.py b/src/aiperf/zmq/zmq_proxy_sockets.py index 4e36ee0e35..bc91af2e96 100644 --- a/src/aiperf/zmq/zmq_proxy_sockets.py +++ b/src/aiperf/zmq/zmq_proxy_sockets.py @@ -65,7 +65,9 @@ async def _initialize_socket(self) -> None: # explicitly (the ZMQ default) to make intent unambiguous. self.socket.setsockopt(zmq.XPUB_VERBOSE, 0) self.debug( - lambda: "XPUB socket configured with XPUB_VERBOSE=0 (de-duplicated subscription forwarding) to avoid event-bus-proxy CPU/memory blow-up" + lambda: ( + "XPUB socket configured with XPUB_VERBOSE=0 (de-duplicated subscription forwarding) to avoid event-bus-proxy CPU/memory blow-up" + ) ) # Dynamically set the class name and qualname based on the socket and end type diff --git a/tests/component_integration/graph/conftest.py b/tests/component_integration/graph/conftest.py new file mode 100644 index 0000000000..98cfd7323a --- /dev/null +++ b/tests/component_integration/graph/conftest.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared fixtures for graph component-integration tests.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from aiperf.common.environment import Environment + + +@pytest.fixture +def mmap_base_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + """Redirect MMAP_BASE_PATH to tmp_path so stores land in a known dir.""" + monkeypatch.setattr(Environment.DATASET, "MMAP_BASE_PATH", tmp_path) + return tmp_path diff --git a/tests/component_integration/graph/test_bare_run_single_pass.py b/tests/component_integration/graph/test_bare_run_single_pass.py new file mode 100644 index 0000000000..d311fcd3fa --- /dev/null +++ b/tests/component_integration/graph/test_bare_run_single_pass.py @@ -0,0 +1,211 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A BARE graph run does a SINGLE PASS OVER N SESSIONS: each loaded trace once. + +With no explicit stop condition (``--request-count`` / ``--num-conversations`` / +``--benchmark-duration`` all unset) the run stays in SINGLE-PASS mode: +``GraphIRReplayStrategy._recycle_has_stop_condition()`` is ``False`` (it reads only +EXPLICIT stop conditions), so the lane fan-out clamps to the corpus size and each +lane does exactly ONE corpus pass -- clean pass-0 plans, no recycle, no +fresh-start, no auto-10 truncation. + +Separately, ``_resolved_num_sessions()`` derives the reported session TARGET +``N = len(self._parsed.traces)`` (the loaded trace count -- mirroring dag_jsonl's +roots->sessions convention). That target drives lane clamping and progress +reporting ONLY; it is NOT a recycle bound (routing it into the recycle gate would +flip the bare run out of single-pass mode). An explicit ``--num-conversations`` +sets a real stop condition and takes the bounded/recycle path instead. + +Component-level (no worker / ZMQ / mock server): a fake ``CreditIssuer`` echoes +each issued credit back to the strategy's return observer, so ``execute_phase`` +drives the REAL dispatch path over REAL weka traces. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import msgspec +import pytest + +from aiperf.common.enums import CacheBustTarget, CreditPhase + +pytestmark = [pytest.mark.component_integration, pytest.mark.asyncio] + +_FIX_DIR = Path(__file__).parents[2] / "unit" / "graph" / "fixtures" +_MIN = _FIX_DIR / "weka_min.json" + + +@dataclass +class _EchoIssuer: + """Fake CreditIssuer echoing each issued credit to the return observer.""" + + observer: Any = None + issued: int = 0 + returned: int = 0 + sending_complete_calls: int = 0 + sent: list[Any] = field(default_factory=list) + + async def issue_graph_credit(self, turn: Any) -> bool: + self.issued += 1 + self.sent.append(turn) + asyncio.get_running_loop().call_soon(self._echo, turn) + return True + + def _echo(self, turn: Any) -> None: + self.returned += 1 + if self.observer is not None: + self.observer(turn, None, False) + + def mark_graph_sending_complete(self) -> None: + self.sending_complete_calls += 1 + + def graph_all_returned(self) -> bool: + return self.returned >= self.issued + + def set_graph_all_returned_event(self) -> None: ... + + +class _BarePhaseCfg: + """Per-phase config stub with NO stop condition (bare graph run).""" + + def __init__(self, *, concurrency: int | None = None) -> None: + self.phase = CreditPhase.PROFILING + self.concurrency = concurrency + self.expected_num_sessions = None + self.total_expected_requests = None + self.expected_duration_sec = None + self.num_dataset_entries = None + self.max_context_length = None + + +def _corpus_distinct(n: int): + """Return a gap-free ParsedGraph whose ``traces`` holds ``n`` clean-id traces. + + The single ``weka_min`` template is cloned into ``n`` traces with ids + ``t-0``..``t-{n-1}`` (no ``#`` so the instance-id split is unambiguous), and + edge delays are zeroed so a dispatching run replays instantly. + """ + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + + base = from_weka_trace(str(_MIN)) + graph = base.graph + zeroed = [ + msgspec.structs.replace( + e, + **{ + f: 0.0 + for f in ("delay_after_predecessor_us", "min_start_delay_us") + if getattr(e, f, None) is not None + }, + ) + for e in graph.edges + ] + base = msgspec.structs.replace( + base, graph=msgspec.structs.replace(graph, edges=zeroed) + ) + t0 = base.traces[0] + clones = [msgspec.structs.replace(t0, id=f"t-{i}") for i in range(n)] + return msgspec.structs.replace(base, traces=clones) + + +def _make_strategy(parsed, issuer: _EchoIssuer, **overrides): + from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy + + overrides.setdefault("start_min_ratio", 0.0) + overrides.setdefault("start_max_ratio", 0.0) + overrides.setdefault("cache_bust", CacheBustTarget.NONE) + return GraphIRReplayStrategy( + config=overrides.pop("config", None), + credit_issuer=issuer, + lifecycle=overrides.pop("lifecycle", None), + parsed_graph=parsed, + register_observer=lambda obs: setattr(issuer, "observer", obs), + **overrides, + ) + + +def _distinct_templates(issuer: _EchoIssuer) -> set[str]: + """Distinct template trace ids (instance ``t-i#p`` -> ``t-i``) among sent turns.""" + templates: set[str] = set() + for turn in issuer.sent: + trace_id = getattr(turn, "trace_id", None) + if trace_id is not None: + templates.add(str(trace_id).split("::", 1)[0]) + return templates + + +@pytest.mark.parametrize("n", [3, 7]) +async def test_bare_run_dispatches_each_trace_exactly_once(n: int) -> None: + """No explicit stop + concurrency==corpus -> single pass over N sessions.""" + parsed = _corpus_distinct(n) + issuer = _EchoIssuer() + strategy = _make_strategy( + parsed, + issuer, + config=_BarePhaseCfg(concurrency=n), + max_concurrent_traces=n, + ) + + # Reported session target == loaded trace count (N), but the run stays in + # SINGLE-PASS mode: no explicit stop condition, so recycle is OFF (each trace + # once via one corpus pass, clean pass-0 plans, no fresh-start). + assert strategy._resolved_num_sessions() == n + assert strategy._recycle_has_stop_condition() is False + + await strategy.setup_phase() + await asyncio.wait_for(strategy.execute_phase(), timeout=30.0) + + # Single pass over N sessions: N instances, each admitted + completed once. + assert strategy._instances_started == n + assert strategy.admitted_traces == n + assert strategy.completed_traces == n + # And every distinct loaded template ran exactly once (no recycle clones). + assert _distinct_templates(issuer) == {f"t-{i}" for i in range(n)} + + +async def test_bare_run_single_lane_covers_full_corpus() -> None: + """Default concurrency 1 (no explicit stop) covers the whole N-session corpus once.""" + n = 5 + parsed = _corpus_distinct(n) + issuer = _EchoIssuer() + strategy = _make_strategy( + parsed, + issuer, + config=_BarePhaseCfg(concurrency=None), + max_concurrent_traces=None, # resolves to the aiperf default of 1 + ) + assert strategy._max_concurrent == 1 + # Reported session target = loaded trace count, and single-pass mode (no + # explicit stop -> recycle OFF), even at concurrency 1. + assert strategy._resolved_num_sessions() == n + assert strategy._recycle_has_stop_condition() is False + + await strategy.setup_phase() + await asyncio.wait_for(strategy.execute_phase(), timeout=30.0) + + assert strategy._instances_started == n + assert strategy.completed_traces == n + assert _distinct_templates(issuer) == {f"t-{i}" for i in range(n)} + + +async def test_explicit_num_conversations_overrides_derived_sessions() -> None: + """An explicit --num-conversations wins over the derived corpus-size target.""" + n = 6 + parsed = _corpus_distinct(n) + issuer = _EchoIssuer() + cfg = _BarePhaseCfg(concurrency=2) + cfg.expected_num_sessions = 4 # explicit bound below the corpus size + strategy = _make_strategy(parsed, issuer, config=cfg, max_concurrent_traces=2) + + assert strategy._resolved_num_sessions() == 4 + + await strategy.setup_phase() + await asyncio.wait_for(strategy.execute_phase(), timeout=30.0) + + # Explicit cap bounds the run at 4 sessions, not the 6-trace corpus. + assert strategy._instances_started == 4 + assert strategy.completed_traces == 4 diff --git a/tests/component_integration/graph/test_dag_jsonl_byte_parity.py b/tests/component_integration/graph/test_dag_jsonl_byte_parity.py new file mode 100644 index 0000000000..e05a510fd6 --- /dev/null +++ b/tests/component_integration/graph/test_dag_jsonl_byte_parity.py @@ -0,0 +1,587 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""In-process wire-parity golden test: ``--graph-format dag_jsonl`` vs legacy +``--custom-dataset-type dag_jsonl``. + +The core acceptance gate for the dag_jsonl graph adapter: BOTH full production +paths run in one process against the SAME fixture with a deterministic fake +transport, and every wire request body the graph path produces must be +payload-identical to the legacy path's (canonical order-insensitive bytes). + +Harness seams: + +* Both sides run the full ``aiperf profile`` CLI in-process via the ``cli`` + fixture (FakeServiceManager + FakeCommunication). The wire seam is + :class:`tests.harness.fake_transport.FakeTransport` -- ``send_request`` is + monkeypatched to capture ``request_info.payload_bytes`` (the canonical + pre-encoded body ``inference_client`` stamps before transport dispatch on + the format_payload, raw_payload, AND raw_payload_bytes paths) and to return + a deterministic non-streaming chat completion whose text is a pure function + of the request's ``messages`` bytes. Identical request messages therefore + produce identical replies on both harnesses, so live-reply splices (FORK + children, multi-turn accumulators) compose inductively: parity of every + parent request implies parity of every child request. +* Requests are keyed ``(session_id, turn_index)``. Legacy: straight off + ``RequestInfo.conversation_id`` / ``turn_index``. Graph: recovered from the + minted ``x_request_id`` (the worker's ``_mint_x_request_id``: + ``{node_id}::{nonce}`` where the nonce is a ``uuid4().hex`` with no ``::``); + split off the trailing ``::`` from the RIGHT, then split the node id + (``"[#n]:"``) on its trailing ``:``. The graph + ``x_correlation_id`` is now an opaque per-trajectory id carrying no node + identity, so it is NOT parsed. Profiling-phase is checked via + ``RequestInfo.credit_phase`` on both planes. + +Pinned environment facts that make parity hold (asserted where cheap): + +* ``AIPERF_GRAPH_MERGE_CONSECUTIVE_USER`` default False stays in effect. +* No ``--cache-bust`` (default NONE -- no ``[rid:...]`` marker, which also + keeps the graph worker on the pre-serialized bytes path for lineage-free + nodes). +* No warmup on either side (PROFILING-only comparison; every capture is + asserted PROFILING via ``credit_phase`` on both planes). +* Single worker on both sides (legacy FORK seeding and the graph dynamic + pool are both worker-local). +* ``--extra-inputs`` parity: the dag graph adapter folds ``endpoint.extra`` + into ``dispatch_overrides`` at parse (legacy position/precedence: turn + ``extra`` wins on overlap, first insertion keeps position) and stamps + ``endpoint_extra_applied`` so the worker skips its re-merge. The + ``mixed-full-extra-inputs`` case proves it: one key overlapping the + fixture's turn ``extra`` (different value) plus one fresh vendor key. +""" + +from __future__ import annotations + +import hashlib +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import orjson +import pytest +from pytest import param + +from aiperf.common.enums import CreditPhase +from aiperf.common.environment import Environment +from aiperf.common.models import RequestInfo, RequestRecord, TextResponse +from tests.component_integration.conftest import ( + ComponentIntegrationTestDefaults as defaults, +) +from tests.harness.fake_transport import FakeTransport +from tests.harness.utils import AIPerfCLI + +pytestmark = pytest.mark.component_integration + +_FIXTURE_DIR = Path(__file__).resolve().parents[2] / "fixtures" / "dag" / "graph_parity" + +# Markers the tool_calls responder keys on (matched against the LAST user +# message content only, so an inherited ancestor marker never re-triggers on a +# child request). ``[tool]`` is intentionally NOT a substring of ``[tool-only]``. +_TOOL_MARKER = "[tool]" +_TOOL_ONLY_MARKER = "[tool-only]" + + +@dataclass(frozen=True, slots=True) +class _WireRequest: + """One captured transport dispatch: identity fields + canonical body bytes.""" + + x_request_id: str + """The credit's per-dispatch request id (legacy: opaque UUID; graph: minted + ``{node_id}::{nonce}`` -- the per-node identity the graph key recovers).""" + x_correlation_id: str + """The credit's correlation id (legacy: per-session UUID; graph: an opaque + per-trajectory ``{conversation}::{nonce}`` carrying NO node identity).""" + conversation_id: str + """``RequestInfo.conversation_id`` (legacy: the dag session id; graph: the + trajectory template id -- NOT per-node, hence the x_request_id parse).""" + turn_index: int + """``RequestInfo.turn_index`` (legacy: turn within the session; graph: the + node's 0-based turn).""" + credit_phase: CreditPhase + """The credit phase this request was dispatched under.""" + body: bytes + """``request_info.payload_bytes`` -- the canonical wire body.""" + + +def _deterministic_response_text(messages_bytes: bytes) -> str: + """Reply text as a pure function of the request's ``messages`` bytes.""" + return f"resp-{hashlib.sha256(messages_bytes).hexdigest()[:8]}" + + +def _deterministic_tool_calls(messages_bytes: bytes) -> list[dict[str, Any]]: + """One ``tool_calls`` entry as a pure function of the request's messages. + + The id/name/arguments all derive from the SAME hash that seeds the reply + text, so identical requests synthesize an identical tool_call on BOTH + planes and child prompts that splice the reply stay byte-identical. + """ + digest = hashlib.sha256(messages_bytes).hexdigest() + return [ + { + "id": f"call_{digest[:12]}", + "type": "function", + "function": { + "name": f"fn_{digest[:8]}", + "arguments": orjson.dumps({"q": digest[8:16]}).decode("utf-8"), + }, + } + ] + + +def _last_user_content(messages: list[dict[str, Any]]) -> str: + """The content of the LAST ``user`` message (``""`` when none / non-str).""" + for msg in reversed(messages): + if msg.get("role") == "user": + content = msg.get("content") + return content if isinstance(content, str) else "" + return "" + + +def _install_capture_transport( + monkeypatch: pytest.MonkeyPatch, + sink: list[_WireRequest], + *, + tool_calls_mode: bool = False, +) -> None: + """Replace ``FakeTransport.send_request`` with a deterministic responder. + + Captures every dispatched body into ``sink`` and answers with a + non-streaming ``chat.completion`` whose content depends only on the + request's ``messages`` bytes, so both harnesses see identical replies for + identical requests (and therefore build identical child prompts). + + When ``tool_calls_mode`` is set (the ``fork_toolcalls`` fixture only), a + request whose LAST user message carries a marker gets a structured reply: + ``[tool-only]`` -> ``content: null`` + ``tool_calls`` (tool-only); + ``[tool]`` -> the usual ``resp-`` content PLUS ``tool_calls`` + (mixed). Determinism is unchanged (pure function of the request bytes), and + the responder is identical on both planes, so parity composes inductively. + """ + + async def _send_request( + self: FakeTransport, + request_info: RequestInfo, + payload: Any, + *, + first_token_callback: Any = None, + ) -> RequestRecord: + body = request_info.payload_bytes + assert body is not None, ( + "inference_client must stamp payload_bytes before transport dispatch" + ) + sink.append( + _WireRequest( + x_request_id=request_info.x_request_id, + x_correlation_id=request_info.x_correlation_id, + conversation_id=request_info.conversation_id, + turn_index=request_info.turn_index, + credit_phase=request_info.credit_phase, + body=bytes(body), + ) + ) + parsed = orjson.loads(body) + messages_bytes = orjson.dumps(parsed["messages"]) + text = _deterministic_response_text(messages_bytes) + message: dict[str, Any] = {"role": "assistant", "content": text} + finish_reason = "stop" + if tool_calls_mode: + marker = _last_user_content(parsed["messages"]) + if _TOOL_ONLY_MARKER in marker: + message = { + "role": "assistant", + "content": None, + "tool_calls": _deterministic_tool_calls(messages_bytes), + } + finish_reason = "tool_calls" + elif _TOOL_MARKER in marker: + message = { + "role": "assistant", + "content": text, + "tool_calls": _deterministic_tool_calls(messages_bytes), + } + finish_reason = "tool_calls" + response_data = { + "id": "chatcmpl-parity", + "object": "chat.completion", + "created": 0, + "model": parsed.get("model", ""), + "choices": [ + { + "index": 0, + "message": message, + "finish_reason": finish_reason, + } + ], + "usage": {"prompt_tokens": 8, "completion_tokens": 4, "total_tokens": 12}, + } + start_perf_ns = time.perf_counter_ns() + end_perf_ns = time.perf_counter_ns() + return RequestRecord( + start_perf_ns=start_perf_ns, + end_perf_ns=end_perf_ns, + timestamp_ns=time.time_ns(), + status=200, + responses=[ + TextResponse( + perf_ns=end_perf_ns, + content_type="application/json", + text=orjson.dumps(response_data).decode("utf-8"), + ) + ], + ) + + monkeypatch.setattr(FakeTransport, "send_request", _send_request) + + +def _reset_inprocess_state() -> None: + """Isolate consecutive in-process CLI runs within ONE test. + + The package conftest clears singletons / reseeds the RNG only BETWEEN + tests; this test runs ``aiperf profile`` twice in one test body, so the + second run must not see the first run's singleton comms/services or RNG + stream. + """ + from aiperf.common import random_generator as rng + from aiperf.common.singleton import SingletonMeta + + SingletonMeta._instances.clear() + rng.reset() + rng.init(42) + + +def _run_and_capture( + cli: AIPerfCLI, sink: list[_WireRequest], cmd: str +) -> list[_WireRequest]: + """Run one in-process profile and return the requests it dispatched.""" + _reset_inprocess_state() + sink.clear() + cli.run_sync(cmd, timeout=120.0, assert_success=True) + return list(sink) + + +def _keyed_legacy(captures: list[_WireRequest]) -> dict[tuple[str, int], list[bytes]]: + """Key legacy captures by ``(session_id, turn_index)`` into a multiset. + + A SPAWN template instantiated N times keys N bodies under the same + ``(session, turn)`` (legacy carries the TEMPLATE ``conversation_id``, so + repeated spawns share the key), so the value is a LIST -- one body per + dispatch under that key. + """ + keyed: dict[tuple[str, int], list[bytes]] = {} + for c in captures: + assert c.credit_phase == CreditPhase.PROFILING, ( + f"legacy request outside PROFILING: {c.credit_phase} " + f"(warmup must stay disabled for parity)" + ) + keyed.setdefault((c.conversation_id, c.turn_index), []).append(c.body) + return keyed + + +def _keyed_graph(captures: list[_WireRequest]) -> dict[tuple[str, int], list[bytes]]: + """Key graph captures by ``(session_id, turn_index)`` from x_request_id. + + ``_mint_x_request_id`` folds the per-dispatch node identity into + ``x_request_id`` as ``{node_id}::{nonce}``; the nonce is a ``uuid4().hex`` + with no ``::``, so split it off from the RIGHT. Node ids are + ``"[#n]:"`` (``#n`` only for repeated SPAWN instances; + ``#`` is gated out of session ids at load time). Repeated SPAWN instances + strip their ``#n`` and accumulate under one key as a multiset -- one body + per instance. Profiling-phase is enforced via ``credit_phase``. + """ + keyed: dict[tuple[str, int], list[bytes]] = {} + for c in captures: + assert c.credit_phase == CreditPhase.PROFILING, ( + f"graph request outside PROFILING: {c.credit_phase} " + f"(warmup must stay disabled for parity)" + ) + assert "::" in c.x_request_id, ( + f"unparsable graph x_request_id {c.x_request_id!r} (no '::' nonce delimiter)" + ) + node_id = c.x_request_id.rsplit("::", 1)[0] + session_part, sep, turn_str = node_id.rpartition(":") + assert sep and turn_str.isdigit(), ( + f"graph node id {node_id!r} does not end in ':'" + ) + key = (session_part.split("#", 1)[0], int(turn_str)) + keyed.setdefault(key, []).append(c.body) + return keyed + + +def _fmt_body(body: bytes) -> str: + return body.decode("utf-8", errors="replace") + + +def _canonical_bytes(body: bytes) -> bytes: + """Canonical order-insensitive form: re-serialize with sorted keys.""" + return orjson.dumps(orjson.loads(body), option=orjson.OPT_SORT_KEYS) + + +def _has_live_reply_message(body: bytes) -> bool: + """True if ``body``'s ``messages`` embed a captured live assistant reply. + + The fake transport answers every request with an assistant message whose + content is ``resp-`` (:func:`_deterministic_response_text`), so a + spliced-in reply is any ``role == "assistant"`` message whose content + starts with ``"resp-"``. + """ + payload = orjson.loads(body) + return any( + msg.get("role") == "assistant" + and str(msg.get("content", "")).startswith("resp-") + for msg in payload.get("messages", []) + ) + + +def _has_tool_calls_message(body: bytes) -> bool: + """True if ``body`` embeds an assistant message carrying ``tool_calls``. + + Anti-vacuous guard for the tool_calls parity case: a spliced-in reply that + reassembled to a ``tool_calls``-bearing assistant message must actually + reach the child's wire body, or a symmetric drop would pass the multiset + comparison while gutting all tool_calls splice coverage. + """ + payload = orjson.loads(body) + return any( + msg.get("role") == "assistant" and "tool_calls" in msg + for msg in payload.get("messages", []) + ) + + +def _has_content_null_tool_calls_message(body: bytes) -> bool: + """True if ``body`` embeds a ``content: null`` tool_calls-only assistant.""" + payload = orjson.loads(body) + return any( + msg.get("role") == "assistant" + and "tool_calls" in msg + and msg.get("content") is None + for msg in payload.get("messages", []) + ) + + +@pytest.mark.parametrize( + "fixture_name,expected_requests,tool_calls_mode,extra_inputs", + [ + param("fork_minimal.dag.jsonl", 5, False, None, id="fork-minimal"), + param("spawn_join.dag.jsonl", 5, False, None, id="spawn-join"), + param("mixed_full.dag.jsonl", 6, False, None, id="mixed-full"), + param("spawn_repeat.dag.jsonl", 4, False, None, id="spawn-repeat"), + param("fork_toolcalls.dag.jsonl", 3, True, None, id="fork-toolcalls"), + # --extra-inputs precedence parity: ``temperature`` OVERLAPS the + # fixture's turn ``extra`` (root:0 authors 0.5, finisher:0 authors 0.1; + # the run supplies 0.9 -- turn extra must win on BOTH planes) and + # ``min_p`` is a fresh vendor key (must land at the endpoint-extra + # position, before the turn-extra-only keys, on BOTH planes). + param( + "mixed_full.dag.jsonl", + 6, + False, + "temperature:0.9 min_p:0.05", + id="mixed-full-extra-inputs", + ), + ], +) # fmt: skip +def test_dag_jsonl_graph_adapter_wire_bytes_match_legacy( + cli: AIPerfCLI, + mmap_base_path: Path, + monkeypatch: pytest.MonkeyPatch, + fixture_name: str, + expected_requests: int, + tool_calls_mode: bool, + extra_inputs: str | None, +) -> None: + """Every PROFILING wire body from the graph path byte-matches legacy. + + Three criteria per request, keyed ``(session_id, turn_index)``: + + 1. ``payload["messages"]`` byte-equal (roots, FORK children with live + parent replies spliced in, SPAWN children, join turns). + 2. Full payload parsed-equal. + 3. Canonical (sorted-keys) body byte-equal -- order-insensitive on key + position, byte-strict on values and types. + + Plus: request COUNTS match and no request matched on only one side. + """ + # Pin the environment facts parity depends on. + assert Environment.GRAPH.MERGE_CONSECUTIVE_USER is False, ( + "AIPERF_GRAPH_MERGE_CONSECUTIVE_USER must stay at its False default" + ) + + fixture = _FIXTURE_DIR / fixture_name + assert fixture.is_file(), f"missing fixture {fixture}" + + sink: list[_WireRequest] = [] + _install_capture_transport(monkeypatch, sink, tool_calls_mode=tool_calls_mode) + + # Shared per-case flags MUST be appended identically to both planes. + extra_args = f"--extra-inputs {extra_inputs} " if extra_inputs else "" + + # The legacy plane needs --concurrency 2: fork fanout dispatches children as + # separate sessions, so a single slot would serialize a parent and its child + # and stall the fork seeding. The graph plane omits it because its replay lanes + # self-schedule fork/spawn children off the parent's completion. + legacy_captures = _run_and_capture( + cli, + sink, + f""" + aiperf profile \ + --model {defaults.model} \ + --custom-dataset-type dag_jsonl \ + --input-file {fixture} \ + --concurrency 2 \ + --num-conversations 1 \ + --record-processor-service-count 1 \ + --workers-max 1 \ + {extra_args}--ui {defaults.ui} + """, + ) + graph_captures = _run_and_capture( + cli, + sink, + f""" + aiperf profile \ + --model {defaults.model} \ + --graph-format dag_jsonl \ + --input-file {fixture} \ + --num-conversations 1 \ + --record-processor-service-count 1 \ + --workers-max 1 \ + {extra_args}--ui {defaults.ui} + """, + ) + + assert len(legacy_captures) == expected_requests, ( + f"legacy dispatched {len(legacy_captures)} requests, " + f"expected {expected_requests}" + ) + assert len(graph_captures) == expected_requests, ( + f"graph dispatched {len(graph_captures)} requests, expected {expected_requests}" + ) + + legacy_by_key = _keyed_legacy(legacy_captures) + graph_by_key = _keyed_graph(graph_captures) + + only_legacy = sorted(set(legacy_by_key) - set(graph_by_key)) + only_graph = sorted(set(graph_by_key) - set(legacy_by_key)) + assert not only_legacy and not only_graph, ( + f"request sets diverge: only-legacy={only_legacy} only-graph={only_graph}" + ) + + for key in sorted(legacy_by_key): + # Repeated SPAWN templates key multiple bodies under one (session, turn); + # compare as a multiset -- sort each side and pair element-wise. Per-key + # multiplicity must also match (a template fired N times on one plane and + # M on the other is a divergence). + legacy_bodies = sorted(legacy_by_key[key], key=_canonical_bytes) + graph_bodies = sorted(graph_by_key[key], key=_canonical_bytes) + assert len(legacy_bodies) == len(graph_bodies), ( + f"{key}: request multiplicity diverges " + f"legacy={len(legacy_bodies)} graph={len(graph_bodies)}" + ) + + for idx, (legacy_body, graph_body) in enumerate( + zip(legacy_bodies, graph_bodies, strict=True) + ): + where = f"{key} #{idx}" if len(legacy_bodies) > 1 else f"{key}" + legacy_payload = orjson.loads(legacy_body) + graph_payload = orjson.loads(graph_body) + + # Criterion 1: messages byte-equal for EVERY request. + assert orjson.dumps(legacy_payload["messages"]) == orjson.dumps( + graph_payload["messages"] + ), ( + f"{where}: messages diverge\n" + f" legacy: {legacy_payload['messages']}\n" + f" graph: {graph_payload['messages']}" + ) + + # Criterion 2: full payload parsed-equal. + assert legacy_payload == graph_payload, ( + f"{where}: parsed payloads diverge\n" + f" legacy: {legacy_payload}\n" + f" graph: {graph_payload}" + ) + + # Criterion 3: canonical (sorted-keys) body byte-equal. Key order is + # deliberately NOT compared -- the graph plane authors stream / the + # token cap / model through native node fields, so their wire + # positions differ from legacy; canonical bytes stay strict on + # values AND types (1 vs 1.0, True vs 1) that criterion 2 conflates. + assert _canonical_bytes(legacy_body) == _canonical_bytes(graph_body), ( + f"{where}: parsed-equal but CANONICALLY byte-different " + f"(value/type drift)\n" + f" legacy: {_fmt_body(legacy_body)}\n" + f" graph: {_fmt_body(graph_body)}" + ) + + # Self-verification: the gate must NOT pass by both paths symmetrically + # dropping the spliced-in live replies (both omit the captured turn when + # build_assistant_turn returns None). At least one compared graph body must + # actually embed a captured assistant reply (content ``resp-*`` from the + # fake transport) -- otherwise a fake-response-shape or parse regression + # would silently gut all splice-position coverage while parity still holds. + assert any( + _has_live_reply_message(body) + for bodies in graph_by_key.values() + for body in bodies + ), ( + "no graph request embedded a live assistant reply (resp-*); reply " + "capture/parse regressed and splice-position coverage was silently lost" + ) + + # Fixture-specific: fork_minimal's branch-a turn-0 is a FORK child whose + # prompt must embed the parent root turn's spliced live reply. + if fixture_name == "fork_minimal.dag.jsonl": + branch_a_bodies = graph_by_key.get(("branch-a", 0)) + assert branch_a_bodies, ( + "fork_minimal: expected a graph request keyed ('branch-a', 0)" + ) + assert _has_live_reply_message(branch_a_bodies[0]), ( + "fork_minimal ('branch-a', 0): FORK-child prompt is missing the " + "spliced parent reply (resp-*)" + ) + + # Fixture-specific: fork_toolcalls proves tool_calls-bearing replies round + # -trip byte-identically. The root reply is mixed (content + tool_calls) and + # must splice into the fork child's turn-0 prompt; the child's own turn-0 + # reply is content:null tool_calls-only and must splice into turn-1's prompt. + if fixture_name == "fork_toolcalls.dag.jsonl": + child_t0_bodies = graph_by_key.get(("branch-a", 0)) + assert child_t0_bodies, ( + "fork_toolcalls: expected a graph request keyed ('branch-a', 0)" + ) + assert _has_tool_calls_message(child_t0_bodies[0]), ( + "fork_toolcalls ('branch-a', 0): FORK-child prompt is missing the " + "spliced parent assistant message carrying tool_calls" + ) + child_t1_bodies = graph_by_key.get(("branch-a", 1)) + assert child_t1_bodies, ( + "fork_toolcalls: expected a graph request keyed ('branch-a', 1)" + ) + assert _has_content_null_tool_calls_message(child_t1_bodies[0]), ( + "fork_toolcalls ('branch-a', 1): prompt is missing the spliced " + "content:null tool_calls-only assistant message" + ) + + # Fixture-specific: the extra-inputs case must not pass vacuously (both + # planes symmetrically dropping the run extras would keep parity). The + # overlap key keeps the TURN value (legacy precedence) at the endpoint-extra + # position, and the fresh vendor key lands on every request. + if extra_inputs is not None: + root_body = orjson.loads(graph_by_key[("root", 0)][0]) + assert root_body["temperature"] == 0.5, ( + "overlap key: the turn 'extra' value must beat --extra-inputs" + ) + assert root_body["min_p"] == 0.05 + keys = list(root_body) + assert keys.index("temperature") < keys.index("min_p") < keys.index("top_p"), ( + f"endpoint-extra position drifted: {keys}" + ) + finisher_body = orjson.loads(graph_by_key[("finisher", 0)][0]) + assert finisher_body["temperature"] == 0.1, ( + "overlap key: the turn 'extra' value must beat --extra-inputs" + ) + helper_body = orjson.loads(graph_by_key[("helper", 0)][0]) + assert helper_body["temperature"] == 0.9, ( + "no turn overlap: the run-level --extra-inputs value must land" + ) + assert helper_body["min_p"] == 0.05 diff --git a/tests/component_integration/graph/test_dag_jsonl_e2e.py b/tests/component_integration/graph/test_dag_jsonl_e2e.py new file mode 100644 index 0000000000..c77b4de116 --- /dev/null +++ b/tests/component_integration/graph/test_dag_jsonl_e2e.py @@ -0,0 +1,559 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""In-process RUNTIME end-to-end proof for ``--graph-format dag_jsonl``. + +Where :mod:`test_dag_jsonl_byte_parity` proves the graph path builds the SAME +wire bytes as legacy, this module proves the graph path RUNS with the right +firing semantics: it drives the full ``aiperf profile`` CLI in-process (the +``GraphIRReplayStrategy`` -> per-trace ``TraceExecutor`` -> ``CreditDispatchAdapter`` +-> worker materialize path) against a deterministic fake transport and asserts +on the observed dispatch/completion ordering and per-turn delays. + +Harness seams (shared with the byte-parity test's pattern): + +* The wire seam is :meth:`tests.harness.fake_transport.FakeTransport.send_request`, + monkeypatched to (a) capture ``request_info.payload_bytes`` (the canonical + pre-encoded body) plus a monotonic dispatch/completion TICK and a + ``perf_counter_ns`` at entry/exit, and (b) answer with a deterministic + non-streaming chat completion whose text is a pure function of the request's + ``messages`` bytes -- so FORK-child live-reply splices are reproducible and + assertable. +* Node identity is recovered from the minted ``x_request_id`` + (``{x_base}|{runtime_trace_id}|{node_id}|{phase_variant}``), anchored from the + RIGHT because ``x_base`` itself may contain ``|``. The node id is + ``"[#n]:"``. + +Causality is asserted with COMPLETION-ORDER ticks (a global monotonic counter +bumped on ``send_request`` entry AND exit), never wall-clock deltas -- so the +ordering assertions are xdist-safe. The single delay assertion is a LOWER BOUND +only (authored 250 ms), measured at the transport seam. + +Asymmetric latency: the responder slows a chosen SESSION's replies by +``_SLOW_RESPONSE_S`` via ``asyncio.sleep`` (deterministic TEXT unchanged) so that +a fast sibling turn can be observed dispatching WHILE a slow turn is still +in-flight (the "intermediate turn does not wait for its spawned child" and +"pre-session child does not wait for its owner" proofs). + +Pinned run flags (mirroring the byte-parity graph run): ``--num-conversations 1`` +(single pass, no recycle -> exactly one dispatch per node), single worker (the +dynamic splice pool is worker-local), single record processor, no warmup. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import itertools +import time +from collections import Counter +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import orjson +import pytest +from pytest import param + +from aiperf.common.constants import IS_WINDOWS +from aiperf.common.enums import CreditPhase +from aiperf.common.models import RequestInfo, RequestRecord, TextResponse +from tests.component_integration.conftest import ( + ComponentIntegrationTestDefaults as defaults, +) +from tests.harness.fake_transport import FakeTransport +from tests.harness.utils import AIPerfCLI + +pytestmark = pytest.mark.component_integration + +_FIXTURE_DIR = Path(__file__).resolve().parents[2] / "fixtures" / "dag" / "graph_parity" + + +# Artificial per-session reply latency for the asymmetric-latency proofs. Large +# enough that a fast sibling turn reliably lands its whole dispatch->completion +# window inside a slow turn's in-flight window, small enough to keep the suite +# quick. +_SLOW_RESPONSE_S = 0.3 + +# Authored ``delay`` on the gated turn of ``delayed_turn.dag.jsonl`` (ms) and the +# lower bound we assert at the transport seam. asyncio.sleep never returns early +# and the executor stamps the predecessor finish AFTER send_request returns, so +# the real gap is provably >= the authored delay; a small slack below absorbs +# only perf_counter sampling granularity. +_DELAY_TURN_MS = 250 +# 5 ms scheduler slack on POSIX; Windows timers tick at ~15.6 ms granularity +# and the loop's timer clock can fire up to two ticks early relative to the +# perf_counter timestamps the seam records, so allow that much there. +_DELAY_LOWER_BOUND_S = 0.218 if IS_WINDOWS else 0.245 + + +def _deterministic_response_text(messages_bytes: bytes) -> str: + """Reply text as a pure function of the request's ``messages`` bytes. + + Byte-identical to the byte-parity harness's responder so a spliced FORK-child + reply is reproducible: ``resp-``. + """ + return f"resp-{hashlib.sha256(messages_bytes).hexdigest()[:8]}" + + +def _node_key(x_request_id: str) -> tuple[str, int]: + """Recover ``(session_id, turn_index)`` from a minted graph request id. + + The worker mints ``{node_id}::{nonce}`` (``_mint_x_request_id``); the node + id is ``"[#n]:"`` (``#n`` only for repeated SPAWN + instances). The correlation id is the trajectory-instance id + (``{conversation}::{nonce}``) and carries no node identity. + """ + node_id, sep, _nonce = x_request_id.rpartition("::") + assert sep, f"unparsable graph x_request_id {x_request_id!r}" + session_part, sep, turn_str = node_id.rpartition(":") + assert sep and turn_str.isdigit(), ( + f"graph node id {node_id!r} does not end in ':'" + ) + return session_part.split("#", 1)[0], int(turn_str) + + +@dataclass(frozen=True, slots=True) +class _Fire: + """One captured transport dispatch with ordering + timing instrumentation.""" + + session: str + """Recovered template session id (``#n`` SPAWN suffix stripped).""" + turn: int + """Recovered zero-based turn index within the session.""" + x_correlation_id: str + """The minted sticky correlation id.""" + credit_phase: CreditPhase + """Credit phase this dispatch ran under (asserted PROFILING everywhere).""" + agent_depth: int + """DAG identity stamped on the credit and carried by ``request_info`` into + the record pipeline (0 = root chain; children owner depth + 1).""" + parent_correlation_id: str | None + """The spawning parent node's correlation id for fork/spawn children; None + for roots and pre-session children. Rides ``request_info`` into records.""" + body: bytes + """``request_info.payload_bytes`` -- the canonical wire body.""" + dispatch_tick: int + """Global monotonic counter value at ``send_request`` ENTRY (fire order).""" + complete_tick: int + """Global monotonic counter value at ``send_request`` EXIT (completion order).""" + dispatch_perf_ns: int + """``perf_counter_ns`` at ``send_request`` entry (delay lower-bound only).""" + complete_perf_ns: int + """``perf_counter_ns`` at ``send_request`` exit (delay lower-bound only).""" + + @property + def key(self) -> tuple[str, int]: + return (self.session, self.turn) + + def messages(self) -> list[dict[str, Any]]: + return orjson.loads(self.body)["messages"] + + +def _install_timing_transport( + monkeypatch: pytest.MonkeyPatch, + sink: list[_Fire], + *, + slow_sessions: dict[str, float], +) -> None: + """Replace ``FakeTransport.send_request`` with an instrumented responder. + + Records a :class:`_Fire` per dispatch (identity + body + entry/exit tick + + entry/exit perf clock) and answers with a deterministic chat completion. + Replies for a session listed in ``slow_sessions`` are delayed by that many + seconds (``asyncio.sleep``) BEFORE the record is finalized, so a concurrent + fast sibling dispatch interleaves ahead of the slow completion tick. The + reply TEXT is unaffected by the delay (pure function of ``messages`` bytes). + """ + tick = itertools.count() + + async def _send_request( + self: FakeTransport, + request_info: RequestInfo, + payload: Any, + *, + first_token_callback: Any = None, + ) -> RequestRecord: + body = request_info.payload_bytes + assert body is not None, ( + "inference_client must stamp payload_bytes before transport dispatch" + ) + session, turn = _node_key(request_info.x_request_id) + dispatch_tick = next(tick) + dispatch_perf_ns = time.perf_counter_ns() + + delay_s = slow_sessions.get(session, 0.0) + if delay_s: + await asyncio.sleep(delay_s) + + parsed = orjson.loads(body) + text = _deterministic_response_text(orjson.dumps(parsed["messages"])) + response_data = { + "id": "chatcmpl-e2e", + "object": "chat.completion", + "created": 0, + "model": parsed.get("model", ""), + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": text}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 8, "completion_tokens": 4, "total_tokens": 12}, + } + complete_perf_ns = time.perf_counter_ns() + complete_tick = next(tick) + sink.append( + _Fire( + session=session, + turn=turn, + x_correlation_id=request_info.x_correlation_id, + credit_phase=request_info.credit_phase, + agent_depth=request_info.agent_depth, + parent_correlation_id=request_info.parent_correlation_id, + body=bytes(body), + dispatch_tick=dispatch_tick, + complete_tick=complete_tick, + dispatch_perf_ns=dispatch_perf_ns, + complete_perf_ns=complete_perf_ns, + ) + ) + return RequestRecord( + start_perf_ns=dispatch_perf_ns, + end_perf_ns=complete_perf_ns, + timestamp_ns=time.time_ns(), + status=200, + responses=[ + TextResponse( + perf_ns=complete_perf_ns, + content_type="application/json", + text=orjson.dumps(response_data).decode("utf-8"), + ) + ], + ) + + monkeypatch.setattr(FakeTransport, "send_request", _send_request) + + +def _run_graph(cli: AIPerfCLI, sink: list[_Fire], fixture: Path) -> list[_Fire]: + """Run one in-process ``--graph-format dag_jsonl`` profile; return its fires. + + Single pass (``--num-conversations 1``), single worker (the dynamic splice + pool is worker-local), single record processor, no warmup -- the same run + shape the byte-parity gate pins. + """ + assert fixture.is_file(), f"missing fixture {fixture}" + sink.clear() + cli.run_sync( + f""" + aiperf profile \ + --model {defaults.model} \ + --graph-format dag_jsonl \ + --input-file {fixture} \ + --num-conversations 1 \ + --record-processor-service-count 1 \ + --workers-max 1 \ + --ui {defaults.ui} + """, + timeout=120.0, + assert_success=True, + ) + return list(sink) + + +def _by_key(fires: list[_Fire]) -> dict[tuple[str, int], _Fire]: + """Index fires by ``(session, turn)``, asserting no node fired twice.""" + keyed: dict[tuple[str, int], _Fire] = {} + for f in fires: + assert f.key not in keyed, f"node {f.key} dispatched more than once" + keyed[f.key] = f + return keyed + + +@pytest.mark.parametrize( + "fixture_name,expected_keys", + [ + param( + "fork_minimal.dag.jsonl", + { + ("root", 0), + ("branch-a", 0), + ("branch-a", 1), + ("branch-b", 0), + ("branch-b", 1), + }, + id="fork-minimal", + ), + param( + "spawn_join.dag.jsonl", + {("parent", 0), ("parent", 1), ("parent", 2), ("worker-a", 0), ("worker-a", 1)}, + id="spawn-join", + ), + param( + "mixed_full.dag.jsonl", + { + ("root", 0), + ("root", 1), + ("root", 2), + ("helper", 0), + ("helper", 1), + ("finisher", 0), + }, + id="mixed-full", + ), + param( + "prespawn.dag.jsonl", + {("root", 0), ("root", 1), ("bg-child", 0)}, + id="prespawn", + ), + param( + "delayed_turn.dag.jsonl", + {("seq", 0), ("seq", 1)}, + id="delayed-turn", + ), + ], +) # fmt: skip +def test_every_node_dispatches_exactly_once( + cli: AIPerfCLI, + mmap_base_path: Path, + monkeypatch: pytest.MonkeyPatch, + fixture_name: str, + expected_keys: set[tuple[str, int]], +) -> None: + """(a) Every tree node dispatches EXACTLY once per single pass. + + No recycles (``--num-conversations 1`` is single-pass), so the captured + ``(session, turn)`` multiset must equal the tree's node set with count 1 + each -- across fork children, spawn children, join turns, pre-session + children, and delayed turns. + """ + sink: list[_Fire] = [] + _install_timing_transport(monkeypatch, sink, slow_sessions={}) + fires = _run_graph(cli, sink, _FIXTURE_DIR / fixture_name) + + assert all(f.credit_phase == CreditPhase.PROFILING for f in fires), ( + "a dispatch ran outside PROFILING (warmup must stay disabled)" + ) + counts = Counter(f.key for f in fires) + assert set(counts) == expected_keys, ( + f"node set mismatch: only-observed={set(counts) - expected_keys} " + f"only-expected={expected_keys - set(counts)}" + ) + assert all(c == 1 for c in counts.values()), ( + f"nodes dispatched more than once: {[k for k, c in counts.items() if c != 1]}" + ) + + +def test_fork_children_fire_after_parent_and_embed_live_reply( + cli: AIPerfCLI, + mmap_base_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """(b) FORK children fire only AFTER the parent's response completes, and + (c) a FORK child's materialized messages embed the parent's ACTUAL + deterministic reply at the correct position (exact message-list shape).""" + sink: list[_Fire] = [] + _install_timing_transport(monkeypatch, sink, slow_sessions={}) + fires = _run_graph(cli, sink, _FIXTURE_DIR / "fork_minimal.dag.jsonl") + by_key = _by_key(fires) + + root = by_key[("root", 0)] + branch_a = by_key[("branch-a", 0)] + branch_b = by_key[("branch-b", 0)] + + # (b) Causality: each fork child's dispatch strictly follows the parent's + # completion (its firing gate is a ChannelRequirement on ``root:0_out``, + # written only after the parent request returns). + assert branch_a.dispatch_tick > root.complete_tick, ( + f"branch-a fired before parent completed: " + f"a.dispatch={branch_a.dispatch_tick} root.complete={root.complete_tick}" + ) + assert branch_b.dispatch_tick > root.complete_tick, ( + f"branch-b fired before parent completed: " + f"b.dispatch={branch_b.dispatch_tick} root.complete={root.complete_tick}" + ) + + # (c) Exact message-list shape: the fork child body is the parent's authored + # messages, then the parent's ACTUAL captured reply as an assistant turn, + # then the child's own authored messages -- in that order. + parent_messages = root.messages() + parent_reply = _deterministic_response_text(orjson.dumps(parent_messages)) + child_authored = [ + {"role": "user", "content": "Expand on the first section in more detail."}, + {"role": "user", "content": "Add a brief counter-argument."}, + ] + expected = [ + *parent_messages, + {"role": "assistant", "content": parent_reply}, + *child_authored, + ] + assert branch_a.messages() == expected, ( + "fork child ('branch-a', 0) message shape/positioning wrong\n" + f" expected: {expected}\n" + f" actual: {branch_a.messages()}" + ) + # The embedded reply is a genuine live capture, not a static echo. + assert parent_reply.startswith("resp-"), parent_reply + + +def test_spawn_join_gates_on_children_but_intermediate_does_not_wait( + cli: AIPerfCLI, + mmap_base_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """(b) SPAWN-join gating and busy-parent independence in one run. + + ``worker-a`` replies are slowed so the parent's intermediate turn (turn 1, + between the spawn turn 0 and the ``join_at: 2`` turn) can be observed + dispatching WHILE the spawned child is still in-flight, while the join turn + (turn 2) is held until the child's leaf response completes. + """ + sink: list[_Fire] = [] + _install_timing_transport( + monkeypatch, sink, slow_sessions={"worker-a": _SLOW_RESPONSE_S} + ) + fires = _run_graph(cli, sink, _FIXTURE_DIR / "spawn_join.dag.jsonl") + by_key = _by_key(fires) + + p0 = by_key[("parent", 0)] + p1 = by_key[("parent", 1)] + p2 = by_key[("parent", 2)] + w0 = by_key[("worker-a", 0)] + w1 = by_key[("worker-a", 1)] + + # Intermediate turn gates only on its own predecessor (parent turn 0), so it + # dispatches after parent turn 0 completes... + assert p1.dispatch_tick > p0.complete_tick, ( + f"intermediate fired before parent turn 0 completed: " + f"p1.dispatch={p1.dispatch_tick} p0.complete={p0.complete_tick}" + ) + # ...WITHOUT waiting for the (deliberately slowed) spawned child: the + # intermediate dispatch lands before the slow child's turn-0 completion. + assert p1.dispatch_tick < w0.complete_tick, ( + f"intermediate waited on the spawned child: " + f"p1.dispatch={p1.dispatch_tick} w0.complete={w0.complete_tick}" + ) + + # Join turn is gated on ALL gating children's LEAF responses (worker-a:1) as + # well as its own predecessor (parent turn 1). + assert p2.dispatch_tick > w1.complete_tick, ( + f"join turn fired before the child leaf completed: " + f"p2.dispatch={p2.dispatch_tick} w1.complete={w1.complete_tick}" + ) + assert p2.dispatch_tick > w0.complete_tick, ( + f"join turn fired before the child's first turn completed: " + f"p2.dispatch={p2.dispatch_tick} w0.complete={w0.complete_tick}" + ) + assert p2.dispatch_tick > p1.complete_tick, ( + f"join turn fired before its own predecessor completed: " + f"p2.dispatch={p2.dispatch_tick} p1.complete={p1.complete_tick}" + ) + + +def test_pre_session_spawn_child_fires_without_waiting_for_owner( + cli: AIPerfCLI, + mmap_base_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """(b) A ``pre_session_spawns`` child fires on the owner turn-0 trigger, NOT + after the owner's turn-0 response. + + First runtime exercise of pre-spawn lowering. The owner (``root``) replies + are slowed; the pre-session child (``bg-child``) shares the owner turn-0's + (empty -> START) firing predecessors and carries no channel gate on the + owner, so it must dispatch before the owner's slowed turn-0 completes. + """ + sink: list[_Fire] = [] + _install_timing_transport( + monkeypatch, sink, slow_sessions={"root": _SLOW_RESPONSE_S} + ) + fires = _run_graph(cli, sink, _FIXTURE_DIR / "prespawn.dag.jsonl") + by_key = _by_key(fires) + + owner_turn0 = by_key[("root", 0)] + pre_child = by_key[("bg-child", 0)] + + assert pre_child.dispatch_tick < owner_turn0.complete_tick, ( + "pre-session child waited on the owner's turn-0 response: " + f"child.dispatch={pre_child.dispatch_tick} " + f"owner0.complete={owner_turn0.complete_tick}" + ) + + +def test_authored_turn_delay_honored_at_transport_seam( + cli: AIPerfCLI, + mmap_base_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """(d) A turn with ``delay: 250`` fires >= 250 ms after its predecessor's + completion, measured at the transport seam (lower bound only).""" + sink: list[_Fire] = [] + _install_timing_transport(monkeypatch, sink, slow_sessions={}) + fires = _run_graph(cli, sink, _FIXTURE_DIR / "delayed_turn.dag.jsonl") + by_key = _by_key(fires) + + turn0 = by_key[("seq", 0)] + turn1 = by_key[("seq", 1)] + + gap_s = (turn1.dispatch_perf_ns - turn0.complete_perf_ns) / 1e9 + assert gap_s >= _DELAY_LOWER_BOUND_S, ( + f"authored delay {_DELAY_TURN_MS}ms not honored: gap was {gap_s * 1000:.1f}ms " + f"(predecessor completion -> delayed-turn dispatch)" + ) + + +def test_fork_children_records_carry_dag_identity( + cli: AIPerfCLI, + mmap_base_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """(e) FORK-child records carry the legacy DAG identity: ``agent_depth=1`` + and ``parent_correlation_id`` equal to the forking parent node's OWN minted + ``x_correlation_id`` (not merely the same shape) -- root turns carry + ``agent_depth=0`` / no parent. Asserted on ``request_info`` at the transport + seam, the same object the record pipeline copies into + ``MetricRecordMetadata.agent_depth`` / ``parent_correlation_id``.""" + sink: list[_Fire] = [] + _install_timing_transport(monkeypatch, sink, slow_sessions={}) + fires = _run_graph(cli, sink, _FIXTURE_DIR / "fork_minimal.dag.jsonl") + by_key = _by_key(fires) + + root = by_key[("root", 0)] + assert root.agent_depth == 0 + assert root.parent_correlation_id is None + + for key in (("branch-a", 0), ("branch-a", 1), ("branch-b", 0), ("branch-b", 1)): + child = by_key[key] + assert child.agent_depth == 1, ( + f"fork child {key} agent_depth={child.agent_depth}, expected 1" + ) + # ALL turns of a child instance carry the SAME triggering parent's + # correlation id -- byte-equal to the parent node's minted one. + assert child.parent_correlation_id == root.x_correlation_id, ( + f"fork child {key} parent_correlation_id=" + f"{child.parent_correlation_id!r}, expected the parent node's " + f"{root.x_correlation_id!r}" + ) + + +def test_prespawn_child_record_carries_depth_one_no_parent( + cli: AIPerfCLI, + mmap_base_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """(e) A pre-session SPAWN child's record carries ``agent_depth=1`` with NO + parent correlation id (legacy ``start_pre_session_child``: no parent session + exists at pre-dispatch), while the owner root's turns stay depth 0.""" + sink: list[_Fire] = [] + _install_timing_transport(monkeypatch, sink, slow_sessions={}) + fires = _run_graph(cli, sink, _FIXTURE_DIR / "prespawn.dag.jsonl") + by_key = _by_key(fires) + + for key in (("root", 0), ("root", 1)): + owner = by_key[key] + assert owner.agent_depth == 0, key + assert owner.parent_correlation_id is None, key + + pre_child = by_key[("bg-child", 0)] + assert pre_child.agent_depth == 1 + assert pre_child.parent_correlation_id is None diff --git a/tests/component_integration/graph/test_dag_jsonl_fidelity.py b/tests/component_integration/graph/test_dag_jsonl_fidelity.py new file mode 100644 index 0000000000..1c08a40aca --- /dev/null +++ b/tests/component_integration/graph/test_dag_jsonl_fidelity.py @@ -0,0 +1,417 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hermetic acceptance gate for the dag_jsonl graph adapter: raw-export parity. + +The in-process byte-parity golden gate +(``tests/component_integration/graph/test_dag_jsonl_byte_parity.py``) proves the +graph adapter emits byte-identical wire bodies to the legacy path at the +transport seam. This file exercises the offline diff tool +(:func:`tools.dag_jsonl_fidelity.prove_parity`) that the external real-run proof +depends on: it drives ``prove_parity`` on hand-crafted synthetic export +fixtures. A matching pair PASSES; corrupted content, a missing request, and +reordered messages each FAIL with a pointed diff (so the tool cannot pass +vacuously). Extra guards prove criterion 3 (body byte-equal) is distinct from +criterion 2 (payload parsed-equal), that an unparsable x_request_id cannot be +silently dropped, and that an empty overlap fails. + +The REAL-subprocess half of this gate -- ONE shared ``aiperf-mock-server`` +answering a legacy (``--custom-dataset-type dag_jsonl``) and a graph +(``--graph-format dag_jsonl``) ``aiperf profile --export-level raw`` run, then +``prove_parity`` on the two exports -- lives in the INTEGRATION lane at +``tests/integration/graph/test_dag_jsonl_fidelity_real.py`` (Linux-only CI, +real subprocesses), mirroring the weka split. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import orjson +import pytest + +from tools.dag_jsonl_fidelity import prove_parity + +# --- synthetic export helpers --------------------------------------------- + +# (session_id, turn_index) requests a faithful pair dispatches, keyed identically +# on both planes. root has a spliced assistant reply so a message-order mutation +# has a role transition to catch. +_SPEC: list[tuple[str, int]] = [ + ("root", 0), + ("branch-a", 0), + ("branch-a", 1), +] + + +def _canonical_payload(session: str, turn: int) -> dict[str, Any]: + """The wire request a FAITHFUL run would export for one ``(session, turn)``. + + Two messages (user + an assistant reply, as a FORK/multi-turn child would + carry) so a reversed-order mutation flips a role, and a stable top-level key + order so criterion 3 (body byte-equal) has a fixed reference. + """ + return { + "model": "m", + "messages": [ + {"role": "user", "content": f"prompt {session} #{turn}"}, + {"role": "assistant", "content": f"resp-{session}-{turn}"}, + ], + "max_tokens": 64, + "stream": True, + } + + +def _legacy_row(session: str, turn: int, payload: dict[str, Any]) -> dict[str, Any]: + """A legacy raw-export line: identity carried directly in metadata.""" + return { + "metadata": { + "session_num": 0, + "conversation_id": session, + "turn_index": turn, + "x_correlation_id": f"legacy-{session}-uuid", + "benchmark_phase": "profiling", + "request_start_ns": 1_000_000_000_000, + }, + "payload": payload, + } + + +def _graph_row( + session: str, turn: int, payload: dict[str, Any], *, fire: int = 0 +) -> dict[str, Any]: + """A graph raw-export line: per-node identity folded into x_request_id. + + ``x_correlation_id`` is now an OPAQUE per-trajectory id (``{conversation}:: + {nonce}``) carrying no node identity, so it deliberately does NOT encode the + node -- the tool must recover identity from ``x_request_id`` (the worker's + ``_mint_x_request_id`` folds ``{node_id}::{nonce}``). ``conversation_id`` is + the ONE trajectory TEMPLATE shared by every node and ``metadata.turn_index`` + is a per-correlation fire counter -- the tool must ignore BOTH. + """ + node_id = f"{session}:{turn}" + return { + "metadata": { + "session_num": 0, + "conversation_id": "traceX", + "turn_index": fire, + "x_correlation_id": f"traceX::corr{fire:08x}", + "x_request_id": f"{node_id}::req{fire:08x}", + "benchmark_phase": "profiling", + "request_start_ns": 1_000_000_000_000, + }, + "payload": payload, + } + + +def _graph_instance_row( + session: str, turn: int, instance: int, payload: dict[str, Any] +) -> dict[str, Any]: + """A graph row for the ``instance``-th SPAWN instantiation of one template. + + Instance 1 keeps the bare ``session:turn`` node id; instance >= 2 carries the + ``#`` suffix the tree mints for repeated SPAWN instances (folded into + ``x_request_id`` by the worker's ``_mint_x_request_id``). Every instance strips + back to the SAME recovered ``(session, turn)`` key, so a template fired N times + contributes multiplicity N to that key. + """ + node_id = f"{session}:{turn}" if instance == 1 else f"{session}#{instance}:{turn}" + return { + "metadata": { + "session_num": 0, + "conversation_id": "traceX", + "turn_index": instance - 1, + "x_correlation_id": f"traceX::corr{instance:08x}", + "x_request_id": f"{node_id}::req{instance:08x}", + "benchmark_phase": "profiling", + "request_start_ns": 1_000_000_000_000, + }, + "payload": payload, + } + + +def _legacy_rows(specs: list[tuple[str, int]]) -> list[dict[str, Any]]: + return [_legacy_row(s, t, _canonical_payload(s, t)) for s, t in specs] + + +def _graph_rows(specs: list[tuple[str, int]]) -> list[dict[str, Any]]: + return [ + _graph_row(s, t, _canonical_payload(s, t), fire=i) + for i, (s, t) in enumerate(specs) + ] + + +def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> Path: + body = b"\n".join(orjson.dumps(r) for r in rows) + path.write_bytes(body + b"\n" if rows else b"") + return path + + +def _write_pair( + tmp_path: Path, + legacy_rows: list[dict[str, Any]], + graph_rows: list[dict[str, Any]], +) -> tuple[Path, Path]: + return ( + _write_jsonl(tmp_path / "legacy.jsonl", legacy_rows), + _write_jsonl(tmp_path / "graph.jsonl", graph_rows), + ) + + +# --- tests ----------------------------------------------------------------- + + +@pytest.mark.component_integration +def test_prove_parity_passes_on_faithful_pair(tmp_path: Path) -> None: + """A byte-identical faithful pair passes every criterion + coverage invariant.""" + legacy, graph = _write_pair(tmp_path, _legacy_rows(_SPEC), _graph_rows(_SPEC)) + report = prove_parity(legacy, graph) + assert report.passed, report.render() + assert report.messages.checked == len(_SPEC) + assert report.payload.checked == len(_SPEC) + assert report.body.checked == len(_SPEC) + # Non-vacuous: every criterion actually matched a request. + assert report.messages.passes == len(_SPEC) + assert report.counts_match and report.keys_match + + +@pytest.mark.component_integration +def test_prove_parity_fails_on_corrupted_content(tmp_path: Path) -> None: + """A single mutated message byte FAILS all three criteria for that key only.""" + graph_rows = _graph_rows(_SPEC) + # Corrupt branch-a turn 0's first user message content. + graph_rows[1]["payload"]["messages"][0]["content"] += "_CORRUPTED" + legacy, graph = _write_pair(tmp_path, _legacy_rows(_SPEC), graph_rows) + report = prove_parity(legacy, graph) + assert not report.passed, report.render() + assert any( + "'branch-a'" in m.where and m.where.endswith("turn=0") + for m in report.messages.mismatches + ), report.messages.render() + assert any("content" in m.detail for m in report.messages.mismatches) + # Only the corrupted request fails; the two untouched ones still pass. + assert report.messages.passes == len(_SPEC) - 1 + + +@pytest.mark.component_integration +def test_prove_parity_fails_on_missing_request(tmp_path: Path) -> None: + """A request present in legacy but missing from graph is an unmatched-key fail.""" + graph_rows = _graph_rows(_SPEC)[:-1] # drop branch-a turn 1 + legacy, graph = _write_pair(tmp_path, _legacy_rows(_SPEC), graph_rows) + report = prove_parity(legacy, graph) + assert not report.passed, report.render() + assert ("branch-a", 1) in report.only_legacy + assert not report.counts_match + + +@pytest.mark.component_integration +def test_prove_parity_fails_on_reordered_messages(tmp_path: Path) -> None: + """Reordering a request's messages (same set, wrong order) FAILS with a role diff.""" + graph_rows = _graph_rows(_SPEC) + graph_rows[0]["payload"][ + "messages" + ].reverse() # root: [user, assistant] -> [assistant, user] + legacy, graph = _write_pair(tmp_path, _legacy_rows(_SPEC), graph_rows) + report = prove_parity(legacy, graph) + assert not report.passed, report.render() + assert any("'root'" in m.where for m in report.messages.mismatches) + assert any("role" in m.detail for m in report.messages.mismatches) + assert report.messages.passes == len(_SPEC) - 1 + + +@pytest.mark.component_integration +def test_prove_parity_key_order_drift_passes_all_criteria(tmp_path: Path) -> None: + """Same keys/values in a different order: EVERY criterion passes. + + The graph plane authors model / stream / the token cap through native + node fields, so wire key POSITION legitimately differs from legacy; the + body criterion compares CANONICAL (sorted-keys) serializations and must + not flag pure order drift. + """ + graph_rows = _graph_rows(_SPEC) + payload = graph_rows[0]["payload"] + graph_rows[0]["payload"] = {k: payload[k] for k in reversed(list(payload))} + legacy, graph = _write_pair(tmp_path, _legacy_rows(_SPEC), graph_rows) + report = prove_parity(legacy, graph) + assert report.passed, report.render() + assert report.body.passed, report.body.render() + + +@pytest.mark.component_integration +def test_prove_parity_type_drift_fails_body_not_payload(tmp_path: Path) -> None: + """Value-TYPE drift Python ``==`` conflates: criterion 2 PASSES, canonical + criterion 3 FAILS. + + ``stream: 1`` vs ``stream: True`` compare equal as parsed dicts but + serialize differently, so the canonical body criterion stays a genuinely + sharper check than parsed equality even without key-order comparison. + """ + graph_rows = _graph_rows(_SPEC) + payload = dict(graph_rows[0]["payload"]) + assert payload["stream"] is True + payload["stream"] = 1 + graph_rows[0]["payload"] = payload + legacy, graph = _write_pair(tmp_path, _legacy_rows(_SPEC), graph_rows) + report = prove_parity(legacy, graph) + assert not report.passed, report.render() + assert report.messages.passed, report.messages.render() + assert report.payload.passed, report.payload.render() + assert not report.body.passed, report.body.render() + assert any("'root'" in m.where for m in report.body.mismatches) + assert any("value/type drift" in m.detail for m in report.body.mismatches) + + +@pytest.mark.component_integration +def test_prove_parity_fails_on_unparsable_x_request_id(tmp_path: Path) -> None: + """A graph row whose x_request_id does not parse cannot be silently dropped. + + Identity-scheme drift is exactly what this gate must catch: the row surfaces + as an ONLY-GRAPH sentinel key while its legacy counterpart is ONLY-LEGACY, so + the proof fails loudly instead of vacuously passing. + """ + graph_rows = _graph_rows(_SPEC) + graph_rows[0]["metadata"]["x_request_id"] = "garbage-no-delimiters" + legacy, graph = _write_pair(tmp_path, _legacy_rows(_SPEC), graph_rows) + report = prove_parity(legacy, graph) + assert not report.passed, report.render() + assert ("root", 0) in report.only_legacy + assert any(k[1] == -1 for k in report.only_graph), report.render() + + +@pytest.mark.component_integration +def test_prove_parity_empty_exports_fail_vacuously(tmp_path: Path) -> None: + """Two empty exports checked nothing -> the proof must FAIL, not pass.""" + legacy, graph = _write_pair(tmp_path, [], []) + report = prove_parity(legacy, graph) + assert not report.passed, report.render() + assert not report.counts_match + + +def _kid_variant(content: str) -> dict[str, Any]: + """A ``("kid", 0)`` payload with a distinguishing user-message content. + + Two distinct variants keyed identically prove the multiset comparison pairs + by SORTED body (order-insensitive), not by dispatch order. + """ + payload = _canonical_payload("kid", 0) + payload["messages"][0]["content"] = content + return payload + + +@pytest.mark.component_integration +def test_prove_parity_passes_on_matching_multiset(tmp_path: Path) -> None: + """A ``(session, turn)`` key with multiplicity 2 passes when the body multisets + match -- even when the two instances are dispatched in OPPOSITE order. + + A repeated SPAWN template fires N times under distinct ``#n`` node ids that all + strip to one ``(session, turn)`` key; the comparator sorts each side's bodies + and compares element-wise, so order-permuted-but-equal multisets pass. + """ + a = _kid_variant("kid instance A") + b = _kid_variant("kid instance B") + legacy_rows = [ + _legacy_row("root", 0, _canonical_payload("root", 0)), + _legacy_row("kid", 0, a), + _legacy_row("kid", 0, b), + ] + graph_rows = [ + _graph_row("root", 0, _canonical_payload("root", 0)), + _graph_instance_row("kid", 0, 1, b), # opposite order from legacy + _graph_instance_row("kid", 0, 2, a), + ] + legacy, graph = _write_pair(tmp_path, legacy_rows, graph_rows) + report = prove_parity(legacy, graph) + assert report.passed, report.render() + # Multiplicity 2 for ("kid", 0) => 3 comparisons per criterion (root + 2 kids). + assert report.messages.checked == 3 + assert report.messages.passes == 3 + assert report.counts_match and report.keys_match + assert report.multiplicity_mismatch == [] + + +@pytest.mark.component_integration +def test_prove_parity_fails_on_corruption_within_multiplicity_key( + tmp_path: Path, +) -> None: + """Corrupting ONE of the two bodies under a multiplicity-2 key FAILS with a + pointed, instance-indexed diff (the other instance still matches). + + Proves the multiset comparator cannot be fooled by a matching sibling: a + single corrupted instance breaks the sorted element-wise equality. + """ + a = _kid_variant("kid instance A") + b = _kid_variant("kid instance B") + b_corrupt = _kid_variant("kid instance B") + b_corrupt["messages"][0]["content"] += "_CORRUPTED" + legacy_rows = [ + _legacy_row("root", 0, _canonical_payload("root", 0)), + _legacy_row("kid", 0, a), + _legacy_row("kid", 0, b), + ] + graph_rows = [ + _graph_row("root", 0, _canonical_payload("root", 0)), + _graph_instance_row("kid", 0, 1, a), + _graph_instance_row("kid", 0, 2, b_corrupt), + ] + legacy, graph = _write_pair(tmp_path, legacy_rows, graph_rows) + report = prove_parity(legacy, graph) + assert not report.passed, report.render() + assert any( + "'kid'" in m.where and m.where.endswith("turn=0") + for m in report.messages.mismatches + ), report.messages.render() + # The failing mismatch is instance-indexed (which of the N instances failed). + assert any( + m.instance is not None and "content" in m.detail + for m in report.messages.mismatches + ), report.messages.render() + # Coverage stays non-vacuous: the matching kid + root still passed. + assert report.messages.passes == 2 + + +@pytest.mark.component_integration +def test_prove_parity_fails_on_multiplicity_mismatch(tmp_path: Path) -> None: + """A per-key multiplicity mismatch (2 vs 1) FAILS even when TOTAL counts match. + + Legacy fires ``("kid", 0)`` twice and ``("root", 0)`` once; graph fires + ``("root", 0)`` twice and ``("kid", 0)`` once. Total profiling counts are + equal (3 == 3) and the key SETS match, so the failure is isolated to the + multiplicity guard -- proving it is a distinct invariant, not a proxy for the + count check. + """ + a = _kid_variant("kid instance A") + b = _kid_variant("kid instance B") + legacy_rows = [ + _legacy_row("root", 0, _canonical_payload("root", 0)), + _legacy_row("kid", 0, a), + _legacy_row("kid", 0, b), + ] + graph_rows = [ + _graph_row("root", 0, _canonical_payload("root", 0)), + _graph_instance_row("root", 0, 2, _canonical_payload("root", 0)), + _graph_instance_row("kid", 0, 1, a), + ] + legacy, graph = _write_pair(tmp_path, legacy_rows, graph_rows) + report = prove_parity(legacy, graph) + assert not report.passed, report.render() + assert report.counts_match, report.render() # total counts DO match (3 == 3) + assert not report.keys_match, report.render() + assert (("kid", 0), 2, 1) in report.multiplicity_mismatch + assert (("root", 0), 1, 2) in report.multiplicity_mismatch + + +@pytest.mark.component_integration +def test_prove_parity_ignores_warmup_records(tmp_path: Path) -> None: + """PROFILING-only comparison: a stray graph warmup row is filtered, not compared. + + Graph auto-warmup can prime a node before profiling; those rows carry a + ``warmup`` phase and must not inflate the graph count or create a spurious + unmatched key. + """ + graph_rows = _graph_rows(_SPEC) + warmup = _graph_row("root", 0, _canonical_payload("root", 0), fire=99) + warmup["metadata"]["benchmark_phase"] = "warmup" + legacy, graph = _write_pair(tmp_path, _legacy_rows(_SPEC), [warmup, *graph_rows]) + report = prove_parity(legacy, graph) + assert report.passed, report.render() + assert report.graph_profiling == len(_SPEC) diff --git a/tests/component_integration/graph/test_dispatch_plumbing.py b/tests/component_integration/graph/test_dispatch_plumbing.py new file mode 100644 index 0000000000..8db367bea2 --- /dev/null +++ b/tests/component_integration/graph/test_dispatch_plumbing.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Resolved dataset-selection values reach ``GraphIRReplayStrategy``. + +Drives the REAL ``TimingConfig.from_run`` -> ``PhaseRunner._build_graph_ir_strategy`` +path for a weka graph workload and asserts the constructed strategy stores the +resolved ``dataset_sampling_strategy`` and ``allow_dataset_wrap`` on its +attributes. This is interface plumbing only: consumption of the values is +covered by ``test_wrap_guard.py`` and ``test_graph_sampling.py``. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from aiperf.common.enums import CreditPhase +from aiperf.config.flags.cli_config import CLIConfig +from aiperf.plugin.enums import DatasetSamplingStrategy, TimingMode +from aiperf.timing.config import TimingConfig +from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy +from tests.unit.conftest import make_run_from_cli + +pytestmark = pytest.mark.component_integration + +_FIX_DIR = Path(__file__).parents[2] / "unit" / "graph" / "fixtures" +_WEKA_MIN = _FIX_DIR / "weka_min.json" + + +def _graph_run_with_resolved_selection(): + """Real weka graph run with the resolved selection values ``GraphDispatchResolver`` publishes. + + ``make_run_from_cli`` leaves ``run.resolved`` at its defaults (the resolver + chain isn't run here), so we set the two fields explicitly to stand in for + what the post-scenario graph-dispatch resolver derives. + """ + cfg = CLIConfig( + model_names=["test-model"], + input_file=str(_WEKA_MIN), + request_count=3, + ) + run = make_run_from_cli(cfg) + run.resolved.dataset_sampling_strategy = DatasetSamplingStrategy.SHUFFLE + run.resolved.allow_dataset_wrap = True + return run + + +def _graph_runner_for(profiling_config): + """A ``PhaseRunner`` wired for the real ``_build_graph_ir_strategy`` seam. + + Mirrors ``tests/unit/graph/test_build_strategy_graph_ir.py``'s harness: + ``PhaseRunner.__new__`` with the collaborators the graph branch reads, and a + capturing ``GraphIRReplayStrategy`` subclass so the forwarded kwargs are + observable. + """ + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + from aiperf.timing.graph_channel import GraphPhaseChannel + from aiperf.timing.phase.runner import PhaseRunner + + parsed = from_weka_trace(str(_WEKA_MIN)) + captured: dict = {} + + class _CapturingStrategy(GraphIRReplayStrategy): + def __init__(self, **kw): + captured.update(kw) + super().__init__(**kw) + + class _Handler: + def set_graph_return_observer(self, obs) -> None: + self._obs = obs + + def set_graph_first_token_observer(self, obs) -> None: + self._ft_obs = obs + + channel = GraphPhaseChannel(parsed_graph=parsed) + runner = PhaseRunner.__new__(PhaseRunner) + runner._config = profiling_config + runner._conversation_source = None + runner._graph_channel = channel + runner._scheduler = None + runner._stop_checker = None + runner._credit_issuer = object() + runner._lifecycle = None + runner._callback_handler = _Handler() + return runner, captured, _CapturingStrategy + + +def test_resolved_selection_reaches_graph_strategy() -> None: + run = _graph_run_with_resolved_selection() + + tc = TimingConfig.from_run(run) + profiling = [p for p in tc.phase_configs if p.phase == CreditPhase.PROFILING] + assert profiling, "expected a graph profiling phase" + cfg = profiling[0] + assert cfg.timing_mode == TimingMode.GRAPH_IR + + # from_run copies the resolved selection onto the CreditPhaseConfig. + assert cfg.dataset_sampling_strategy == DatasetSamplingStrategy.SHUFFLE + assert cfg.allow_dataset_wrap is True + + runner, captured, StrategyClass = _graph_runner_for(cfg) + strategy = runner._build_graph_ir_strategy(StrategyClass) + + assert isinstance(strategy, GraphIRReplayStrategy) + # Stored on the strategy for Tasks 11/12 to consume. + assert strategy._dataset_sampling_strategy == DatasetSamplingStrategy.SHUFFLE + assert strategy._allow_dataset_wrap is True + # Forwarded as explicit kwargs by the runner, not left to defaults. + assert captured["dataset_sampling_strategy"] == DatasetSamplingStrategy.SHUFFLE + assert captured["allow_dataset_wrap"] is True diff --git a/tests/component_integration/graph/test_duplication_report.py b/tests/component_integration/graph/test_duplication_report.py new file mode 100644 index 0000000000..e4f22fe1ac --- /dev/null +++ b/tests/component_integration/graph/test_duplication_report.py @@ -0,0 +1,350 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dispatch duplication report (covers serial recycle). + +Component-level: no real worker / ZMQ / mock server. A 5-trace corpus is +replayed at concurrency 5 with a session cap (10) that forces every lane to +recycle once, so ``total_instances_started`` (10) exceeds +``distinct_loaded_traces`` (5) and the duplication factor is 2.0. + +Asserts the report contract: +- factor > 1 with cache-bust OFF emits a WARNING (clones collide on identical + prefixes -- duplication without the antidote); +- the SAME duplicated run with cache-bust ON emits NO warning (per-instance + markers make the duplication safe); +- the warning ALSO fires on the DURATION-cancel path (the most recycle-heavy + mode: lanes recycle until the timer, so the lane future is cancelled -- the + report must still emit, exactly once); +- the resolved ``endpoint.cache_bust`` reaches ``strategy._cache_bust`` through + the REAL ``TimingConfig.from_run`` -> ``PhaseRunner._build_graph_ir_strategy`` + seam (guards the documented three-touch-wiring trap). +""" + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import msgspec +import pytest + +from aiperf.common.enums import CacheBustTarget, CreditPhase + +pytestmark = [pytest.mark.component_integration, pytest.mark.asyncio] + +_FIX_DIR = Path(__file__).parents[2] / "unit" / "graph" / "fixtures" +_MIN = _FIX_DIR / "weka_min.json" + +_DISTINCT_TRACES = 5 +_SESSION_CAP = 10 # 2x the corpus -> every lane recycles once (factor == 2.0) + + +@dataclass +class _EchoIssuer: + """Fake CreditIssuer: each ``issue_graph_credit`` schedules an echoed return. + + Mirrors the harness in ``test_graph_ir_replay_strategy`` -- the echoed + return is delivered on a later loop tick to the installed graph-return + observer, simulating the worker round-trip without a worker. + """ + + observer: Any = None + issued: int = 0 + returned: int = 0 + sending_complete_calls: int = 0 + all_returned_event_set: bool = False + + async def issue_graph_credit(self, turn: Any) -> bool: + self.issued += 1 + asyncio.get_running_loop().call_soon(self._echo, turn) + return True + + def _echo(self, turn: Any) -> None: + self.returned += 1 + if self.observer is not None: + self.observer(turn, None, False) + + def mark_graph_sending_complete(self) -> None: + self.sending_complete_calls += 1 + + def graph_all_returned(self) -> bool: + return self.returned >= self.issued + + def set_graph_all_returned_event(self) -> None: + self.all_returned_event_set = True + + +class _PhaseCfg: + """Per-phase config stub carrying the fields the strategy reads. + + ``expected_num_sessions`` caps distinct roots ever started (the recycle + gate), so a cap above the corpus size forces recycles -- serial duplication. + ``sessions=None`` leaves recycle uncapped so a duration budget is the only + bound (the duration-cancel scenario). + """ + + def __init__(self, *, concurrency: int, sessions: int | None) -> None: + self.phase = CreditPhase.PROFILING + self.concurrency = concurrency + self.expected_num_sessions = sessions + self.total_expected_requests = None + self.expected_duration_sec = None + + +class _DurationLifecycle: + """Lifecycle stub whose ``time_left_in_seconds`` reports a fixed budget. + + A float means ``--benchmark-duration`` is set: the duration stop condition + bounds the phase and ``_run_traces_under_duration_budget`` wraps the lane + dispatch in ``asyncio.wait_for(timeout=...)``, cancelling it when the budget + elapses -- the recycle-heavy duration-cancel path. + """ + + def __init__(self, remaining: float) -> None: + self._remaining = remaining + + def time_left_in_seconds(self, include_grace_period: bool = False) -> float: + return self._remaining + + +def _five_trace_corpus(*, gap_free: bool = False): + """Return a ParsedGraph whose ``traces`` holds 5 distinct-id instances. + + Replicates the single ``weka_min`` template with ``#N`` suffixes (the same + id-replace pattern the multi-trace strategy test uses); every clone resolves + to the base graph (``parsed.graphs`` is empty -> single-graph fallback) and + keys back to the base template id (the worker strips everything after the + first ``#``). + + ``gap_free`` zeros every edge delay so instances replay instantly -- used by + the duration-cancel scenario, where a tiny real-time budget must let the + lanes recycle (factor >> 1) BEFORE the timer cancels them, rather than + parking on ``weka_min``'s recorded ~1s inter-turn idle gaps. + """ + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + + base = from_weka_trace(str(_MIN)) + if gap_free: + graph = base.graph + zeroed_edges = [ + msgspec.structs.replace( + e, + **{ + f: 0.0 + for f in ("delay_after_predecessor_us", "min_start_delay_us") + if getattr(e, f, None) is not None + }, + ) + for e in graph.edges + ] + base = msgspec.structs.replace( + base, graph=msgspec.structs.replace(graph, edges=zeroed_edges) + ) + t0 = base.traces[0] + clones = [t0] + clones.extend( + msgspec.structs.replace(t0, id=f"{t0.id}#{i}") + for i in range(1, _DISTINCT_TRACES) + ) + return msgspec.structs.replace(base, traces=clones) + + +def _make_strategy( + parsed, + issuer: _EchoIssuer, + cache_bust: CacheBustTarget, + *, + sessions: int | None = _SESSION_CAP, + lifecycle: Any = None, +): + from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy + + return GraphIRReplayStrategy( + config=_PhaseCfg(concurrency=_DISTINCT_TRACES, sessions=sessions), + conversation_source=None, + scheduler=None, + stop_checker=None, + credit_issuer=issuer, + lifecycle=lifecycle, + parsed_graph=parsed, + register_observer=lambda obs: setattr(issuer, "observer", obs), + max_concurrent_traces=_DISTINCT_TRACES, + start_min_ratio=0.0, + start_max_ratio=0.0, + cache_bust=cache_bust, + ) + + +def _duplication_warnings(records: list[logging.LogRecord]) -> list[str]: + return [ + r.getMessage() + for r in records + if r.levelno >= logging.WARNING and "duplication" in r.getMessage().lower() + ] + + +async def test_cache_bust_target_threads_to_strategy(): + """The resolved cache-bust target lands on ``self._cache_bust`` (T11 reuses it).""" + parsed = _five_trace_corpus() + strategy = _make_strategy( + parsed, _EchoIssuer(), cache_bust=CacheBustTarget.FIRST_TURN_PREFIX + ) + assert strategy._cache_bust == CacheBustTarget.FIRST_TURN_PREFIX + + +async def test_duplication_warning_emitted_when_cache_bust_off(caplog): + """factor > 1 with cache-bust OFF warns that clones collide on shared prefixes.""" + parsed = _five_trace_corpus() + assert len(parsed.traces) == _DISTINCT_TRACES + + issuer = _EchoIssuer() + strategy = _make_strategy(parsed, issuer, cache_bust=CacheBustTarget.NONE) + await strategy.setup_phase() + with caplog.at_level(logging.WARNING, logger="GraphIRReplayTiming"): + # Generous ceiling: the phase ends on its own budget almost immediately + # on an idle machine, but loaded Windows CI runners have been observed + # to need well over 15s of wall clock end to end. + await asyncio.wait_for(strategy.execute_phase(), timeout=60.0) + + # 5 lanes recycle up to the 10-session cap -> 10 instances over 5 traces. + assert strategy._admitted_traces == _SESSION_CAP + warnings = _duplication_warnings(caplog.records) + assert warnings, "duplication warning must be emitted when cache-bust is OFF" + assert any("cache-bust" in w.lower() for w in warnings) + + +async def test_duplication_warning_absent_when_cache_bust_on(caplog): + """The SAME duplicated run with cache-bust ON emits NO duplication warning.""" + parsed = _five_trace_corpus() + + issuer = _EchoIssuer() + strategy = _make_strategy( + parsed, issuer, cache_bust=CacheBustTarget.FIRST_TURN_PREFIX + ) + await strategy.setup_phase() + with caplog.at_level(logging.WARNING, logger="GraphIRReplayTiming"): + # Generous ceiling: the phase ends on its own budget almost immediately + # on an idle machine, but loaded Windows CI runners have been observed + # to need well over 15s of wall clock end to end. + await asyncio.wait_for(strategy.execute_phase(), timeout=60.0) + + # Duplication still happened (same recycle dynamics)... + assert strategy._admitted_traces == _SESSION_CAP + # ...but the per-instance markers make it safe, so no warning fires. + assert not _duplication_warnings(caplog.records), ( + "duplication warning must be suppressed when cache-bust is ON" + ) + + +async def test_duplication_warning_emitted_on_duration_cancel(caplog): + """The report fires on the DURATION-cancel path -- the most recycle-heavy mode. + + No session cap: lanes recycle until the duration budget elapses, at which + point ``_run_traces_under_duration_budget`` CANCELS the lane future. The + report must still emit (from the caller's ``finally``), EXACTLY once, so the + warning is not silently absent precisely where factor is largest. + """ + parsed = _five_trace_corpus(gap_free=True) + + issuer = _EchoIssuer() + strategy = _make_strategy( + parsed, + issuer, + cache_bust=CacheBustTarget.NONE, + sessions=None, + lifecycle=_DurationLifecycle(0.1), + ) + await strategy.setup_phase() + with caplog.at_level(logging.WARNING, logger="GraphIRReplayTiming"): + # Generous ceiling: the phase ends on its own budget almost immediately + # on an idle machine, but loaded Windows CI runners have been observed + # to need well over 15s of wall clock end to end. + await asyncio.wait_for(strategy.execute_phase(), timeout=60.0) + + # Recycle-heavy: far more instances started than the 5-trace corpus (factor + # >> 1), even though the lane future was cancelled by the timer. + assert strategy._instances_started > _DISTINCT_TRACES + warnings = _duplication_warnings(caplog.records) + # Fires despite the cancel, and EXACTLY once (the caller's finally runs once). + assert len(warnings) == 1, f"expected exactly one duplication warning: {warnings}" + assert "cache-bust" in warnings[0].lower() + + +def _capturing_graph_runner(profiling_config): + """A ``PhaseRunner`` wired for the real ``_build_graph_ir_strategy`` seam. + + Mirrors ``test_dispatch_plumbing``'s harness: ``PhaseRunner.__new__`` with the + collaborators the graph branch reads, and a capturing ``GraphIRReplayStrategy`` + subclass so the forwarded kwargs are observable. + """ + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + from aiperf.timing.graph_channel import GraphPhaseChannel + from aiperf.timing.phase.runner import PhaseRunner + from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy + + parsed = from_weka_trace(str(_MIN)) + captured: dict = {} + + class _CapturingStrategy(GraphIRReplayStrategy): + def __init__(self, **kw): + captured.update(kw) + super().__init__(**kw) + + class _Handler: + def set_graph_return_observer(self, obs) -> None: + self._obs = obs + + def set_graph_first_token_observer(self, obs) -> None: + self._ft_obs = obs + + channel = GraphPhaseChannel(parsed_graph=parsed) + runner = PhaseRunner.__new__(PhaseRunner) + runner._config = profiling_config + runner._conversation_source = None + runner._graph_channel = channel + runner._scheduler = None + runner._stop_checker = None + runner._credit_issuer = object() + runner._lifecycle = None + runner._callback_handler = _Handler() + return runner, captured, _CapturingStrategy + + +async def test_cache_bust_seam_reaches_strategy_end_to_end(): + """endpoint.cache_bust -> from_run -> runner -> strategy._cache_bust (3-touch seam). + + Drives the REAL wiring so a regression in ANY of the three touches + (CreditPhaseConfig field, runner forwarding, strategy storage) is caught -- + this repo has a documented three-touch-wiring trap. + """ + from aiperf.config.flags.cli_config import CLIConfig + from aiperf.plugin.enums import TimingMode + from aiperf.timing.config import TimingConfig + from tests.unit.conftest import make_run_from_cli + + cfg = CLIConfig( + model_names=["test-model"], + input_file=str(_MIN), + request_count=3, + cache_bust=CacheBustTarget.FIRST_TURN_PREFIX, + ) + run = make_run_from_cli(cfg) + assert run.cfg.endpoint.cache_bust == CacheBustTarget.FIRST_TURN_PREFIX + + tc = TimingConfig.from_run(run) + profiling = [p for p in tc.phase_configs if p.phase == CreditPhase.PROFILING] + assert profiling, "expected a graph profiling phase" + phase_cfg = profiling[0] + assert phase_cfg.timing_mode == TimingMode.GRAPH_IR + # Touch 1: from_run copies endpoint.cache_bust onto the CreditPhaseConfig. + assert phase_cfg.cache_bust == CacheBustTarget.FIRST_TURN_PREFIX + + runner, captured, StrategyClass = _capturing_graph_runner(phase_cfg) + strategy = runner._build_graph_ir_strategy(StrategyClass) + # Touch 3: stored on the strategy (T11 reuses ``self._cache_bust``). + assert strategy._cache_bust == CacheBustTarget.FIRST_TURN_PREFIX + # Touch 2: forwarded as an explicit kwarg by the runner, not left to default. + assert captured["cache_bust"] == CacheBustTarget.FIRST_TURN_PREFIX diff --git a/tests/component_integration/graph/test_dynamo_build_wiring.py b/tests/component_integration/graph/test_dynamo_build_wiring.py new file mode 100644 index 0000000000..945d93aeab --- /dev/null +++ b/tests/component_integration/graph/test_dynamo_build_wiring.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DatasetManager routes a dynamo graph into the unified store (unconditional). + +The dynamo trie parse stamps per-node ``prompt_segment_ids`` + a +``segment_pool`` at parse time (shared LCP segment-trie core), so the store +build streams the stamped parse's per-trace payloads into a +:class:`GraphSegmentUnifiedBackingStore` via +:func:`build_unified_trie_store_from_payloads` (the same drain the weka pool +path uses). This test drives the whole path through +``_configure_graph_workload`` and asserts a known node's manifest materializes +from the unified store (no ``GraphEnvelopeMissing`` / empty store). +""" + +from __future__ import annotations + +from pathlib import Path + +import orjson +import pytest + +from aiperf.config.flags.cli_config import CLIConfig +from aiperf.dataset.dataset_manager import DatasetManager +from aiperf.dataset.graph_segment_unified_store import GraphSegmentUnifiedClient +from aiperf.plugin.enums import EndpointType +from tests.unit.conftest import make_run_from_cli + +pytestmark = [pytest.mark.component_integration, pytest.mark.asyncio] + + +def _dynamo_record(ts: int, sid: str, input_tokens: int, hashes: list[int]) -> dict: + """One current-schema ``dynamo.request.trace.v1`` ``request_end`` record.""" + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": {"session_id": sid}, + "request": { + "request_id": f"r{ts}", + "model": "m", + "input_tokens": input_tokens, + "output_tokens": 8, + "cached_tokens": 0, + "replay": { + "trace_block_size": 16, + "input_length": input_tokens, + "input_sequence_hashes": hashes, + }, + }, + } + + +@pytest.fixture +def dyn_fixture(tmp_path: Path) -> Path: + """A minimal single-session two-turn dynamo trace (recorded replay hashes).""" + p = tmp_path / "dyn_min.jsonl" + records = [ + _dynamo_record(1000, "s1", 32, [111, 222]), + _dynamo_record(2000, "s1", 64, [111, 222, 333, 444]), + ] + p.write_bytes(b"\n".join(orjson.dumps(r) for r in records)) + return p + + +async def test_dynamo_build_populates_unified_store( + dyn_fixture: Path, + mmap_base_path: Path, +) -> None: + run = make_run_from_cli( + CLIConfig( + model_names=["m"], + endpoint_type=EndpointType.CHAT, + streaming=False, + url="http://localhost:8000", + input_file=str(dyn_fixture), + random_seed=1234, + ) + ) + dm = DatasetManager(run=run, service_id="test") + + convs = await dm._configure_graph_workload(dyn_fixture) + assert convs.trace_ids, "dynamo build produced no graph traces" + + # Known node: trace 's1', ordinal 0 (the first LlmNode). Its profiling manifest + # must materialize from the unified store -- proving the dynamo build wrote the + # replay-derived content pool + node manifests, never left empty. + with GraphSegmentUnifiedClient(mmap_base_path, dm.run.benchmark_id).open() as c: + raw = c.get_node_envelope("s1", 0, "profiling") + assert raw is not None, "node (s1, 0) missing from unified store" + envelope = orjson.loads(raw) + assert "handles" in envelope, f"unexpected manifest envelope: {envelope!r}" + assert envelope["handles"], "node manifest has no segment handles" + msgs = c.materialize_handles(envelope["handles"]) + assert msgs and all(set(m) == {"role", "content"} for m in msgs) diff --git a/tests/component_integration/graph/test_dynamo_e2e_materialize.py b/tests/component_integration/graph/test_dynamo_e2e_materialize.py new file mode 100644 index 0000000000..d32d5d6637 --- /dev/null +++ b/tests/component_integration/graph/test_dynamo_e2e_materialize.py @@ -0,0 +1,292 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end proof the dynamo trie IR is RUNNABLE + DEDUPS. + +Drives the REAL build path -- ``DatasetManager._configure_graph_workload`` -> +``parse_graph_workload`` (dynamo lowers through the shared LCP segment-trie +core at parse time) -> ``build_unified_trie_store_from_payloads`` -> +``GraphSegmentUnifiedBackingStore`` -- on a MULTI-TURN + SUBAGENT dynamo +trace, then materializes every node the SAME way the worker does +(``materialize_graph_request_unified`` reading the interned int-handle +manifests) and asserts the materialized messages equal the +:class:`SegmentPool` ground truth (``pool.materialize(path)``). No +``GraphEnvelopeMissing``, no empty prompts. + +Beyond byte-parity it validates the KV-cache dedup the trie IR exists for: +the root session's strictly-nested replay hashes make every later turn's +``prompt_segment_ids`` EXTEND the earlier turn's (shared segment-id prefix), +and the child session dedups internally the same way. The old channel-replay +mixed-K fan-in break is gone -- the flat lowering shares the prefix across +ALL root turns, including the one after the tool + subagent window. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import orjson +import pytest + +from aiperf.config.flags.cli_config import CLIConfig +from aiperf.dataset.dataset_manager import DatasetManager +from aiperf.dataset.graph.graph_path_catalog import ( + build_catalog_context, + node_ordinal_for, +) +from aiperf.dataset.graph.models import LlmNode +from aiperf.dataset.graph.workload_detect import parse_graph_workload +from aiperf.dataset.graph_segment_unified_store import GraphSegmentUnifiedClient +from aiperf.graph.worker_materialize import materialize_graph_request_unified +from aiperf.plugin.enums import EndpointType +from tests.unit.conftest import make_run_from_cli + +pytestmark = [pytest.mark.component_integration, pytest.mark.asyncio] + +_SEED = 1234 + + +def _request_end( + *, + ts: int, + session_id: str, + parent_session_id: str | None = None, + hashes: list[int] | None = None, +) -> dict: + """One ``dynamo.request.trace.v1`` ``request_end`` with recorded replay hashes. + + One 16-token block per replay hash so ``input_length == 16 * len(hashes)`` + stays block-aligned for the parse-time alignment gate at ``block_size=16`` + (``(n-1)*16 < input_length <= n*16``). + """ + ctx: dict = {"session_id": session_id} + if parent_session_id is not None: + ctx["parent_session_id"] = parent_session_id + input_length = 16 * len(hashes) if hashes else 32 + req: dict = { + "request_id": f"{session_id}-{ts}", + "model": "m", + "input_tokens": input_length, + "output_tokens": 16, + "cached_tokens": 0, + } + if hashes is not None: + req["replay"] = { + "trace_block_size": 16, + "input_length": input_length, + "input_sequence_hashes": hashes, + } + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": ctx, + "request": req, + } + + +def _tool_event( + *, ts: int, session_id: str, event_type: str, tool_call_id: str +) -> dict: + tool: dict = {"tool_call_id": tool_call_id, "tool_class": "search"} + if event_type == "tool_end": + tool["duration_ms"] = 40.0 + tool["status"] = "succeeded" + return { + "schema": "dynamo.request.trace.v1", + "event_type": event_type, + "event_time_unix_ms": ts, + "event_source": "harness", + "agent_context": {"session_id": session_id}, + "tool": tool, + } + + +@pytest.fixture +def dyn_subagent_fixture(tmp_path: Path) -> Path: + """Multi-turn root session (4 nested-hash turns) + a 2-turn subagent child. + + The root ``parent`` records strictly NESTED replay hashes + (``[111,222]`` prefix-of ``[..,333]`` prefix-of ``[..,444]`` prefix-of + ``[..,555]``), so the LCP content trie re-derives the SAME leading pool + segments turn over turn -- a genuine multi-element cross-turn shared + prefix. The child session (``child``, nested hashes ``[900,901]`` -> + ``[900,901,902]``) flattens into the same graph and dedups internally. + Tool events between root turns 3 and 4 land on turn 3's tool_breakdown + metadata (the trie IR has no tool nodes). + """ + p = tmp_path / "dyn_subagent.jsonl" + records = [ + _request_end(ts=1000, session_id="parent", hashes=[111, 222]), + _request_end(ts=1100, session_id="parent", hashes=[111, 222, 333]), + _request_end(ts=1200, session_id="parent", hashes=[111, 222, 333, 444]), + _tool_event( + ts=1220, + session_id="parent", + event_type="tool_start", + tool_call_id=":subagent:child:invoke", + ), + _tool_event( + ts=1260, + session_id="parent", + event_type="tool_end", + tool_call_id=":subagent:child:invoke", + ), + _request_end( + ts=1280, session_id="child", parent_session_id="parent", hashes=[900, 901] + ), + _request_end( + ts=1300, + session_id="child", + parent_session_id="parent", + hashes=[900, 901, 902], + ), + _request_end(ts=1400, session_id="parent", hashes=[111, 222, 333, 444, 555]), + ] + with p.open("wb") as f: + for r in records: + f.write(orjson.dumps(r)) + f.write(b"\n") + return p + + +def _build_dm(fixture: Path) -> DatasetManager: + run = make_run_from_cli( + CLIConfig( + model_names=["m"], + endpoint_type=EndpointType.CHAT, + streaming=False, + url="http://localhost:8000", + input_file=str(fixture), + random_seed=_SEED, + ) + ) + return DatasetManager(run=run, service_id="test") + + +def _reparse_pool_and_paths(dm: DatasetManager, fixture: Path): + """Re-derive the parse + pool through the SAME shared seam the build used. + + ``_configure_graph_workload`` builds the store but does not hand back the + parse. ``parse_graph_workload`` is deterministic from ``(run, path, env)``, + so re-running it yields the SAME pool + ``prompt_segment_ids`` per node -- + the ground truth the store must round-trip. + """ + parsed = parse_graph_workload(dm.run, fixture) + assert parsed.segment_pool is not None + return parsed, parsed.segment_pool + + +def _prompt_path(node: LlmNode) -> list[str]: + return node.metadata["trie"]["prompt_segment_ids"] + + +def _shared_prefix_len(a: list[str], b: list[str]) -> int: + n = 0 + for x, y in zip(a, b, strict=False): + if x != y: + break + n += 1 + return n + + +def _session_turns(parsed, session_id: str) -> list[tuple[str, LlmNode]]: + """The session's ``(node_id, node)`` turns in turn-index order.""" + turns = [ + (nid, node) + for nid, node in parsed.graph.nodes.items() + if isinstance(node, LlmNode) and nid.startswith(f"{session_id}:") + ] + return sorted(turns, key=lambda kv: int(kv[0].rsplit(":", 1)[-1])) + + +async def test_dynamo_build_store_worker_materialize_byte_equal( + dyn_subagent_fixture: Path, + mmap_base_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Every dynamo node (root + child session) materializes byte-equal via the + worker path.""" + dm = _build_dm(dyn_subagent_fixture) + benchmark_id = dm.run.benchmark_id + + convs = await dm._configure_graph_workload(dyn_subagent_fixture) + assert convs.trace_ids, "dynamo build produced no graph traces" + + stamped, pool = _reparse_pool_and_paths(dm, dyn_subagent_fixture) + ctx = build_catalog_context(stamped) + trace = stamped.traces[0] + + checked = 0 + with GraphSegmentUnifiedClient(mmap_base_path, benchmark_id).open() as client: + for node_id, node in stamped.graph.nodes.items(): + assert isinstance(node, LlmNode) + ordinal = node_ordinal_for(ctx, trace.id, node_id) + assert ordinal is not None, f"node {node_id!r} has no ordinal" + req = materialize_graph_request_unified( + client, trace.id, ordinal, "profiling" + ) + assert req is not None, ( + f"node {node_id!r} (ord {ordinal}) did not materialize " + "(GraphEnvelopeMissing)" + ) + assert req["messages"] == pool.materialize(_prompt_path(node)), node_id + checked += 1 + + assert checked >= 6, "root (4 turns) + child (2 turns) must all materialize" + + +async def test_dynamo_multiturn_shares_multi_element_prefix_and_dedups( + dyn_subagent_fixture: Path, + mmap_base_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """CONFIRM the KV-cache dedup the trie IR exists for -- across-turn shared + prefix on the root session, internal dedup on the child session, and + pool-level sharing overall.""" + dm = _build_dm(dyn_subagent_fixture) + await dm._configure_graph_workload(dyn_subagent_fixture) + + stamped, _pool = _reparse_pool_and_paths(dm, dyn_subagent_fixture) + + root = _session_turns(stamped, "parent") + assert len(root) == 4, "root session must contribute 4 turns" + # Every later root turn extends the previous turn's whole path -- including + # turn 4, which follows the tool + subagent window (the old channel-replay + # mixed-K fan-in used to break this). + for (_prev_id, prev), (_curr_id, curr) in zip(root, root[1:], strict=False): + p_prev, p_curr = _prompt_path(prev), _prompt_path(curr) + assert p_curr[: len(p_prev)] == p_prev, ( + f"turn {_curr_id} does not extend {_prev_id}: {p_prev} vs {p_curr}" + ) + shared = _shared_prefix_len(_prompt_path(root[1][1]), _prompt_path(root[2][1])) + assert shared >= 2, ( + "CONTENT-MODEL CONCERN: clean-chain root turns do NOT share a " + f"multi-element segment prefix (shared={shared}); each turn would be " + "one monolithic message with no cross-turn KV dedup." + ) + + # The child session ALSO dedups internally: turn 2 extends turn 1. + child = _session_turns(stamped, "child") + assert len(child) == 2, "fixture must produce a 2-turn child session" + c1, c2 = _prompt_path(child[0][1]), _prompt_path(child[1][1]) + assert c2[: len(c1)] == c1, f"child turn 2 does not extend turn 1: {c1} vs {c2}" + + # Pool-level dedup: total unique segments < sum of every node's path length. + all_paths = [ + _prompt_path(n) for n in stamped.graph.nodes.values() if isinstance(n, LlmNode) + ] + total = sum(len(p) for p in all_paths) + unique = len({sid for p in all_paths for sid in p}) + assert unique < total, ( + f"CONTENT-MODEL CONCERN: no dedup -- {unique} unique segments == " + f"{total} total across paths; the pool is not sharing any content." + ) + + # Emit the evidence when explicitly requested. + if os.environ.get("AIPERF_TASK10_EVIDENCE"): + for nid, node in root + child: + print(f"\nEVIDENCE {nid} path: {_prompt_path(node)}") + print(f"EVIDENCE shared(turn2,turn3) prefix len: {shared}") + print(f"EVIDENCE unique={unique} total={total}") diff --git a/tests/component_integration/graph/test_dynamo_subagent_only_dispatch.py b/tests/component_integration/graph/test_dynamo_subagent_only_dispatch.py new file mode 100644 index 0000000000..949ce8c240 --- /dev/null +++ b/tests/component_integration/graph/test_dynamo_subagent_only_dispatch.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A parent + subagent dynamo trace dispatches through the schedule-plane +``TraceExecutor``. + +In the flat trie IR every session's turns are ordinary ``LlmNode``s writing +their scalar +dispatch placeholder to a per-node ``{node_id}_out`` scratch channel (TEXT / +overwrite, scalar-tolerant), with concurrency emergent from the recorded +intervals. This drives the REAL schedule-plane parse (``parse_graph_workload`` +-> ``from_dynamo_trace`` trie lowering) through a ``TraceExecutor`` with a +scalar-returning stub issuer (the exact ``""`` the real +``CreditDispatchAdapter`` returns) and asserts the executor COMPLETES with +every node -- parent AND child session -- dispatched exactly once. +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path +from typing import Any + +import orjson +import pytest + +from aiperf.config.flags.cli_config import CLIConfig +from aiperf.dataset.graph.models import LlmNode +from aiperf.dataset.graph.workload_detect import parse_graph_workload +from aiperf.graph.executor import TraceExecutor +from aiperf.plugin.enums import EndpointType +from tests.unit.conftest import make_run_from_cli + +pytestmark = [pytest.mark.component_integration, pytest.mark.asyncio] + +_SEED = 1234 + + +def _request_end( + *, + ts: int, + session_id: str, + hashes: list[int], + parent_session_id: str | None = None, +) -> dict[str, Any]: + """One ``dynamo.request.trace.v1`` ``request_end`` with recorded replay hashes. + + One 16-token block per replay hash so ``input_length == 16 * len(hashes)`` + stays block-aligned for the parse-time alignment gate. + """ + input_length = 16 * len(hashes) + ctx: dict[str, Any] = {"session_id": session_id} + if parent_session_id is not None: + ctx["parent_session_id"] = parent_session_id + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": ctx, + "request": { + "request_id": f"{session_id}-{ts}", + "model": "m", + "input_tokens": input_length, + "output_tokens": 16, + "cached_tokens": 0, + "replay": { + "trace_block_size": 16, + "input_length": input_length, + "input_sequence_hashes": hashes, + }, + }, + } + + +@pytest.fixture +def subagent_fixture(tmp_path: Path) -> Path: + """Parent A (3 turns) + child B whose turns interleave A's K=2..3 window. + + B's turns (ts 1150 / 1170) fall between A's K=2 (1100) and K=3 (1200) + recorded points, so the flat lowering chains them by finished-before order, + with every scalar placeholder write landing on a per-node scratch channel. + """ + p = tmp_path / "subagent_flat.jsonl" + records = [ + _request_end(ts=1000, session_id="A", hashes=[11, 22]), + _request_end(ts=1100, session_id="A", hashes=[11, 22, 33]), + _request_end(ts=1150, session_id="B", parent_session_id="A", hashes=[90, 91]), + _request_end( + ts=1170, session_id="B", parent_session_id="A", hashes=[90, 91, 92] + ), + _request_end(ts=1200, session_id="A", hashes=[11, 22, 33, 44]), + ] + with p.open("wb") as f: + for r in records: + f.write(orjson.dumps(r)) + f.write(b"\n") + return p + + +def _parse(fixture: Path): + run = make_run_from_cli( + CLIConfig( + model_names=["m"], + endpoint_type=EndpointType.CHAT, + streaming=False, + url="http://localhost:8000", + input_file=str(fixture), + random_seed=_SEED, + ) + ) + return parse_graph_workload(run, fixture) + + +class _ScalarIssuer: + """Stub credit issuer returning the scalar placeholder the real + ``CreditDispatchAdapter.dispatch`` returns (``""``) -- the exact value that + would unwind a list-reducer channel if a placeholder write ever reached one.""" + + def __init__(self) -> None: + self.n = 0 + + async def dispatch(self, node, request, ctx, **kw) -> str: + self.n += 1 + return "" + + +async def test_flat_parent_child_dispatches_through_executor( + subagent_fixture: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The flat parent+child graph runs to completion through the real executor.""" + + parsed = _parse(subagent_fixture) + + # Coverage guard: the flat lowering really produced parent + child session + # LlmNodes in ONE graph, each writing a declared per-node scratch channel. + llm_nodes = { + nid: n for nid, n in parsed.graph.nodes.items() if isinstance(n, LlmNode) + } + assert len(llm_nodes) == 5, sorted(parsed.graph.nodes) + assert "messages" not in parsed.graph.state + for nid, node in llm_nodes.items(): + assert node.output == f"{nid}_out" + assert f"{nid}_out" in parsed.graph.state + + # Drive the REAL executor with a scalar-returning issuer over every trace. + # Must complete without raising: every scalar placeholder write lands on a + # scalar-tolerant scratch channel. + issuer = _ScalarIssuer() + ex = TraceExecutor(parsed, credit_issuer=issuer) + async with asyncio.TaskGroup(): + for trace in parsed.traces: + await ex.run(trace) + + assert issuer.n == 5, ( + f"expected 3 parent + 2 child dispatches; got {issuer.n} -- some flat " + "node never fired through the executor" + ) diff --git a/tests/component_integration/graph/test_graph_ir_completion_and_accounting.py b/tests/component_integration/graph/test_graph_ir_completion_and_accounting.py new file mode 100644 index 0000000000..6a33130a44 --- /dev/null +++ b/tests/component_integration/graph/test_graph_ir_completion_and_accounting.py @@ -0,0 +1,322 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fix C — GraphIRReplayStrategy completion + accounting against the REAL bridge. + +These tests wire the strategy to the production ``CreditIssuer`` + +``PhaseProgressTracker`` + ``CreditCallbackHandler`` (the actual event bridge the +``PhaseRunner`` waits on), with a fake ``CreditRouter`` that echoes each issued +``Credit`` back as a ``CreditReturn`` on a later tick (simulating the worker +round-trip). They prove the two PROVEN bugs are fixed: + +* F1/A3 — no-cap completion: with NO request/session/duration cap the phase + STILL completes -- ``execute_phase`` returns, ``all_credits_sent_event`` is + set on executor-drain (the strategy owns it), and ``all_credits_returned_event`` + is set once every issued graph credit returns. (RED: the strategy never set + the sent event, so a no-cap run hung on ``event.wait()`` forever.) +* F2/A2 — node-as-session: graph credits BYPASS ``CreditCounter`` session + arithmetic; ``--num-conversations N`` runs N WHOLE traces (not N nodes) and + the sent-count is NOT frozen after node 1 of trace 1. ``sent_sessions`` stays + 0 for graph credits. +* ``--request-count N`` caps total node dispatches at N then stops cleanly. + +A direct ``CreditCounter`` unit test (no-trace path) guards that non-graph +session accounting is byte-for-byte unchanged. +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path +from typing import Any + +import pytest + +pytestmark = [pytest.mark.component_integration, pytest.mark.asyncio] + +_FIX_DIR = Path(__file__).parents[2] / "unit" / "graph" / "fixtures" +_MIN = _FIX_DIR / "weka_min.json" + + +def _parsed(path: Path): + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + + return from_weka_trace(str(path)) + + +async def _expected_dispatch_count(parsed) -> int: + """Count the LLM/Tool dispatches the executor fires for the whole corpus.""" + from aiperf.graph.executor import TraceExecutor + + class _Rec: + def __init__(self) -> None: + self.n = 0 + + async def dispatch(self, node, request, ctx, **kw): + self.n += 1 + return "x" + + rec = _Rec() + ex = TraceExecutor(parsed, credit_issuer=rec) + async with asyncio.TaskGroup(): + for trace in parsed.traces: + await ex.run(trace) + return rec.n + + +def _multi_trace_parsed(n: int): + """Replicate the single-trace corpus into ``n`` distinct trace instances.""" + import msgspec + + base = _parsed(_MIN) + t0 = base.traces[0] + traces = [msgspec.structs.replace(t0, id=f"{t0.id}#{i}") for i in range(n)] + return msgspec.structs.replace(base, traces=traces) + + +class _EchoRouter: + """Fake CreditRouter: echoes each sent Credit back via the callback handler. + + ``send_credit`` schedules ``CreditReturn`` delivery to the real + ``CreditCallbackHandler.on_credit_return`` on a later tick -- exactly the + worker round-trip, exercising the production graph-return observer + + counter + event bridge. + """ + + def __init__(self) -> None: + self.handler: Any = None + self.sent: list[Any] = [] + + async def wait_for_workers(self, timeout: float) -> None: + return None + + async def cancel_all_credits(self) -> None: + return None + + async def send_credit(self, credit: Any) -> None: + from aiperf.credit.messages import CreditReturn + + self.sent.append(credit) + ret = CreditReturn(credit=credit, first_token_sent=True) + asyncio.get_running_loop().call_soon(self._deliver, ret) + + def _deliver(self, ret: Any) -> None: + asyncio.ensure_future(self.handler.on_credit_return("w0", ret)) + + +def _build_real_stack(parsed, *, requests=None, sessions=None, duration=None): + """Build the production issuer/progress/lifecycle/callback stack + strategy. + + Returns ``(strategy, issuer, progress, lifecycle, router)``. + """ + from aiperf.common.enums import CreditPhase + from aiperf.credit.callback_handler import CreditCallbackHandler + from aiperf.credit.issuer import CreditIssuer + from aiperf.plugin.enums import TimingMode as _TM + from aiperf.timing.concurrency import ConcurrencyManager + from aiperf.timing.config import CreditPhaseConfig + from aiperf.timing.phase.lifecycle import PhaseLifecycle + from aiperf.timing.phase.progress_tracker import PhaseProgressTracker + from aiperf.timing.phase.stop_conditions import StopConditionChecker + from aiperf.timing.request_cancellation import ( + RequestCancellationConfig, + RequestCancellationSimulator, + ) + from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy + + config = CreditPhaseConfig( + phase=CreditPhase.PROFILING, + timing_mode=_TM.GRAPH_IR, + total_expected_requests=requests, + expected_num_sessions=sessions, + expected_duration_sec=duration, + ) + lifecycle = PhaseLifecycle(config) + progress = PhaseProgressTracker(config) + stop_checker = StopConditionChecker( + config=config, lifecycle=lifecycle, counter=progress.counter + ) + concurrency = ConcurrencyManager() + concurrency.configure_for_phase(CreditPhase.PROFILING, None, None) + cancellation = RequestCancellationSimulator(RequestCancellationConfig()) + router = _EchoRouter() + issuer = CreditIssuer( + phase=CreditPhase.PROFILING, + stop_checker=stop_checker, + progress=progress, + concurrency_manager=concurrency, + credit_router=router, + cancellation_policy=cancellation, + lifecycle=lifecycle, + ) + handler = CreditCallbackHandler(concurrency) + strategy = GraphIRReplayStrategy( + config=config, + conversation_source=None, + scheduler=None, + stop_checker=stop_checker, + credit_issuer=issuer, + lifecycle=lifecycle, + parsed_graph=parsed, + register_observer=handler.set_graph_return_observer, + max_concurrent_traces=8, + start_min_ratio=0.0, + start_max_ratio=0.0, + ) + handler.register_phase( + phase=CreditPhase.PROFILING, + progress=progress, + lifecycle=lifecycle, + stop_checker=stop_checker, + strategy=strategy, + ) + router.handler = handler + lifecycle.start() + return strategy, issuer, progress, lifecycle, router + + +async def _drive(strategy, progress): + """Run the strategy and then wait for the returned event, all under a guard. + + Mirrors the PhaseRunner: execute_phase sets all_credits_sent_event on drain, + then we await all_credits_returned_event (set by the callback handler). + """ + await strategy.setup_phase() + await asyncio.wait_for(strategy.execute_phase(), timeout=5.0) + assert progress.all_credits_sent_event.is_set() + await asyncio.wait_for(progress.all_credits_returned_event.wait(), timeout=5.0) + + +async def test_no_cap_run_completes_and_sets_both_events(): + """F1/A3: no request/session/duration cap -> phase still completes. + + Proves the strategy bridges executor-drain to ``all_credits_sent_event`` and + that every issued graph credit returns (``all_credits_returned_event`` set). + """ + parsed = _parsed(_MIN) + expected = await _expected_dispatch_count(parsed) + strategy, issuer, progress, lifecycle, router = _build_real_stack(parsed) + + await _drive(strategy, progress) + + assert len(router.sent) == expected + assert progress.counter.requests_sent == expected + assert progress.counter.final_requests_sent == expected + # Node-as-session bug fixed: NO graph credit bumps session counters. + assert progress.counter.sent_sessions == 0 + assert progress.counter.requests_completed == expected + + +async def test_num_conversations_runs_n_whole_traces(): + """F2/A2: --num-conversations N runs N WHOLE traces, not N nodes. + + The proven bug froze the sent count after node 1 of trace 1. Here N=2 over a + 4-trace corpus must dispatch exactly 2 traces' worth of nodes and complete. + """ + parsed = _multi_trace_parsed(4) + per_trace = await _expected_dispatch_count(_parsed(_MIN)) + strategy, issuer, progress, lifecycle, router = _build_real_stack( + parsed, sessions=2 + ) + + await _drive(strategy, progress) + + assert strategy.admitted_traces == 2 + assert strategy.completed_traces == 2 + # records == sum of the 2 admitted traces' node dispatches (NOT 2 nodes) + assert len(router.sent) == per_trace * 2 + assert progress.counter.requests_sent == per_trace * 2 + assert progress.counter.final_requests_sent == per_trace * 2 + # NOT frozen after node 1: with per_trace == 3, the proven bug froze at 1. + assert per_trace > 1 + assert progress.counter.final_requests_sent != 1 + assert progress.counter.sent_sessions == 0 + + +async def test_num_conversations_unset_runs_whole_corpus(): + """Unset --num-conversations replays every trace in the corpus.""" + parsed = _multi_trace_parsed(3) + per_trace = await _expected_dispatch_count(_parsed(_MIN)) + strategy, issuer, progress, lifecycle, router = _build_real_stack(parsed) + + await _drive(strategy, progress) + + assert strategy.admitted_traces == 3 + assert len(router.sent) == per_trace * 3 + + +async def test_request_count_caps_node_dispatches_then_stops_cleanly(): + """--request-count N caps total node dispatches at N then stops (no hang). + + The issuer's RequestCountStopCondition gate refuses issuance once + requests_sent >= N; ``issue_graph_credit`` returns False, which the adapter + turns into a clean per-trace stop. The phase must NOT hang. + """ + parsed = _multi_trace_parsed(4) + per_trace = await _expected_dispatch_count(_parsed(_MIN)) + cap = per_trace + 1 # straddles trace 1/2 boundary + strategy, issuer, progress, lifecycle, router = _build_real_stack( + parsed, requests=cap + ) + + await strategy.setup_phase() + await asyncio.wait_for(strategy.execute_phase(), timeout=5.0) + assert progress.all_credits_sent_event.is_set() + await asyncio.wait_for(progress.all_credits_returned_event.wait(), timeout=5.0) + + # Exactly ``cap`` node dispatches were placed on the wire, no more. + assert len(router.sent) == cap + assert progress.counter.requests_sent == cap + assert progress.counter.final_requests_sent == cap + + +async def test_drained_traces_release_all_per_trace_state(): + """A drained trace releases ALL of its per-trace state. + + This test PINS the invariant that the per-trace + resource (the ``CreditDispatchAdapter`` and its parked-Future ``_waiters``) + is FULLY released the moment a trace drains, and the phase's completion is + driven purely by ``sent == returned`` -- no "ready-but-queued" style + accounting term may hold + a drained trace open. A regression that introduced such an + accounting term would either leak adapters past drain or hang the phase. + + Parity-neutral: this is resource-cleanup correctness, NOT a wire/dispatch + behavior change, so it touches no parity oracle axis. + """ + parsed = _multi_trace_parsed(4) + per_trace = await _expected_dispatch_count(_parsed(_MIN)) + strategy, issuer, progress, lifecycle, router = _build_real_stack(parsed) + + # Record every adapter the strategy builds so we can assert post-drain that + # NONE retains an outstanding waiter (the per-tree ``outstanding`` analog). + built: list[Any] = [] + original_build = strategy._build_adapter + + def _record(trace_id: str, instance_id: str, **kwargs: Any): + adapter = original_build(trace_id, instance_id, **kwargs) + built.append(adapter) + return adapter + + strategy._build_adapter = _record # type: ignore[method-assign] + + await _drive(strategy, progress) + + assert strategy.completed_traces == 4 + assert len(built) == 4 + # Every trace's adapter was removed from the live registry on drain: no + # finished tree lingers (the "tree drained -> release slot" invariant). + assert strategy._adapters == {}, ( + f"drained traces leaked adapters: {sorted(strategy._adapters)}" + ) + # And no drained adapter still holds a parked dispatch Future -- every + # outstanding request settled (returned/cancelled), the ONLY thing that + # keeps a tree from draining now that the queued term is gone. + for adapter in built: + assert adapter.inflight_count == 0, ( + f"adapter {adapter._trace_id!r} retained " + f"{adapter.inflight_count} outstanding waiter(s) past drain" + ) + # Completion is sent == returned, no residual outstanding work. + assert progress.counter.requests_sent == per_trace * 4 + assert progress.counter.requests_completed == per_trace * 4 + assert progress.counter.sent_sessions == 0 diff --git a/tests/component_integration/graph/test_graph_ir_replay_strategy.py b/tests/component_integration/graph/test_graph_ir_replay_strategy.py new file mode 100644 index 0000000000..ac0fa10f10 --- /dev/null +++ b/tests/component_integration/graph/test_graph_ir_replay_strategy.py @@ -0,0 +1,355 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""R5 — GraphIRReplayStrategy drives a TraceExecutor per trace over a FAKE credit loop. + +Component-level: no real worker, no ZMQ, no mock server. A fake ``CreditIssuer`` +whose ``issue_graph_credit`` schedules an echoed ``CreditReturn`` delivered to the +strategy's installed graph-return observer (simulating the worker round-trip). + +Asserts the corrected D2 contract: +- the phase COMPLETES (``execute_phase`` returns; no hang/deadlock); +- the number of graph credits issued equals the trace's executor LLM/Tool + dispatch count; +- concurrent multi-trace runs all complete; +- a cancelled/errored return does NOT hang the phase. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import pytest + +pytestmark = [pytest.mark.component_integration, pytest.mark.asyncio] + +_FIX_DIR = Path(__file__).parents[2] / "unit" / "graph" / "fixtures" +_MIN = _FIX_DIR / "weka_min.json" + + +async def _expected_dispatch_count(parsed) -> int: + """Drive the executor with a recording stub to learn the LLM/Tool dispatch + count for a trace (the executor decides firing/fan-out, so we count what it + actually fires rather than re-deriving the topology).""" + from aiperf.graph.executor import TraceExecutor + + class _Rec: + def __init__(self) -> None: + self.n = 0 + + async def dispatch(self, node, request, ctx, **kw): + self.n += 1 + return "x" + + rec = _Rec() + ex = TraceExecutor(parsed, credit_issuer=rec) + async with asyncio.TaskGroup(): + for trace in parsed.traces: + await ex.run(trace) + return rec.n + + +@dataclass +class _FakeReturn: + """Drives the graph-return observer; carries (credit, error, cancelled).""" + + error: str | None = None + cancelled: bool = False + + +@dataclass +class _EchoIssuer: + """Fake CreditIssuer: each ``issue_graph_credit`` schedules an echoed return. + + The echoed ``CreditReturn`` is delivered on a later event-loop tick to the + strategy's installed graph-return observer, keyed by the issued turn's + ``(x_correlation_id, turn_index)`` -- exactly the correlation the real + ``CreditCallbackHandler`` graph-return hook forwards. This simulates the + worker round-trip without a worker. + """ + + observer: Any = None + behavior: Any = None # Callable[[turn], _FakeReturn] | None + issued: int = 0 + returned: int = 0 + sent: list[Any] = field(default_factory=list) + sending_complete_calls: int = 0 + all_returned_event_set: bool = False + + async def issue_graph_credit(self, turn: Any) -> bool: + self.issued += 1 + self.sent.append(turn) + ret = self.behavior(turn) if self.behavior is not None else _FakeReturn() + asyncio.get_running_loop().call_soon(self._echo, turn, ret) + return True + + def _echo(self, turn: Any, ret: _FakeReturn) -> None: + self.returned += 1 + if self.observer is None: + return + self.observer(turn, ret.error, ret.cancelled) + + # New D2 completion contract: the strategy drives these once its executors + # drain. The real ``CreditIssuer`` freezes counts + sets the phase events; + # the fake just records that the strategy honored the contract. + def mark_graph_sending_complete(self) -> None: + self.sending_complete_calls += 1 + + def graph_all_returned(self) -> bool: + # All echoed returns are scheduled via call_soon; at the moment the + # TaskGroup drains they have all been delivered (the executor awaited + # each dispatch Future), so every issued credit has returned. + return self.returned >= self.issued + + def set_graph_all_returned_event(self) -> None: + self.all_returned_event_set = True + + +def _parsed(path: Path): + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + + return from_weka_trace(str(path)) + + +def _make_strategy(parsed, issuer: _EchoIssuer, **overrides): + from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy + + # Pin the full-replay window so 0-arg call sites assert issued == full count. + overrides.setdefault("start_min_ratio", 0.0) + overrides.setdefault("start_max_ratio", 0.0) + strategy = GraphIRReplayStrategy( + config=overrides.pop("config", None), + conversation_source=None, + scheduler=None, + stop_checker=None, + credit_issuer=issuer, + lifecycle=overrides.pop("lifecycle", None), + parsed_graph=parsed, + register_observer=lambda obs: setattr(issuer, "observer", obs), + max_concurrent_traces=overrides.pop("max_concurrent_traces", 8), + **overrides, + ) + return strategy + + +async def test_strategy_completes_single_trace_and_issues_per_node_credits(): + parsed = _parsed(_MIN) + expected = await _expected_dispatch_count(parsed) + assert expected == 3 + + issuer = _EchoIssuer() + strategy = _make_strategy(parsed, issuer) + await strategy.setup_phase() + await asyncio.wait_for(strategy.execute_phase(), timeout=5.0) + + assert issuer.issued == expected + assert strategy.completed_traces == len(parsed.traces) + # D2: the strategy OWNS completion -- it must signal sending-complete once + # its executors drain (the PhaseRunner's all_credits_sent_event bridge). + assert issuer.sending_complete_calls >= 1 + + +async def test_strategy_completes_concurrent_multi_trace(): + # Two trace instances of the same template -> concurrent admission. + import msgspec + + base = _parsed(_MIN) + t0 = base.traces[0] + t1 = msgspec.structs.replace(t0, id=t0.id + "#1") + parsed = msgspec.structs.replace(base, traces=[t0, t1]) + per_trace = await _expected_dispatch_count(_parsed(_MIN)) + + issuer = _EchoIssuer() + strategy = _make_strategy(parsed, issuer, max_concurrent_traces=2) + await strategy.setup_phase() + await asyncio.wait_for(strategy.execute_phase(), timeout=5.0) + + assert strategy.completed_traces == 2 + assert issuer.issued == per_trace * 2 + + +async def test_strategy_does_not_hang_on_errored_return(): + parsed = _parsed(_MIN) + # A dispatch error on a node is CONTAINED (mid-conversation resilience): the + # executor does NOT unwind the trace's TaskGroup -- the failed request is + # recorded as an error record, and the conversation continues past the + # failed turn. So the trace COMPLETES (no hang) and is NOT counted as an + # errored trace. + issuer = _EchoIssuer(behavior=lambda turn: _FakeReturn(error="boom")) + strategy = _make_strategy(parsed, issuer) + await strategy.setup_phase() + + # Must NOT hang: the phase completes; the per-node failure is contained. + await asyncio.wait_for(strategy.execute_phase(), timeout=5.0) + assert strategy.completed_traces == len(parsed.traces) + assert strategy.errored_traces == 0 + + +async def test_strategy_does_not_hang_on_cancelled_return(): + parsed = _parsed(_MIN) + issuer = _EchoIssuer(behavior=lambda turn: _FakeReturn(cancelled=True)) + strategy = _make_strategy(parsed, issuer) + await strategy.setup_phase() + + # Cancelled return is contained the same way -- conversation continues, the + # trace completes, no errored-trace unwind. + await asyncio.wait_for(strategy.execute_phase(), timeout=5.0) + assert strategy.completed_traces == len(parsed.traces) + assert strategy.errored_traces == 0 + + +class _Phase: + """Minimal per-phase config stub carrying the phase + concurrency slots the + strategy reads (``phase`` for the warmup/profiling t* disposition).""" + + def __init__(self, phase: Any) -> None: + self.phase = phase + self.concurrency = None + self.expected_num_sessions = None + + +async def test_unset_concurrency_resolves_to_one(): + """An unset phase ``concurrency`` resolves to 1 (aiperf default), not 64. + + Regression for the removed ``AIPERF_GRAPH_MAX_CONCURRENT_TRACES`` fallback: + with no explicit override and no phase concurrency, the trace-admission + bound falls back to the plain aiperf default of 1. + """ + from aiperf.common.enums import CreditPhase + + parsed = _parsed(_MIN) + strategy = _make_strategy( + parsed, + _EchoIssuer(), + config=_Phase(CreditPhase.PROFILING), + max_concurrent_traces=None, + ) + assert strategy._max_concurrent == 1 + + +async def test_positive_window_profiling_phase_resumes_after_tstar(): + """PROFILING phase + the SAME t*=50% window replays the post-t* turns only + (the warmup<->profiling boundary), strictly fewer than the full 3-turn replay.""" + from aiperf.common.enums import CreditPhase + + parsed = _parsed(_MIN) + issuer = _EchoIssuer() + strategy = _make_strategy( + parsed, + issuer, + config=_Phase(CreditPhase.PROFILING), + start_min_ratio=0.5, + start_max_ratio=0.5, + t_star_random_seed=42, + ) + gt = next(iter(strategy._plans.values())) + assert gt.t_star_us > 0 + + await strategy.setup_phase() + await asyncio.wait_for(strategy.execute_phase(), timeout=5.0) + + full = await _expected_dispatch_count(parsed) + assert full == 3 + # Post-t* turns at offsets 1.5s and 3.0s -> 2 profiled dispatches (< full). + assert issuer.issued == 2, "profiling must replay only at/after-t* turns" + assert 0 < issuer.issued < full + assert strategy.completed_traces == 1 + + +class _DurationLifecycle: + """Lifecycle stub whose ``time_left_in_seconds`` mimics a configured budget. + + ``remaining`` is returned verbatim: a float means ``--benchmark-duration`` is + set (the duration stop condition bounds the phase); ``None`` means no duration + (count/session/bare mode), exactly the real ``PhaseLifecycle`` contract. + """ + + def __init__(self, remaining: float | None) -> None: + self._remaining = remaining + + def time_left_in_seconds(self, include_grace_period: bool = False) -> float | None: + return self._remaining + + +def _graph_with_max_gap(parsed, gap_us: float): + """Return ``parsed`` with every static edge stamped with ``gap_us`` delay. + + Lets the idle-gap advisory be exercised at the unit level without the + minute-scale integration fixture: ``_max_inter_turn_gap_seconds`` scans these + ``delay_after_predecessor_us`` values. + """ + import msgspec + + graph = parsed.graph + new_edges = [ + msgspec.structs.replace(edge, delay_after_predecessor_us=gap_us) + if hasattr(edge, "delay_after_predecessor_us") + else edge + for edge in graph.edges + ] + new_graph = msgspec.structs.replace(graph, edges=new_edges) + return msgspec.structs.replace(parsed, graph=new_graph) + + +async def test_max_inter_turn_gap_seconds_reads_edge_and_node_delays(): + """``_max_inter_turn_gap_seconds`` reports the largest recorded gap (seconds).""" + parsed = _graph_with_max_gap(_parsed(_MIN), 59_000_000.0) # 59s in us + strategy = _make_strategy(parsed, _EchoIssuer()) + assert strategy._max_inter_turn_gap_seconds() == pytest.approx(59.0) + + +async def test_max_inter_turn_gap_seconds_zero_for_gap_free_corpus(): + """A gap-free corpus (no edge/node delays) reports 0.0s -> advisory suppressed.""" + strategy = _make_strategy(_graph_with_max_gap(_parsed(_MIN), 0.0), _EchoIssuer()) + assert strategy._max_inter_turn_gap_seconds() == pytest.approx(0.0) + + +async def test_has_benchmark_duration_tracks_lifecycle_and_config(): + """``_has_benchmark_duration`` is True iff a duration budget is wired.""" + parsed = _parsed(_MIN) + no_dur = _make_strategy(parsed, _EchoIssuer(), lifecycle=_DurationLifecycle(None)) + assert no_dur._has_benchmark_duration() is False + + with_dur = _make_strategy(parsed, _EchoIssuer(), lifecycle=_DurationLifecycle(10.0)) + assert with_dur._has_benchmark_duration() is True + + +async def test_advisory_fires_for_idle_gap_corpus_without_duration(caplog): + """An idle-gap corpus with no duration emits the once-per-run NOTICE advisory.""" + import logging + + parsed = _graph_with_max_gap(_parsed(_MIN), 59_000_000.0) + strategy = _make_strategy(parsed, _EchoIssuer(), lifecycle=_DurationLifecycle(None)) + with caplog.at_level(logging.INFO, logger="GraphIRReplayTiming"): + strategy._advise_if_idle_gap_corpus_without_duration() + assert any("--benchmark-duration" in r.getMessage() for r in caplog.records), ( + "idle-gap-without-duration advisory must mention --benchmark-duration" + ) + + +async def test_advisory_suppressed_when_duration_set(caplog): + """A duration-bounded run never logs the advisory (the budget IS the bound).""" + import logging + + parsed = _graph_with_max_gap(_parsed(_MIN), 59_000_000.0) + strategy = _make_strategy(parsed, _EchoIssuer(), lifecycle=_DurationLifecycle(10.0)) + with caplog.at_level(logging.INFO, logger="GraphIRReplayTiming"): + strategy._advise_if_idle_gap_corpus_without_duration() + assert not any("--benchmark-duration" in r.getMessage() for r in caplog.records), ( + "advisory must be suppressed when a duration budget is configured" + ) + + +async def test_advisory_suppressed_for_gap_free_corpus(caplog): + """A gap-free corpus (sub-threshold gaps) never logs the advisory.""" + import logging + + parsed = _graph_with_max_gap(_parsed(_MIN), 0.0) + strategy = _make_strategy(parsed, _EchoIssuer(), lifecycle=_DurationLifecycle(None)) + with caplog.at_level(logging.INFO, logger="GraphIRReplayTiming"): + strategy._advise_if_idle_gap_corpus_without_duration() + assert not any("--benchmark-duration" in r.getMessage() for r in caplog.records), ( + "advisory must be suppressed for a gap-free corpus" + ) diff --git a/tests/component_integration/graph/test_graph_sampling.py b/tests/component_integration/graph/test_graph_sampling.py new file mode 100644 index 0000000000..e1fb262ee7 --- /dev/null +++ b/tests/component_integration/graph/test_graph_sampling.py @@ -0,0 +1,234 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GraphIRReplayStrategy honors ``--dataset-sampling-strategy``. + +Every cross-trace draw in the lane fan-out / recycle loop routes through +``GraphIRReplayStrategy._draw_index(x, total)``, which remaps the historical +monotonic ``x % total`` draw to a strategy-selected trace index: + +* ``sequential`` (or ``None``) -> ``x % total`` byte-for-byte (the historical + cursor-with-wrap draw; a golden sequence pins this). +* ``shuffle`` -> a seeded per-pass permutation: each pass of ``total`` draws + covers every index exactly once (without replacement), deterministic under a + fixed ``t_star_random_seed`` and sensitive to it. +* ``random`` -> coerced to ``shuffle`` semantics in this single-pass-per-cycle + context (no dup/omission within a pass); random == shuffle here. + +These construct the REAL strategy over REAL weka traces (no mocks) and exercise +the real ``_draw_index`` / ``_resolve_pass0_lanes`` path. Deselected by default; +run with ``-m component_integration``. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import msgspec +import pytest + +pytestmark = [pytest.mark.component_integration] + +_FIX_DIR = Path(__file__).parents[2] / "unit" / "graph" / "fixtures" +_MIN = _FIX_DIR / "weka_min.json" + + +class _Issuer: + """Minimal CreditIssuer stub: the sampling path never issues a credit.""" + + def issue_graph_credit(self, *a: Any, **k: Any) -> bool: + return True + + def mark_graph_sending_complete(self) -> None: ... + + def graph_all_returned(self) -> bool: + return True + + def set_graph_all_returned_event(self) -> None: ... + + +def _parsed_with_n_traces(n: int): + """Real weka parse whose corpus is ``n`` distinct traces (id-cloned). + + The single ``weka_min`` trace is cloned into ``n`` distinct trace ids so the + draw abstraction sees a corpus of size ``n`` to permute over. + """ + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + + base = from_weka_trace(str(_MIN)) + t0 = base.traces[0] + traces = [msgspec.structs.replace(t0, id=f"{t0.id}#{i}") for i in range(n)] + return msgspec.structs.replace(base, traces=list(traces)) + + +def _make_strategy(parsed, *, sampling: Any, seed: int = 42): + from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy + + return GraphIRReplayStrategy( + credit_issuer=_Issuer(), + parsed_graph=parsed, + register_observer=lambda obs: None, + start_min_ratio=0.0, + start_max_ratio=0.0, + t_star_random_seed=seed, + dataset_sampling_strategy=sampling, + max_concurrent_traces=8, + ) + + +# --------------------------------------------------------------------------- # +# (a) SEQUENTIAL / None: byte-identical to the historical ``x % total`` draw. +# --------------------------------------------------------------------------- # + + +def test_sequential_draw_index_is_byte_identical_golden(): + """``sequential`` maps ``x -> x % total`` exactly (pinned golden sequence).""" + n = 5 + parsed = _parsed_with_n_traces(n) + from aiperf.plugin.enums import DatasetSamplingStrategy + + strategy = _make_strategy(parsed, sampling=DatasetSamplingStrategy.SEQUENTIAL) + + drawn = [strategy._draw_index(x, n) for x in range(3 * n)] + # Golden: pure sequential-with-wrap over 3 full passes. + assert drawn == [0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4] + assert drawn == [x % n for x in range(3 * n)] + + +def test_none_sampling_draw_index_matches_sequential(): + """An unset (``None``) strategy is the historical sequential draw.""" + n = 4 + parsed = _parsed_with_n_traces(n) + strategy = _make_strategy(parsed, sampling=None) + + drawn = [strategy._draw_index(x, n) for x in range(2 * n)] + assert drawn == [x % n for x in range(2 * n)] + + +def test_sequential_pass0_trace_ids_are_corpus_order(): + """Real ``_resolve_pass0_lanes`` under ``sequential`` yields ``traces[i]``.""" + n = 5 + parsed = _parsed_with_n_traces(n) + from aiperf.plugin.enums import DatasetSamplingStrategy + + strategy = _make_strategy(parsed, sampling=DatasetSamplingStrategy.SEQUENTIAL) + traces = list(parsed.traces) + + pass0, cursor = strategy._resolve_pass0_lanes(traces, n) + assert [t.id for t in pass0] == [t.id for t in traces] + assert cursor == n + + # Draw-counter -> trace id over two passes is byte-identical to traces[x % N]. + drawn_ids = [traces[strategy._draw_index(x, n)].id for x in range(2 * n)] + assert drawn_ids == [traces[x % n].id for x in range(2 * n)] + + +# --------------------------------------------------------------------------- # +# (b) SHUFFLE (seeded): without-replacement permutation, deterministic + salted. +# --------------------------------------------------------------------------- # + + +def test_shuffle_single_pass_is_permutation_without_replacement(): + """A single ``shuffle`` pass covers all N indices exactly once (no dup/omit).""" + n = 8 + parsed = _parsed_with_n_traces(n) + from aiperf.plugin.enums import DatasetSamplingStrategy + + strategy = _make_strategy(parsed, sampling=DatasetSamplingStrategy.SHUFFLE) + + perm = [strategy._draw_index(x, n) for x in range(n)] + assert sorted(perm) == list(range(n)) # permutation: every index exactly once + # Cheap + consistent: repeated draws within the same pass are stable. + assert [strategy._draw_index(x, n) for x in range(n)] == perm + + +def test_shuffle_second_pass_is_a_fresh_permutation(): + """Pass 1 (indices N..2N-1) is its own without-replacement permutation.""" + n = 8 + parsed = _parsed_with_n_traces(n) + from aiperf.plugin.enums import DatasetSamplingStrategy + + strategy = _make_strategy(parsed, sampling=DatasetSamplingStrategy.SHUFFLE) + + pass1 = [strategy._draw_index(x, n) for x in range(n, 2 * n)] + assert sorted(pass1) == list(range(n)) + + +def test_shuffle_is_deterministic_for_same_seed(): + """Two constructions with the same seed produce the SAME permutation.""" + n = 8 + from aiperf.plugin.enums import DatasetSamplingStrategy + + a = _make_strategy( + _parsed_with_n_traces(n), + sampling=DatasetSamplingStrategy.SHUFFLE, + seed=42, + ) + b = _make_strategy( + _parsed_with_n_traces(n), + sampling=DatasetSamplingStrategy.SHUFFLE, + seed=42, + ) + assert [a._draw_index(x, n) for x in range(n)] == [ + b._draw_index(x, n) for x in range(n) + ] + + +def test_shuffle_differs_for_different_seed(): + """A different seed yields a different permutation order.""" + n = 8 + from aiperf.plugin.enums import DatasetSamplingStrategy + + a = _make_strategy( + _parsed_with_n_traces(n), + sampling=DatasetSamplingStrategy.SHUFFLE, + seed=42, + ) + b = _make_strategy( + _parsed_with_n_traces(n), + sampling=DatasetSamplingStrategy.SHUFFLE, + seed=1234, + ) + perm_a = [a._draw_index(x, n) for x in range(n)] + perm_b = [b._draw_index(x, n) for x in range(n)] + # Both are valid permutations, but seeded differently -> different order. + assert sorted(perm_a) == list(range(n)) + assert sorted(perm_b) == list(range(n)) + assert perm_a != perm_b + + +# --------------------------------------------------------------------------- # +# (c) RANDOM: single-pass -> coerced to without-replacement (== shuffle). +# --------------------------------------------------------------------------- # + + +def test_random_single_pass_is_without_replacement(): + """``random`` in single-pass mode covers all N indices exactly once.""" + n = 8 + parsed = _parsed_with_n_traces(n) + from aiperf.plugin.enums import DatasetSamplingStrategy + + strategy = _make_strategy(parsed, sampling=DatasetSamplingStrategy.RANDOM) + + perm = [strategy._draw_index(x, n) for x in range(n)] + assert sorted(perm) == list(range(n)) # no dup, no omission + + +def test_random_equals_shuffle_for_same_seed(): + """``random`` is coerced to ``shuffle`` semantics: identical order per seed.""" + n = 8 + from aiperf.plugin.enums import DatasetSamplingStrategy + + rnd = _make_strategy( + _parsed_with_n_traces(n), + sampling=DatasetSamplingStrategy.RANDOM, + seed=42, + ) + shf = _make_strategy( + _parsed_with_n_traces(n), + sampling=DatasetSamplingStrategy.SHUFFLE, + seed=42, + ) + assert [rnd._draw_index(x, n) for x in range(n)] == [ + shf._draw_index(x, n) for x in range(n) + ] diff --git a/tests/component_integration/graph/test_lane_ramp.py b/tests/component_integration/graph/test_lane_ramp.py new file mode 100644 index 0000000000..152b86b3a6 --- /dev/null +++ b/tests/component_integration/graph/test_lane_ramp.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Lane-level concurrency ramp on the graph replay path. + +``--concurrency-ramp-duration`` cannot throttle graph lanes through session +slots (graph credits bypass them), so the strategy exposes +``LaneSettableProtocol.set_lane_limit`` and the runner's concurrency ramper +drives it: lanes above the live limit PARK before their first instance and +are admitted as the ramper raises the limit 1 -> ``--concurrency``. + +Component-level: fake ``CreditIssuer`` echoing returns, no worker/ZMQ. The +decisive proofs: parked lanes dispatch NOTHING until admitted; raising the +limit releases exactly the newly-admitted lanes; the phase completes once the +limit reaches the target (the Ramper always reaches its target, and +duration-cancel cancels parked waiters cleanly through the TaskGroup). +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import msgspec +import pytest + +from aiperf.common.enums import CreditPhase +from aiperf.timing.strategies.core import LaneSettableProtocol + +pytestmark = [pytest.mark.component_integration, pytest.mark.asyncio] + +_FIX_DIR = Path(__file__).parents[2] / "unit" / "graph" / "fixtures" +_MIN = _FIX_DIR / "weka_min.json" + + +@dataclass +class _EchoIssuer: + """Fake CreditIssuer echoing each issued credit to the return observer.""" + + observer: Any = None + issued: int = 0 + returned: int = 0 + sending_complete_calls: int = 0 + sent: list[Any] = field(default_factory=list) + + async def issue_graph_credit(self, turn: Any) -> bool: + self.issued += 1 + self.sent.append(turn) + asyncio.get_running_loop().call_soon(self._echo, turn) + return True + + def _echo(self, turn: Any) -> None: + self.returned += 1 + if self.observer is not None: + self.observer(turn, None, False) + + def mark_graph_sending_complete(self) -> None: + self.sending_complete_calls += 1 + + def graph_all_returned(self) -> bool: + return self.returned >= self.issued + + def set_graph_all_returned_event(self) -> None: + return None + + +class _PhaseCfg: + def __init__(self, *, concurrency: int | None = None) -> None: + self.phase = CreditPhase.PROFILING + self.concurrency = concurrency + self.expected_num_sessions = None + self.total_expected_requests = None + self.expected_duration_sec = None + self.num_dataset_entries = None + self.max_context_length = None + + +def _corpus(n: int): + """``n`` distinct-id clones of the gap-free ``weka_min`` template.""" + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + + base = from_weka_trace(str(_MIN)) + graph = base.graph + zeroed = [ + msgspec.structs.replace( + e, + **{ + f: 0.0 + for f in ("delay_after_predecessor_us", "min_start_delay_us") + if getattr(e, f, None) is not None + }, + ) + for e in graph.edges + ] + base = msgspec.structs.replace( + base, graph=msgspec.structs.replace(graph, edges=zeroed) + ) + t0 = base.traces[0] + clones = [t0] + clones.extend(msgspec.structs.replace(t0, id=f"{t0.id}#{i}") for i in range(1, n)) + return msgspec.structs.replace(base, traces=clones) + + +def _make_strategy(parsed, issuer: _EchoIssuer, **overrides): + from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy + + overrides.setdefault("start_min_ratio", 0.0) + overrides.setdefault("start_max_ratio", 0.0) + return GraphIRReplayStrategy( + config=overrides.pop("config", None), + conversation_source=None, + scheduler=None, + stop_checker=None, + credit_issuer=issuer, + lifecycle=overrides.pop("lifecycle", None), + parsed_graph=parsed, + register_observer=lambda obs: setattr(issuer, "observer", obs), + **overrides, + ) + + +async def test_strategy_satisfies_lane_settable_protocol(): + """The runner's ramper fan-out gates on this isinstance check.""" + strategy = _make_strategy(_corpus(1), _EchoIssuer(), max_concurrent_traces=2) + assert isinstance(strategy, LaneSettableProtocol) + + +async def test_lane_admission_mechanics(): + """Admitted lanes pass immediately; parked lanes release on a raise; + the limit clamps to [1, concurrency].""" + strategy = _make_strategy(_corpus(4), _EchoIssuer(), max_concurrent_traces=4) + + strategy.set_lane_limit(2) + # Lanes 0 and 1 are admitted: the waits complete without a raise. + await asyncio.wait_for(strategy._wait_for_lane_admission(0), timeout=1.0) + await asyncio.wait_for(strategy._wait_for_lane_admission(1), timeout=1.0) + + parked = asyncio.create_task(strategy._wait_for_lane_admission(3)) + await asyncio.sleep(0) + assert not parked.done(), "lane 3 must park while the limit is 2" + + strategy.set_lane_limit(3) + await asyncio.sleep(0) + assert not parked.done(), "lane 3 must stay parked at limit 3" + + strategy.set_lane_limit(4) + await asyncio.wait_for(parked, timeout=1.0) + + # Clamps: never below 1, never above the resolved concurrency. + strategy.set_lane_limit(0) + assert strategy._lane_limit == 1 + strategy.set_lane_limit(99) + assert strategy._lane_limit == 4 + + +async def test_parked_lanes_dispatch_nothing_until_admitted(): + """E2E through execute_phase: with the limit held at 1, only lane 0's + instance dispatches; raising the limit to the target releases the parked + lanes and the phase completes covering the whole corpus.""" + parsed = _corpus(3) + issuer = _EchoIssuer() + strategy = _make_strategy( + parsed, + issuer, + config=_PhaseCfg(concurrency=3), + max_concurrent_traces=3, + allow_dataset_wrap=False, + ) + # Simulate the ramper's initial setter(1), which lands before execution. + strategy.set_lane_limit(1) + + await strategy.setup_phase() + run = asyncio.create_task(strategy.execute_phase()) + + # Let lane 0 finish its single-pass instance; lanes 1-2 stay parked. + for _ in range(200): + await asyncio.sleep(0) + assert not run.done(), "phase must not complete while lanes are parked" + dispatched_traces = {t.trace_id.split("::", 1)[0] for t in issuer.sent} + assert len(dispatched_traces) == 1, ( + f"only lane 0's template may dispatch under limit 1, got " + f"{sorted(dispatched_traces)}" + ) + + strategy.set_lane_limit(3) + await asyncio.wait_for(run, timeout=15.0) + assert strategy.completed_traces == 3 + assert issuer.sending_complete_calls >= 1 + + +async def test_no_ramp_admits_all_lanes_immediately(): + """Without a ramp the limit is born at the resolved concurrency and the + run is byte-identical to before the feature.""" + parsed = _corpus(3) + issuer = _EchoIssuer() + strategy = _make_strategy( + parsed, + issuer, + config=_PhaseCfg(concurrency=3), + max_concurrent_traces=3, + allow_dataset_wrap=False, + ) + assert strategy._lane_limit == 3 + await strategy.setup_phase() + await asyncio.wait_for(strategy.execute_phase(), timeout=15.0) + assert strategy.completed_traces == 3 diff --git a/tests/component_integration/graph/test_weka_trace_fidelity.py b/tests/component_integration/graph/test_weka_trace_fidelity.py new file mode 100644 index 0000000000..740fcbbec7 --- /dev/null +++ b/tests/component_integration/graph/test_weka_trace_fidelity.py @@ -0,0 +1,1200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Acceptance gate for the weka segment-trie IR: raw-export fidelity vs the REAL trace. + +Two layers: + +* COMPONENT lane (fast, hermetic): drive each :mod:`tools.weka_trace_fidelity` + Report function on small hand-crafted raw-export + trace inputs, including a + deliberately-CORRUPTED export that MUST fail each check (so the tool cannot pass + vacuously). These build a tiny real trie graph (``build_trie_graph`` + + ``SegmentPool``) so the expected content is the genuine pool materialization, not + a stub. + +* INTEGRATION lane (slow, real subprocess): run an ACTUAL + ``aiperf profile --export-level raw`` of the trie path against the in-repo + ``aiperf-mock-server`` over the subagent fixture, then assert BOTH + :func:`content_vs_real_trace` and :func:`causality_timing_vs_real_trace` PASS on + the produced ``profile_export_raw.jsonl``. This is the empirical proof the + dispatched prompts and the dispatch causality/timing match the recorded trace. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import orjson +import pytest + +from tools.weka_trace_fidelity import ( + _ExportRecord, + _RecordedNode, + _RecordedTrace, + build_recorded_trace, + causality_timing_vs_real_trace, + content_byte_exact_vs_v04, + content_vs_real_trace, + prove_corpus, +) + +_REPO = Path(__file__).resolve().parents[3] +_FIX = _REPO / "tests" / "unit" / "graph" / "fixtures" / "weka_subagent.json" +_MODEL = "claude-opus-4-5-20251101" + + +# --- synthetic-input helpers ---------------------------------------------- + + +def _build_expected(trace_file: Path) -> dict[str, list[dict[str, str]]]: + """The genuine per-node materialized prompt for a trace, keyed by node id. + + Rebuilds the trie graph + pool exactly as the tool does, so a synthetic raw + export carrying these messages is what a FAITHFUL run would have exported. + """ + from aiperf.dataset.graph.adapters.weka.trace_models import WekaTrace + from aiperf.dataset.graph.adapters.weka.trie_build import build_trie_graph + from aiperf.dataset.graph.models import LlmNode + + trace = WekaTrace.model_validate(json.loads(trace_file.read_text())) + parsed, pool = build_trie_graph( + trace, tokenizer_name="gpt2", prompt_corpus="coding", root_seed=None + ) + out: dict[str, list[dict[str, str]]] = {} + for nid, node in parsed.graph.nodes.items(): + if isinstance(node, LlmNode): + out[nid] = pool.materialize(node.metadata["trie"]["prompt_segment_ids"]) + return out + + +def _recorded_timing(trace_file: Path) -> dict[str, tuple[float, float]]: + """``{node_id: (start_s, end_s)}`` of RAW recorded ``request.t`` / ``+api_time``. + + Reads the raw request fields directly: ``_flatten_requests`` here does NOT + apply the idle-gap warp, so ``TrieNode.start`` / ``.end`` would return the 0.0 + ``warped_start`` default rather than the recorded timeline. The faithful + export's edge delay must be the genuine recorded end-to-start gap. + """ + from aiperf.dataset.graph.adapters.weka.trace_models import WekaTrace + from aiperf.dataset.graph.adapters.weka.trie_build import _flatten_requests + + trace = WekaTrace.model_validate(json.loads(trace_file.read_text())) + return { + n.node_id: (n.request.t, n.request.t + (n.request.api_time or 0.0)) + for n in _flatten_requests(trace.requests, root_scope=trace.id) + } + + +def _record( + *, + node_id: str, + phase: str, + request_start_ns: int, + messages: list[dict[str, str]], + conversation_id: str = "trace_sub_n2s1#0.0", +) -> dict: + """A single raw-export JSONL record shaped like the production writer emits. + + ``x_request_id`` folds the node id ahead of its ``::`` nonce, matching the + worker's graph mint -- the node-identity channel the tool reads. + """ + return { + "metadata": { + "conversation_id": conversation_id, + "x_request_id": f"{node_id}::deadbeefdeadbeefdeadbeefdeadbeef", + "benchmark_phase": phase, + "request_start_ns": request_start_ns, + "credit_issued_ns": request_start_ns - 1_000_000, + "session_num": 0, + "turn_index": 0, + }, + "payload": {"messages": messages, "model": _MODEL, "stream": False}, + } + + +def _write_jsonl(path: Path, records: list[dict]) -> Path: + path.write_bytes(b"\n".join(orjson.dumps(r) for r in records) + b"\n") + return path + + +def _faithful_records(trace_file: Path) -> list[dict]: + """A synthetic profiling export that a FAITHFUL trie run would have produced. + + Profiles the two leaf nodes ``agent_001:1`` and ``trace_sub_n2s1:1`` (the chopped + survivors in the real run), each with its genuine materialized prompt and a + dispatch time whose start-to-start gap equals the recorded END-to-start edge + delay (the zero-latency mock collapses ``api_time`` to ~0, so observed + start-to-start == recorded end-to-start). + """ + expected = _build_expected(trace_file) + timing = _recorded_timing(trace_file) + t0 = 1_000_000_000_000 + # Recorded end-to-start edge delay agent_001:1 -> trace_sub_n2s1:1 == + # trace_sub_n2s1:1.start - agent_001:1.end. + edge_delay_s = timing["trace_sub_n2s1:1"][0] - timing["agent_001:1"][1] + ns = 1_000_000_000 + return [ + _record( + node_id="agent_001:1", + phase="profiling", + request_start_ns=t0, + messages=expected["agent_001:1"], + ), + _record( + node_id="trace_sub_n2s1:1", + phase="profiling", + request_start_ns=t0 + int(edge_delay_s * ns), + messages=expected["trace_sub_n2s1:1"], + ), + ] + + +# --- component lane: content_vs_real_trace -------------------------------- + + +@pytest.mark.component_integration +def test_content_vs_real_trace_passes_on_faithful_export(tmp_path: Path) -> None: + # The synthetic export was built with gpt2 (``_build_expected``); the tool + # defaults to the live-run builtin tokenizer, so the run's knob is passed + # explicitly. + raw = _write_jsonl(tmp_path / "raw.jsonl", _faithful_records(_FIX)) + report = content_vs_real_trace(raw, _FIX, tokenizer_name="gpt2") + assert report.passed, report.render() + assert report.checked == 2 + + +@pytest.mark.component_integration +def test_content_vs_real_trace_fails_on_corrupted_export(tmp_path: Path) -> None: + """A single mutated byte in one prompt message must FAIL the content check.""" + records = _faithful_records(_FIX) + # Corrupt trace_sub_n2s1:1's first user message content. + records[1]["payload"]["messages"][0]["content"] += "_CORRUPTED" + raw = _write_jsonl(tmp_path / "raw.jsonl", records) + report = content_vs_real_trace(raw, _FIX, tokenizer_name="gpt2") + assert not report.passed, report.render() + assert any("node=trace_sub_n2s1:1" in m.where for m in report.mismatches) + # Only the corrupted record fails; the untouched one still passes. + assert report.checked == 2 + assert report.passes == 1 + + +@pytest.mark.component_integration +def test_content_vs_real_trace_strips_rid_marker(tmp_path: Path) -> None: + """A ``[rid:...]`` cache-bust prefix on the first user msg is stripped, not failed.""" + records = _faithful_records(_FIX) + first_user = next( + m for m in records[0]["payload"]["messages"] if m["role"] == "user" + ) + first_user["content"] = "[rid:0123456789ab]\n\n" + first_user["content"] + raw = _write_jsonl(tmp_path / "raw.jsonl", records) + report = content_vs_real_trace(raw, _FIX, tokenizer_name="gpt2") + assert report.passed, report.render() + + +# --- component lane: causality_timing_vs_real_trace ----------------------- + + +@pytest.mark.component_integration +def test_causality_timing_passes_on_faithful_export(tmp_path: Path) -> None: + raw = _write_jsonl(tmp_path / "raw.jsonl", _faithful_records(_FIX)) + report = causality_timing_vs_real_trace(raw, _FIX) + assert report.passed, report.render() + assert report.checked == 2 + + +@pytest.mark.component_integration +def test_causality_timing_fails_on_reordered_dispatch(tmp_path: Path) -> None: + """Dispatching trace_sub_n2s1:1 BEFORE its recorded predecessor agent_001:1 is a + causal-order fail.""" + records = _faithful_records(_FIX) + # Swap the dispatch times so trace_sub_n2s1:1 fires 3s BEFORE agent_001:1 + # (predecessor inversion). + r11_ns = records[0]["metadata"]["request_start_ns"] + records[1]["metadata"]["request_start_ns"] = r11_ns - 3_000_000_000 + raw = _write_jsonl(tmp_path / "raw.jsonl", records) + report = causality_timing_vs_real_trace(raw, _FIX) + assert not report.passed, report.render() + assert any("causal-order" in m.detail for m in report.mismatches) + + +@pytest.mark.component_integration +def test_causality_timing_fails_on_wrong_relative_offset(tmp_path: Path) -> None: + """An inter-request gap far from the recorded edge delay must FAIL timing.""" + records = _faithful_records(_FIX) + # Push trace_sub_n2s1:1 10s LATER than its recorded edge delay prescribes + # (well past tol). + records[1]["metadata"]["request_start_ns"] += 10_000_000_000 + raw = _write_jsonl(tmp_path / "raw.jsonl", records) + report = causality_timing_vs_real_trace(raw, _FIX) + assert not report.passed, report.render() + assert any("relative offset" in m.detail for m in report.mismatches) + + +# --- component lane: executor max-gate AND-join + START-root origin ------- +# +# These drive the timing model directly with hand-built ``_RecordedTrace`` / +# ``_ExportRecord`` objects so the executor's firing-gate semantics are exercised +# in isolation (no trie rebuild): the dispatch of a join node is the MAX over its +# profiled preds of ``pred_observed_dispatch + warped_delay`` (the binding pred is +# the argmax, not the recorded-latest pred), and a START-rooted node fires at +# ``origin + min_start_delay`` -- never relative to another root's END. + +_NS = 1_000_000_000 +_MSG = [{"role": "user", "content": "hi"}] + + +def _rnode( + node_id: str, + *, + preds: dict[str, float] | None = None, + min_start_delay_us: float | None = None, + raw_start_s: float = 0.0, + raw_end_s: float = 0.0, +) -> _RecordedNode: + """A timing-only recorded node: ``preds`` maps each pred id to its WARPED edge + delay (us); ``min_start_delay_us`` set => START-rooted. ``raw_start_s`` / + ``raw_end_s`` are the unwarped recorded instants used ONLY to classify the + binding edge exact-vs-idle-capped. Content fields are inert (timing ignores).""" + pred_ids = list(preds or {}) + return _RecordedNode( + node_id=node_id, + messages=_MSG, + start_s=0.0, + end_s=0.0, + raw_start_s=raw_start_s, + raw_end_s=raw_end_s, + predecessors=pred_ids, + pred_delay_us=dict(preds or {}), + rooted_at_start=min_start_delay_us is not None, + min_start_delay_us=min_start_delay_us, + ) + + +def _erec(node_id: str, request_start_ns: int) -> _ExportRecord: + """One profiling export projection at an absolute dispatch instant (ns).""" + return _ExportRecord( + conversation_id="t#0.0", + node_id=node_id, + phase="profiling", + request_start_ns=request_start_ns, + credit_issued_ns=request_start_ns - 1_000_000, + messages=_MSG, + ) + + +@pytest.mark.component_integration +def test_timing_multipred_binds_to_max_gate_not_latest_recorded() -> None: + """A join's expected dispatch is the MAX gate over preds; the binding pred is + the argmax (here B's 60s-capped edge), NOT the recorded-latest pred A.""" + # A dispatches at origin+0, B at origin+10s. A's warped edge into J is 1s + # (small), B's is 60s (idle-capped). Executor gate(J) = max(A+1, B+60) = B+60. + origin = 1_000_000_000_000 + recorded = _RecordedTrace( + trace_id="t", + nodes={ + "A": _rnode("A", min_start_delay_us=0.0), + "B": _rnode("B", min_start_delay_us=10_000_000.0, raw_end_s=0.0), + # J's raw gap from B (100s) exceeds the 60s warped edge -> idle-capped. + "J": _rnode( + "J", preds={"A": 1_000_000.0, "B": 60_000_000.0}, raw_start_s=100.0 + ), + }, + ) + a_ns = origin + b_ns = origin + 10 * _NS + j_ns = b_ns + 60 * _NS # binds to B's capped gate, not A's 1s gate + records = [_erec("A", a_ns), _erec("B", b_ns), _erec("J", j_ns)] + report = causality_timing_vs_real_trace( + None, None, recorded=recorded, records=records + ) + assert report.passed, report.render() + # B is the binding argmax and its warped (60s) < raw -> classified idle-capped. + assert report.idle_capped_edges == 1, report.render() + + +@pytest.mark.component_integration +def test_timing_multipred_too_early_at_other_preds_gate_fails() -> None: + """Moving the join to A's (non-binding) gate dispatches it too early vs the + executor's max gate -> FAIL (non-vacuous).""" + origin = 1_000_000_000_000 + recorded = _RecordedTrace( + trace_id="t", + nodes={ + "A": _rnode("A", min_start_delay_us=0.0), + "B": _rnode("B", min_start_delay_us=10_000_000.0), + "J": _rnode("J", preds={"A": 1_000_000.0, "B": 60_000_000.0}), + }, + ) + a_ns = origin + b_ns = origin + 10 * _NS + j_ns = a_ns + 1 * _NS # A's gate: ~69s BEFORE the binding B+60s gate + records = [_erec("A", a_ns), _erec("B", b_ns), _erec("J", j_ns)] + report = causality_timing_vs_real_trace( + None, None, recorded=recorded, records=records + ) + assert not report.passed, report.render() + assert any("node=J" in m.where for m in report.mismatches), report.render() + + +@pytest.mark.component_integration +def test_timing_and_join_zero_delay_argmax_not_classified_idle_capped() -> None: + """A fan-in whose OBSERVED argmax gate is a NON-BINDING AND-join (delay-0.0) + edge is not classified at all: the delay-0 edge carries no think-time to + compare, even though its raw end-to-start gap is positive. + + Recorded shape: A=[0,10] is the binding pred (warped delay 2s into J), + B=[8,9] is an AND-join wait (delay 0.0 by construction), J raw_start=12. + Observed: B dispatches 8s after A, so gate(B)=B+0 wins the argmax over + gate(A)=A+2s. Before the fix J's raw gap from B (12-9=3s > 0) misread as + "idle-capped" although no recorded gap exceeded the cap. + """ + origin = 1_000_000_000_000 + recorded = _RecordedTrace( + trace_id="t", + nodes={ + "A": _rnode("A", min_start_delay_us=0.0, raw_end_s=10.0), + "B": _rnode( + "B", min_start_delay_us=8_000_000.0, raw_start_s=8.0, raw_end_s=9.0 + ), + "J": _rnode("J", preds={"A": 2_000_000.0, "B": 0.0}, raw_start_s=12.0), + }, + ) + a_ns = origin + b_ns = origin + 8 * _NS + j_ns = b_ns # dispatches at the AND-join gate (B+0), past A's 2s gate + records = [_erec("A", a_ns), _erec("B", b_ns), _erec("J", j_ns)] + report = causality_timing_vs_real_trace( + None, None, recorded=recorded, records=records + ) + assert report.passed, report.render() + assert report.timing_checks == 3, report.render() # J's gate WAS compared + assert report.idle_capped_edges == 0, report.render() + assert report.exact_edges == 0, report.render() + + +@pytest.mark.component_integration +def test_timing_two_start_roots_fire_at_origin_plus_delay() -> None: + """Two independent START-roots dispatch at origin+their own min_start_delay; + the second is NOT timed relative to the first root's END (no spurious negative).""" + origin = 1_000_000_000_000 + recorded = _RecordedTrace( + trace_id="t", + nodes={ + "R0": _rnode("R0", min_start_delay_us=0.0), + "R1": _rnode("R1", min_start_delay_us=2_000_000.0), + }, + ) + records = [ + _erec("R0", origin), + _erec("R1", origin + 2 * _NS), + ] + report = causality_timing_vs_real_trace( + None, None, recorded=recorded, records=records + ) + assert report.passed, report.render() + assert report.checked == 2, report.render() + + +@pytest.mark.component_integration +def test_timing_start_root_shifted_beyond_tolerance_fails() -> None: + """Shifting one START-root's dispatch well past origin+delay FAILS (non-vacuous).""" + origin = 1_000_000_000_000 + recorded = _RecordedTrace( + trace_id="t", + nodes={ + "R0": _rnode("R0", min_start_delay_us=0.0), + "R1": _rnode("R1", min_start_delay_us=2_000_000.0), + }, + ) + records = [ + _erec("R0", origin), + _erec("R1", origin + 2 * _NS + 5 * _NS), # 5s past origin+2s, well past tol + ] + report = causality_timing_vs_real_trace( + None, None, recorded=recorded, records=records + ) + assert not report.passed, report.render() + assert any("node=R1" in m.where for m in report.mismatches), report.render() + + +# --- component lane: content_byte_exact_vs_v04 ---------------------------- + + +@pytest.mark.component_integration +def test_byte_exact_vs_v04_passes_when_identical_after_rid_strip( + tmp_path: Path, +) -> None: + """Identical payloads differing only by the rid marker are byte-exact-equal.""" + ours = _faithful_records(_FIX) + v04 = _faithful_records(_FIX) + # v0.4 carries a DIFFERENT rid marker on the first user turn; ours carries none. + v04_user = next(m for m in v04[0]["payload"]["messages"] if m["role"] == "user") + v04_user["content"] = "[rid:ffffffffffff]\n\n" + v04_user["content"] + ours_raw = _write_jsonl(tmp_path / "ours.jsonl", ours) + v04_raw = _write_jsonl(tmp_path / "v04.jsonl", v04) + report = content_byte_exact_vs_v04(ours_raw, v04_raw) + assert report.passed, report.render() + assert report.checked == 2 + + +@pytest.mark.component_integration +def test_byte_exact_vs_v04_fails_on_content_drift(tmp_path: Path) -> None: + """A non-marker content difference must FAIL the byte-exact check.""" + ours = _faithful_records(_FIX) + v04 = _faithful_records(_FIX) + v04[1]["payload"]["messages"][-1]["content"] += "_DRIFT" + ours_raw = _write_jsonl(tmp_path / "ours.jsonl", ours) + v04_raw = _write_jsonl(tmp_path / "v04.jsonl", v04) + report = content_byte_exact_vs_v04(ours_raw, v04_raw) + assert not report.passed, report.render() + + +@pytest.mark.component_integration +def test_byte_exact_vs_v04_fails_on_missing_coverage(tmp_path: Path) -> None: + """A node present in one export but missing from the other is a coverage fail. + + Guards against a vacuous PASS on an empty key overlap. + """ + ours = _faithful_records(_FIX) + v04 = _faithful_records(_FIX)[:1] # drop trace_sub_n2s1:1 from v0.4 + ours_raw = _write_jsonl(tmp_path / "ours.jsonl", ours) + v04_raw = _write_jsonl(tmp_path / "v04.jsonl", v04) + report = content_byte_exact_vs_v04(ours_raw, v04_raw) + assert not report.passed, report.render() + assert any("missing" in m.detail for m in report.mismatches) + + +# --- cap-aware classification + corpus driver ----------------------------- + +_CAP_S = 5.0 + + +def _linear_trace(trace_id: str, gaps_s: list[float], api_s: float = 1.0) -> dict: + """A linear single-chain Weka trace: each turn extends the prior hash prefix. + + ``gaps_s[i]`` is the RAW recorded END-to-start gap between turn ``i`` and turn + ``i+1`` (so turn ``i+1`` starts at ``end_of_turn_i + gaps_s[i]``). A linear + hash-prefix chain yields one single-predecessor ``StaticEdge`` per hop, the + cleanest input for the exact-vs-idle-capped classification. + """ + t = 0.0 + requests: list[dict] = [] + hash_ids: list[int] = [] + for i, _ in enumerate([0.0, *gaps_s]): + hash_ids = [*hash_ids, i + 1] + requests.append( + { + "t": t, + "type": "n", + "model": _MODEL, + "in": 100, + "out": 10, + "hash_ids": list(hash_ids), + "input_types": ["text"], + "output_types": ["text"], + "stop": "end_turn", + "api_time": api_s, + "think_time": 0.0, + } + ) + if i < len(gaps_s): + t = t + api_s + gaps_s[i] + return { + "id": trace_id, + "models": [_MODEL], + "block_size": 64, + "hash_id_scope": "local", + "requests": requests, + } + + +def _warped_records( + trace_file: Path, conv: str, *, phase: str = "profiling" +) -> list[dict]: + """A faithful capped-run export: observed gaps == the WARPED edge delays. + + Builds the recorded trace WITH the cap, then dispatches each node at + ``predecessor_dispatch + warped_edge_delay`` (accumulated down the chain) so + the observed start-to-start gap of every edge equals its warped + ``delay_after_predecessor_us`` -- exactly what the zero-latency mock produces + (the predecessor returns ~instantly, so observed start-to-start collapses onto + the END-to-start edge delay). Roots dispatch at ``t0``. ``phase`` stamps every + record's ``benchmark_phase`` (``"warmup"`` simulates auto-warmup priming). + """ + recorded = build_recorded_trace(trace_file, idle_gap_cap_seconds=_CAP_S) + t0 = 1_000_000_000_000 + ns = 1_000_000_000 + # Walk nodes in recorded start order so a predecessor's dispatch is assigned + # before its successor reads it (the linear chain is already topologically + # ordered by start_s). + order = sorted(recorded.nodes, key=lambda nid: recorded.nodes[nid].start_s) + dispatch_ns: dict[str, int] = {} + out: list[dict] = [] + for nid in order: + node = recorded.nodes[nid] + preds = [p for p in node.predecessors if p in dispatch_ns] + if preds: + cause = max(preds, key=lambda p: recorded.nodes[p].end_s) + edge_delay_s = node.pred_delay_us.get(cause, 0.0) / 1e6 + start_ns = dispatch_ns[cause] + int(edge_delay_s * ns) + else: + start_ns = t0 + dispatch_ns[nid] = start_ns + out.append( + _record( + node_id=nid, + phase=phase, + request_start_ns=start_ns, + messages=node.messages, + conversation_id=conv, + ) + ) + return out + + +@pytest.mark.component_integration +def test_timing_classifies_exact_and_idle_capped(tmp_path: Path) -> None: + """One sub-cap edge => 'exact'; one over-cap edge => 'idle-capped'; PASSES.""" + # gaps: r_0->r_1 raw 1s (< cap, exact), r_1->r_2 raw 97s (> cap, idle-capped). + trace = _linear_trace("lin_cap", gaps_s=[1.0, 97.0]) + trace_file = tmp_path / "lin_cap.json" + trace_file.write_text(json.dumps(trace)) + raw = _write_jsonl( + tmp_path / "raw.jsonl", _warped_records(trace_file, "lin_cap#0.0") + ) + report = causality_timing_vs_real_trace( + raw, trace_file, idle_gap_cap_seconds=_CAP_S + ) + assert report.passed, report.render() + assert report.exact_edges == 1, report.render() + assert report.idle_capped_edges == 1, report.render() + + +@pytest.mark.component_integration +def test_timing_fails_when_observed_gap_violates_warped(tmp_path: Path) -> None: + """Mutating an observed gap beyond tolerance FAILS timing (non-vacuous proof).""" + trace = _linear_trace("lin_cap", gaps_s=[1.0, 97.0]) + trace_file = tmp_path / "lin_cap.json" + trace_file.write_text(json.dumps(trace)) + records = _warped_records(trace_file, "lin_cap#0.0") + # Push the LAST node 30s past its warped-expected edge delay (well past tol); + # the over-cap edge is bounded by the cap, so 30s is a genuine violation, not + # a tolerated raw think-time. + records[-1]["metadata"]["request_start_ns"] += 30_000_000_000 + raw = _write_jsonl(tmp_path / "raw.jsonl", records) + report = causality_timing_vs_real_trace( + raw, trace_file, idle_gap_cap_seconds=_CAP_S + ) + assert not report.passed, report.render() + assert any("relative offset" in m.detail for m in report.mismatches), ( + report.render() + ) + + +@pytest.mark.component_integration +def test_prove_corpus_aggregates_two_traces(tmp_path: Path) -> None: + """Two-trace corpus: aggregate counts correct; both traces pass; coverage 100%.""" + trace_dir = tmp_path / "traces" + trace_dir.mkdir() + records: list[dict] = [] + for tid in ("traceA", "traceB"): + trace = _linear_trace(tid, gaps_s=[1.0, 97.0]) + tf = trace_dir / f"{tid}.json" + tf.write_text(json.dumps(trace)) + records.extend(_warped_records(tf, f"{tid}#0.0")) + raw = _write_jsonl(tmp_path / "raw.jsonl", records) + rc = prove_corpus(raw, trace_dir, idle_gap_cap_seconds=_CAP_S) + assert rc == 0 + + +@pytest.mark.component_integration +def test_prove_corpus_content_mutation_fails_only_that_trace(tmp_path: Path) -> None: + """A content mutation in ONE trace fails the corpus (and only that trace).""" + trace_dir = tmp_path / "traces" + trace_dir.mkdir() + records: list[dict] = [] + for tid in ("traceA", "traceB"): + trace = _linear_trace(tid, gaps_s=[1.0, 97.0]) + tf = trace_dir / f"{tid}.json" + tf.write_text(json.dumps(trace)) + recs = _warped_records(tf, f"{tid}#0.0") + if tid == "traceB": + recs[-1]["payload"]["messages"][0]["content"] += "_CORRUPTED" + records.extend(recs) + raw = _write_jsonl(tmp_path / "raw.jsonl", records) + rc = prove_corpus(raw, trace_dir, idle_gap_cap_seconds=_CAP_S) + assert rc == 1 + + +@pytest.mark.component_integration +def test_prove_corpus_reports_coverage_for_undispatched_node(tmp_path: Path) -> None: + """An undispatched deep turn is reported as COVERAGE, not a failure.""" + trace = _linear_trace("traceA", gaps_s=[1.0, 97.0]) + trace_dir = tmp_path / "traces" + trace_dir.mkdir() + tf = trace_dir / "traceA.json" + tf.write_text(json.dumps(trace)) + records = _warped_records(tf, "traceA#0.0") + # The bounded run never reached the deepest turn (drop traceA:2's dispatch). Its + # coverage drops below 100% but the proof still PASSES (no checked failure). + records = [ + r for r in records if not r["metadata"]["x_request_id"].startswith("traceA:2::") + ] + raw = _write_jsonl(tmp_path / "raw.jsonl", records) + rc = prove_corpus(raw, trace_dir, idle_gap_cap_seconds=_CAP_S) + assert rc == 0 + + +@pytest.mark.component_integration +def test_prove_corpus_warmup_only_trace_skips_and_passes( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + """A trace only warmup-primed (zero PROFILING records) is SKIP, not FAIL. + + Graph auto-warmup bursts priming credits before profiling, so a bounded + profiling run can leave some traces warmup-only. Both criteria hard-fail on + zero profiling records, so counting any-phase records used to flip such a + trace to FAIL and falsely fail the whole corpus; it must SKIP with a + distinct reason while the fully-profiled trace carries the PASS. + """ + trace_dir = tmp_path / "traces" + trace_dir.mkdir() + records: list[dict] = [] + for tid, phase in (("traceA", "profiling"), ("traceB", "warmup")): + trace = _linear_trace(tid, gaps_s=[1.0, 97.0]) + tf = trace_dir / f"{tid}.json" + tf.write_text(json.dumps(trace)) + records.extend(_warped_records(tf, f"{tid}#0.0", phase=phase)) + raw = _write_jsonl(tmp_path / "raw.jsonl", records) + rc = prove_corpus(raw, trace_dir, idle_gap_cap_seconds=_CAP_S) + out = capsys.readouterr().out + assert rc == 0, out + assert "SKIP (warmup-primed, not profiled)" in out + assert "corpus proof: PASS" in out + assert "FAIL" not in out + + +@pytest.mark.component_integration +def test_prove_corpus_all_warmup_vacuous_exit_nonzero( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + """An export with ONLY warmup records checked nothing -> vacuous, exit 1. + + Every trace is SKIP (warmup-primed), so no per-trace failure fires, but the + corpus-level VACUOUS gate must still fail an export that profiled nothing. + """ + trace_dir, tf, _records = _corpus_with_records(tmp_path) + records = _warped_records(tf, "traceA#0.0", phase="warmup") + raw = _write_jsonl(tmp_path / "raw.jsonl", records) + rc = prove_corpus(raw, trace_dir, idle_gap_cap_seconds=_CAP_S) + out = capsys.readouterr().out + assert rc == 1 + assert "SKIP (warmup-primed, not profiled)" in out + assert "VACUOUS: nothing checked" in out + + +# --- exit-gate integrity + vacuous-proof rejection (T1/T3) ----------------- + + +def _corpus_with_records(tmp_path: Path) -> tuple[Path, Path, list[dict]]: + """A one-trace corpus dir + its faithful export records (not yet written).""" + trace = _linear_trace("traceA", gaps_s=[1.0]) + trace_dir = tmp_path / "traces" + trace_dir.mkdir() + tf = trace_dir / "traceA.json" + tf.write_text(json.dumps(trace)) + return trace_dir, tf, _warped_records(tf, "traceA#0.0") + + +@pytest.mark.component_integration +def test_prove_corpus_unresolvable_node_ids_exit_nonzero(tmp_path: Path) -> None: + """Records whose correlation ids resolve to NO trie node fail the exit code. + + Correlation-scheme drift is exactly what this gate exists to catch; before + the exit-gate fix these rows printed MISMATCH yet exited 0 because + ``dispatched_nodes == 0`` / ``timing.checked == 0`` masked them. + """ + trace_dir, _tf, records = _corpus_with_records(tmp_path) + for r in records: + r["metadata"]["x_request_id"] = "bogus_node:0::deadbeefdeadbeefdeadbeefdeadbeef" + raw = _write_jsonl(tmp_path / "raw.jsonl", records) + rc = prove_corpus(raw, trace_dir, idle_gap_cap_seconds=_CAP_S) + assert rc == 1 + + +@pytest.mark.component_integration +def test_prove_corpus_missing_request_start_exit_nonzero(tmp_path: Path) -> None: + """Records with no ``request_start_ns`` cannot vacuously pass the timing gate.""" + trace_dir, _tf, records = _corpus_with_records(tmp_path) + for r in records: + r["metadata"]["request_start_ns"] = None + r["metadata"]["credit_issued_ns"] = None + raw = _write_jsonl(tmp_path / "raw.jsonl", records) + rc = prove_corpus(raw, trace_dir, idle_gap_cap_seconds=_CAP_S) + assert rc == 1 + + +@pytest.mark.component_integration +def test_prove_corpus_vacuous_empty_inputs_exit_nonzero( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + """An empty export against an empty trace dir checked nothing -> exit 1.""" + trace_dir = tmp_path / "traces" + trace_dir.mkdir() + raw = tmp_path / "raw.jsonl" + raw.write_text("") + rc = prove_corpus(raw, trace_dir, idle_gap_cap_seconds=_CAP_S) + assert rc == 1 + assert "VACUOUS: nothing checked" in capsys.readouterr().out + + +@pytest.mark.component_integration +def test_prove_corpus_zero_overlap_vacuous_exit_nonzero( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + """An export whose records map to NO recorded trace checked nothing -> exit 1. + + The wrong-artifact-paired-with-wrong-corpus case: every trace row is 0% + coverage, so a proof that asserted nothing must not exit green. + """ + trace_dir, _tf, _records = _corpus_with_records(tmp_path) + other = [ + _record( + node_id="r_0", + phase="profiling", + request_start_ns=1_000_000_000_000, + messages=[{"role": "user", "content": "hi"}], + conversation_id="otherTrace#0.0", + ) + ] + raw = _write_jsonl(tmp_path / "raw.jsonl", other) + rc = prove_corpus(raw, trace_dir, idle_gap_cap_seconds=_CAP_S) + assert rc == 1 + assert "VACUOUS: nothing checked" in capsys.readouterr().out + + +# --- summary arithmetic + order/timing separation (T2) --------------------- + + +@pytest.mark.component_integration +def test_timing_report_arithmetic_no_negative_passes() -> None: + """One record failing BOTH causal-order and relative-timing appends two + mismatches for ONE checked record; pass counts come from the dedicated + ``passes`` counter (the old ``checked - len(mismatches)`` went negative), + and the causal-order vs relative-timing tallies stay separated.""" + origin = 1_000_000_000_000 + recorded = _RecordedTrace( + trace_id="t", + nodes={ + "A": _rnode("A", min_start_delay_us=0.0), + "B": _rnode("B", preds={"A": 1_000_000.0}), + }, + ) + records = [ + _erec("A", origin), + _erec("B", origin - 10 * _NS), # 10s BEFORE its recorded predecessor + ] + report = causality_timing_vs_real_trace( + None, None, recorded=recorded, records=records + ) + assert not report.passed + assert report.checked == 2 + assert report.passes == 1 # A passes; B fails (order + timing) + assert len(report.mismatches) == 2 + assert report.passes == report.checked - 1 # never negative, one fail-record + assert (report.order_checks, report.order_failures) == (1, 1) + assert (report.timing_checks, report.timing_failures) == (2, 1) + + +# Content-knob threading (--tokenizer/--corpus/--seed) is proven in +# tests/unit/graph/test_weka_fidelity_tool_gates.py: this lane's FakeTokenizer +# patch flattens tokenizer/seed-dependent content, so the knobs are only +# distinguishable against the real builtin synthesizer the unit lane uses. + + +# --- component lane: end-to-end cache-safety invariants ------------------- +# +# The invariant the interval-order + message-unit redesign exists for, proven +# through the REAL gpt2 tokenizer + real ``CorpusContentSynthesizer`` (no stub +# callbacks): any two requests sharing a block-aligned prefix render an +# IDENTICAL leading per-message pool-id chain (role + message boundaries frozen +# at block creation, inherited immutably, never relabeled/coalesced). The unit +# tests in ``tests/unit/graph/test_weka_trie_interval_order.py`` cover this +# only with collision-free stub decoders; these drive the production path. + +from aiperf.dataset.graph.adapters.weka.trace_models import ( # noqa: E402 + WekaTrace as _WekaTrace, +) +from aiperf.dataset.graph.adapters.weka.trie_build import ( # noqa: E402 + build_trie_graph as _build_trie_graph, +) + + +def _n_req( + t: float, + *, + in_tokens: int, + out_tokens: int, + hash_ids: list[int], + api_time: float = 1.0, + think_time: float = 0.0, +) -> dict: + """A single recorded normal ("n") Weka request dict (``_linear_trace`` shape).""" + return { + "t": t, + "type": "n", + "model": _MODEL, + "in": in_tokens, + "out": out_tokens, + "hash_ids": list(hash_ids), + "input_types": ["text"], + "output_types": ["text"], + "stop": "end_turn", + "api_time": api_time, + "think_time": think_time, + } + + +def _subagent(t: float, agent_id: str, requests: list[dict]) -> dict: + """A completed (blocking) subagent marker wrapping inner recorded requests.""" + return { + "t": t, + "type": "subagent", + "agent_id": agent_id, + "subagent_type": "X", + "status": "completed", + "models": [_MODEL], + "requests": requests, + } + + +def _trace_dict(requests: list[dict], *, block_size: int = 64) -> dict: + """A trace-level Weka dict (``_linear_trace`` field shape) wrapping ``requests``.""" + return { + "id": "invariant", + "models": [_MODEL], + "block_size": block_size, + "hash_id_scope": "local", + "requests": requests, + } + + +def _build_real(requests: list[dict], *, block_size: int = 64): + """Validate + build the trie graph through the REAL gpt2 synthesizer (no stubs).""" + trace = _WekaTrace.model_validate(_trace_dict(requests, block_size=block_size)) + return _build_trie_graph( + trace, + tokenizer_name="gpt2", + prompt_corpus="coding", + root_seed=None, + idle_gap_cap_seconds=None, + ) + + +def _llm_nodes(parsed) -> dict: + """``{node_id: LlmNode}`` for every LlmNode in the built graph.""" + from aiperf.dataset.graph.models import LlmNode + + return {nid: n for nid, n in parsed.graph.nodes.items() if isinstance(n, LlmNode)} + + +def _hash_ids_by_node( + requests: list[dict], *, block_size: int = 64 +) -> dict[str, list[int]]: + """``{node_id: recorded hash_ids}`` recovered from the flatten pass. + + The built ``LlmNode`` no longer carries its recorded ``hash_ids`` in + ``metadata["trie"]`` (the envelope is slimmed to ``prompt_segment_ids`` only), + so tests that identify a node by its recorded hash prefix recover it from the + same ``_flatten_requests`` pass the builder uses -- byte-identical node ids and + the genuine recorded prefix, no reliance on a build-plane metadata copy. + """ + from aiperf.dataset.graph.adapters.weka.trie_build import _flatten_requests + + trace = _WekaTrace.model_validate(_trace_dict(requests, block_size=block_size)) + return { + n.node_id: list(n.request.hash_ids) + for n in _flatten_requests(trace.requests, root_scope=trace.id) + } + + +def _block_prefix_len(a: list[int], b: list[int]) -> int: + """Length of the longest common leading run of two hash-id lists.""" + n = 0 + for x, y in zip(a, b, strict=False): + if x != y: + break + n += 1 + return n + + +def _msg_block_counts_from_tags(block_tags: list[tuple[str, bool]]) -> list[int]: + """Block count per message from a node's frozen per-block tags. + + Mirrors ``assemble_messages`` grouping exactly: a new group opens at the + first block and at every ``starts_new_message`` block. + """ + groups: list[list[int]] = [] + for j, (_role, starts) in enumerate(block_tags): + if starts or not groups: + groups.append([j]) + else: + groups[-1].append(j) + return [len(g) for g in groups] + + +def _tags_by_node( + requests: list[dict], block_size: int +) -> dict[str, list[tuple[str, bool]]]: + """Frozen per-block tags keyed by node id for a trace's requests. + + Reuses the production tag pass so the test's notion of message->block spans + is byte-identical to what the builder froze. + """ + from aiperf.dataset.graph.adapters.weka.trie_build import ( + _flatten_requests, + ) + from aiperf.dataset.graph.segment_ir.trie_content import ( + assign_block_tags, + compute_asst_caps, + resolve_content_parents, + ) + + trace = _WekaTrace.model_validate(_trace_dict(requests, block_size=block_size)) + nodes = _flatten_requests(trace.requests, root_scope=trace.id) + resolve_content_parents(nodes) + caps = compute_asst_caps(nodes, block_size) + return assign_block_tags(nodes, block_size, caps) + + +def _messages_within_blocks(msg_block_counts: list[int], n_blocks: int) -> int: + """Number of leading messages whose blocks lie WHOLLY within the first ``n_blocks``. + + A message spanning a block past ``n_blocks`` is excluded (it is only partly + shared and its id may legitimately differ). + """ + covered = 0 + count = 0 + for bc in msg_block_counts: + if covered + bc > n_blocks: + break + covered += bc + count += 1 + return count + + +@pytest.mark.component_integration +def test_shared_block_prefix_identical_leading_prompt_ids_real_tokenizer() -> None: + """Core invariant (real gpt2): any two nodes sharing ``L>0`` leading blocks + render byte-IDENTICAL pool ids for every message wholly inside those ``L`` + blocks. A single divergent shared-prefix id would trip the equality. + + Fixture (block_size=64, every ``in`` block-aligned): + * ``invariant:0`` root user turn, hash_ids ``[1,2,3]`` (``in=64*3``), big ``out``. + * ``a:0`` a completed subagent forking off the root: inherits the FULL + ``[1,2,3]`` then extends ``[4,5]`` (``in=64*5``) -> new blocks become + assistant (parent had user + out>0). + * ``invariant:1`` a sibling top-level turn sharing the full ``[1,2,3]`` then + diverging at block 3 (``[1,2,3,90]``, ``in=64*4``). + + All three share ``L=3`` blocks; the root's single leading user message + (blocks 0..2, wholly inside L) is byte-identical across all three, and the + shorter node's fully-in-``L`` ids are a prefix of the longer node's. + """ + block_size = 64 + requests = [ + _n_req(0.0, in_tokens=64 * 3, out_tokens=64 * 4, hash_ids=[1, 2, 3]), + _subagent( + 1.0, + "a", + [_n_req(1.1, in_tokens=64 * 5, out_tokens=0, hash_ids=[1, 2, 3, 4, 5])], + ), + _n_req(3.0, in_tokens=64 * 4, out_tokens=0, hash_ids=[1, 2, 3, 90]), + ] + parsed, _pool = _build_real(requests, block_size=block_size) + nodes = _llm_nodes(parsed) + tags = _tags_by_node(requests, block_size) + + ids = {nid: n.metadata["trie"]["prompt_segment_ids"] for nid, n in nodes.items()} + hashes = _hash_ids_by_node(requests, block_size=block_size) + spans = {nid: _msg_block_counts_from_tags(tags[nid]) for nid in nodes} + + shared_pairs = 0 + for u in nodes: + for v in nodes: + if u == v: + continue + length = _block_prefix_len(hashes[u], hashes[v]) + if length <= 0: + continue + shared_pairs += 1 + # Messages of u/v that live WHOLLY inside the first L shared blocks + # must be byte-identical ids (frozen-tag inheritance). The count of + # such messages matches for both because the shared L blocks are + # tagged identically -> identical grouping -> identical span prefix. + mu = _messages_within_blocks(spans[u], length) + mv = _messages_within_blocks(spans[v], length) + k = min(mu, mv) + assert k >= 1, (u, v, length, spans[u], spans[v]) + assert ids[u][:k] == ids[v][:k], ( + f"{u} vs {v}: shared-prefix ids diverge within L={length} blocks: " + f"{ids[u][:k]} != {ids[v][:k]}" + ) + # Non-vacuity: at least one ordered pair actually shared a block prefix. + assert shared_pairs >= 1 + + +@pytest.mark.component_integration +def test_57f2a77e_no_relabel_shared_assistant_block_frozen_real_tokenizer() -> None: + """``57f2a77e``-shape no-relabel regression (real gpt2). + + The corpus receipt ``57f2a77e2d33...`` was a cache MISS because the per-turn + pass RELABELED a shared block's role on a later turn: a parent turn whose + output made part of the shared prefix ASSISTANT, and a subagent sharing that + exact block prefix but recorded pure-user, disagreed at the role-transition + block. The redesign FREEZES each shared block's role at its first creator and + the subagent inherits it verbatim -> the shared leading id chain is IDENTICAL. + + Fixture (block_size=64): + * ``invariant:0`` root user turn ``[1,2,3]`` with big ``out`` (5 blocks) so its + child's new blocks are attributed to the assistant. + * ``invariant:1`` a follow-up sharing ``[1,2,3]`` and extending to ``[1..8]`` + (``in=64*8``); its new blocks 3..7 split assistant-then-user, so its + 8-block prefix materializes as user / assistant / user (a role transition + INSIDE the shared span). + * ``a:0`` a completed subagent sharing the SAME 8-block prefix + (``[1..8, 9]``) recorded pure-user (``out=0``). + + Observed role layout of the shared 8-block prefix (both invariant:1 and a:0): + ``["user", "assistant", "user"]`` (3 messages). Asserts a:0's leading 3 ids + == invariant:1's full 3 ids -> the assistant block was NOT relabeled to user on + the subagent. A relabel would change a:0's middle message role -> different + content-addressed id -> assertion fails. + """ + block_size = 64 + parent_prefix = list(range(1, 9)) # 8 shared blocks + requests = [ + _n_req(0.0, in_tokens=64 * 3, out_tokens=64 * 5, hash_ids=[1, 2, 3]), + _n_req(2.0, in_tokens=64 * 8, out_tokens=64 * 2, hash_ids=parent_prefix), + _subagent( + 4.0, + "a", + [_n_req(4.1, in_tokens=64 * 9, out_tokens=0, hash_ids=[*parent_prefix, 9])], + ), + ] + parsed, pool = _build_real(requests, block_size=block_size) + nodes = _llm_nodes(parsed) + hashes = _hash_ids_by_node(requests, block_size=block_size) + + def by_hash(h: list[int]) -> str: + return next(nid for nid, hids in hashes.items() if hids == h) + + r1 = by_hash(parent_prefix) + sub = by_hash([*parent_prefix, 9]) + r1_ids = nodes[r1].metadata["trie"]["prompt_segment_ids"] + sub_ids = nodes[sub].metadata["trie"]["prompt_segment_ids"] + + r1_roles = [pool.materialize([i])[0]["role"] for i in r1_ids] + # Non-vacuity: the parent's shared 8-block span genuinely contains an + # assistant message (so the test proves inheritance ACROSS a role transition, + # not a trivially-all-user case). + assert "assistant" in r1_roles, r1_roles + assert len(r1_ids) >= 1 + + # The subagent shares exactly r_1's 8 blocks -> its leading id chain covering + # those blocks is byte-identical to r_1's full prompt-id chain (frozen tags). + assert sub_ids[: len(r1_ids)] == r1_ids, ( + f"57f2a77e relabel regression: subagent shared-prefix ids diverge from " + f"parent: {sub_ids[: len(r1_ids)]} != {r1_ids} (roles={r1_roles})" + ) + + +@pytest.mark.component_integration +def test_block_aligned_isl_covered_count_real_tokenizer() -> None: + """Block-aligned ISL through the real gpt2 tokenizer. + + Every LlmNode's frozen covered-block count equals + ``min(len(hash_ids), in // block_size)`` -- the message-unit emitter covers + only that many whole blocks (no partial tail, no synthesis of missing whole + blocks). The build already HARD-ABORTS (``TrieISLMismatchError``) if a node's + materialized token count misses this target, so a successful build is the + load-bearing proof; this test additionally pins the covered-block count so a + DROPPED block (over/under-coverage) would trip the equality. + + Non-vacuity: ``[1,2,3,4]`` records FEWER hash blocks than ``in // block_size`` + (a truncated hash list: ``in=64*5`` but only 4 blocks). The covered count must + use the ``min`` (4), not over-demand 5 -- an over-demand would fail here. + """ + block_size = 64 + requests = [ + _n_req(0.0, in_tokens=64 * 2, out_tokens=64, hash_ids=[1, 2]), + _n_req(2.0, in_tokens=64 * 3, out_tokens=0, hash_ids=[1, 2, 3]), + # Truncated: in // bs == 5 but only 4 recorded hash blocks -> covered = 4. + _n_req(4.0, in_tokens=64 * 5, out_tokens=0, hash_ids=[1, 2, 3, 4]), + ] + parsed, _pool = _build_real(requests, block_size=block_size) + nodes = _llm_nodes(parsed) + tags = _tags_by_node(requests, block_size) + + hashes = _hash_ids_by_node(requests, block_size=block_size) + saw_truncated = False + for nid in nodes: + hash_ids = hashes[nid] + in_blocks = _recorded_in_blocks(requests, hash_ids) + # The frozen tag count IS the covered-block count (the build asserts + # len(tags) * bs == covered-count * bs before emitting the node). + covered = len(tags[nid]) + expected_cover = min(len(hash_ids), in_blocks) + assert covered == expected_cover, (nid, covered, expected_cover, hash_ids) + if in_blocks > len(hash_ids): + saw_truncated = True + assert covered == len(hash_ids) + assert saw_truncated, "fixture must include a truncated-hash node" + + +def _recorded_in_blocks(requests: list[dict], hash_ids: list[int]) -> int: + """``in // 64`` of the (possibly subagent-nested) request whose ``hash_ids`` match. + + Recovers the raw ``in`` block count for a node so a test can distinguish + covered vs truncated coverage; block_size is 64 in these fixtures. + """ + + def _find(reqs: list[dict]) -> int | None: + for r in reqs: + if r["type"] == "subagent": + inner = _find(r["requests"]) + if inner is not None: + return inner + elif r["hash_ids"] == hash_ids: + return r["in"] // 64 + return None + + got = _find(requests) + assert got is not None + return got + + +@pytest.mark.component_integration +def test_two_user_turn_boundary_preserved_real_tokenizer() -> None: + """Boundary preservation end-to-end (real gpt2). + + A single request spanning two consecutive user turns (block 1 opens a new + message, ``starts_new_message=True``) must materialize as >=2 messages (NOT + coalesced into one) and the LAST materialized message role is ``"user"`` + (trailing-user frozen at creation). Coalescing the boundary would drop to 1 + message and trip the count; a relabeled trailing block would trip the role. + """ + block_size = 64 + requests = [ + _n_req(0.0, in_tokens=64, out_tokens=0, hash_ids=[1]), + _n_req(2.0, in_tokens=64 * 2, out_tokens=0, hash_ids=[1, 2]), + ] + parsed, pool = _build_real(requests, block_size=block_size) + nodes = _llm_nodes(parsed) + + hashes = _hash_ids_by_node(requests, block_size=block_size) + two_turn_id = next(nid for nid, hids in hashes.items() if hids == [1, 2]) + ids = nodes[two_turn_id].metadata["trie"]["prompt_segment_ids"] + roles = [pool.materialize([i])[0]["role"] for i in ids] + assert len(ids) >= 2, roles # boundary preserved, not coalesced + assert roles[-1] == "user", roles # trailing-user frozen at creation diff --git a/tests/component_integration/graph/test_weka_trie_dispatch.py b/tests/component_integration/graph/test_weka_trie_dispatch.py new file mode 100644 index 0000000000..66ca21f47c --- /dev/null +++ b/tests/component_integration/graph/test_weka_trie_dispatch.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""T6.6 e2e: the schedule plane DISPATCHES every weka segment-trie ``LlmNode``. + +Drives a real ``GraphIRReplayStrategy.execute_phase`` over a trie ``ParsedGraph`` +built from ``weka_subagent.json`` with +a fake/echo credit issuer (no worker, no ZMQ). Proves the catalog/ordinal +resolution the build plane wrote its unified segment store at is the SAME ordinal +the dispatch adapter resolves a fired node to -- so credits actually flow and the +worker would read the right envelope: + +* (a) EVERY flat trie ``LlmNode`` is dispatched exactly once; +* (b) each dispatched credit's ``node_ordinal`` equals the ordinal + ``trie_node_ordinals`` assigns to that node id (the build<->schedule contract); +* (c) the phase completes with no error and no ``GraphEnvelopeMissing`` / unresolved + (``node_ordinal is None``) dispatch. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import pytest + +pytestmark = [pytest.mark.component_integration, pytest.mark.asyncio] + +_FIX_DIR = Path(__file__).parents[2] / "unit" / "graph" / "fixtures" +_SUBAGENT = _FIX_DIR / "weka_subagent.json" + + +@pytest.fixture(autouse=True) +def _offline_hf(monkeypatch): + """Pin the tokenizer load to the local HuggingFace cache (offline HF).""" + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + + +@dataclass +class _EchoIssuer: + """Fake issuer recording each issued turn; echoes its return next tick. + + Mirrors ``test_graph_ir_overlap_barrier._OrderedEchoIssuer`` (the established + no-worker credit-loop pattern): every ``issue_graph_credit`` records the turn + and schedules the strategy's return observer to fire on the next loop tick, + so the parked dispatch Future resolves exactly as a real worker round-trip + would. The echoed object IS the ``TurnToSend`` -- it carries + ``x_correlation_id`` / ``turn_index`` / ``trace_id`` / ``node_ordinal``, the + only fields the observer + adapter read. + """ + + observer: Any = None + issued_turns: list[Any] = field(default_factory=list) + issued: int = 0 + returned: int = 0 + + async def issue_graph_credit(self, turn: Any) -> bool: + self.issued += 1 + self.issued_turns.append(turn) + asyncio.get_running_loop().call_soon(self._echo, turn) + return True + + def _echo(self, turn: Any) -> None: + self.returned += 1 + if self.observer is not None: + self.observer(turn, None, False) + + def mark_graph_sending_complete(self) -> None: + pass + + def graph_all_returned(self) -> bool: + return self.returned >= self.issued + + def set_graph_all_returned_event(self) -> None: + pass + + +def _parsed_trie(): + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + + parsed = from_weka_trace(str(_SUBAGENT)) + assert parsed.segment_pool is not None, "fixture did not build a trie graph" + return parsed + + +def _make_strategy(parsed, issuer): + from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy + + # Full-replay window (t*=0): assert the WHOLE trie graph dispatches; a + # positive t* would chop pre-t* nodes and the every-node count would not hold. + return GraphIRReplayStrategy( + credit_issuer=issuer, + parsed_graph=parsed, + register_observer=lambda obs: setattr(issuer, "observer", obs), + max_concurrent_traces=8, + start_min_ratio=0.0, + start_max_ratio=0.0, + ) + + +async def test_trie_dispatch_every_node_resolves_build_plane_ordinal(): + """Every trie node dispatches once, each at its ``trie_node_ordinals`` ordinal.""" + from aiperf.dataset.graph.models import LlmNode + from aiperf.dataset.graph.segment_ir.store_builder import trie_node_ordinals + + parsed = _parsed_trie() + llm_nodes = { + nid: n for nid, n in parsed.graph.nodes.items() if isinstance(n, LlmNode) + } + expected_ordinals = trie_node_ordinals(llm_nodes) + assert expected_ordinals, "trie fixture produced no LlmNodes" + + issuer = _EchoIssuer() + strategy = _make_strategy(parsed, issuer) + await strategy.setup_phase() + await asyncio.wait_for(strategy.execute_phase(), timeout=10.0) + + # (c) phase completed cleanly, no errored trace. + assert strategy.completed_traces == 1 + assert strategy.errored_traces == 0 + + # (a) every trie LlmNode dispatched exactly once. + assert issuer.issued == len(llm_nodes) + ordinals_issued = sorted(t.node_ordinal for t in issuer.issued_turns) + assert ordinals_issued == sorted(expected_ordinals.values()) + + # (b) each dispatched credit's ordinal is a real build-plane ordinal (never + # the unresolved None that produces GraphEnvelopeMissing at the worker). + assert all(t.node_ordinal is not None for t in issuer.issued_turns) + assert set(ordinals_issued) == set(expected_ordinals.values()) + + +async def test_trie_dispatch_catalog_matches_build_plane_store(): + """The strategy's catalog is byte-identical to the build plane's ordinal map. + + The build plane's unified trie builders key each node's manifest by + ``trie_node_ordinals``; the schedule plane's ``build_catalog_context`` must + produce the SAME ``{node_id: ordinal}`` so the dispatched credit reads the + right manifest. A drift here is exactly the GraphEnvelopeMissing failure T6.6 + closes. + """ + from aiperf.dataset.graph.graph_path_catalog import build_catalog_context + from aiperf.dataset.graph.models import LlmNode + from aiperf.dataset.graph.segment_ir.store_builder import trie_node_ordinals + + parsed = _parsed_trie() + trace = parsed.traces[0] + llm_nodes = { + nid: n for nid, n in parsed.graph.nodes.items() if isinstance(n, LlmNode) + } + + catalog = build_catalog_context(parsed).catalog[trace.id] + assert catalog == trie_node_ordinals(llm_nodes) diff --git a/tests/component_integration/graph/test_weka_trie_e2e_materialize.py b/tests/component_integration/graph/test_weka_trie_e2e_materialize.py new file mode 100644 index 0000000000..4f52a7918f --- /dev/null +++ b/tests/component_integration/graph/test_weka_trie_e2e_materialize.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end proof the weka segment-trie IR is RUNNABLE. + +The build plane parses a weka trace into a trie ``ParsedGraph`` (every +``LlmNode`` carrying a ``prompt_segment_ids`` pool path) and drains the +:class:`SegmentPool` plus every node's manifest into the ONE interned unified +store (``aiperf_graph_segments_/``) -- the sole trie store shape. + +This test drives that REAL build+persist path through +:meth:`DatasetManager._configure_graph_workload`, then materializes a node the +SAME way the worker does -- via +:func:`aiperf.graph.worker_materialize.materialize_graph_request_unified` +reading the persisted unified store -- and asserts the materialized prompt is +BYTE-EQUAL to ``pool.materialize(node.prompt_segment_ids)``. That is the +deliverable's proof the trie graph runs with correct prompts. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from aiperf.config.flags.cli_config import CLIConfig +from aiperf.dataset.dataset_manager import DatasetManager +from aiperf.dataset.graph.models import LlmNode +from aiperf.dataset.graph.segment_ir.store_builder import trie_node_ordinals +from aiperf.dataset.graph.workload_detect import parse_graph_workload +from aiperf.dataset.graph_segment_unified_store import GraphSegmentUnifiedClient +from aiperf.graph.worker_materialize import materialize_graph_request_unified +from aiperf.plugin.enums import EndpointType +from tests.unit.conftest import make_run_from_cli + +_FIX = Path(__file__).parents[2] / "unit" / "graph" / "fixtures" / "weka_subagent.json" + + +@pytest.fixture +def trie_dataset_manager( + mmap_base_path: Path, # noqa: ARG001 # side-effect: patches MMAP_BASE_PATH +) -> DatasetManager: + """A DatasetManager pointed at the subagent fixture.""" + cli_config = CLIConfig( + model_names=["test-model"], + endpoint_type=EndpointType.CHAT, + streaming=False, + url="http://localhost:8000", + input_file=str(_FIX), + ) + run = make_run_from_cli(cli_config) + return DatasetManager(run=run, service_id="test") + + +@pytest.mark.asyncio +@pytest.mark.component_integration +async def test_trie_build_persists_and_worker_materializes_byte_equal( + trie_dataset_manager: DatasetManager, + mmap_base_path: Path, +) -> None: + """Every trie node's worker-materialized prompt == ``pool.materialize(path)``.""" + dm = trie_dataset_manager + benchmark_id = dm.run.benchmark_id + + # REAL build+persist path: writes the ONE unified store (content pool + + # per-node interned manifests). + convs = await dm._configure_graph_workload(_FIX) + assert convs.trace_ids, "build must yield at least one graph trace" + + # The unified store artifacts exist under the shared base path. + unified_dir = mmap_base_path / f"aiperf_graph_segments_{benchmark_id}" + for name in ("content.blob", "content.idx", "nodes.blob", "nodes.idx"): + assert (unified_dir / name).exists(), f"unified store missing {name}" + + # Re-parse to recover the in-memory trie graph + pool the build persisted + # (deterministic from the same run), so we know each node's expected + # prompt path and the SegmentPool ground truth. + parsed = parse_graph_workload(dm.run, _FIX) + assert parsed.segment_pool is not None, "trie parse must surface a SegmentPool" + pool = parsed.segment_pool + + trace = parsed.traces[0] + llm_nodes = { + nid: n for nid, n in parsed.graph.nodes.items() if isinstance(n, LlmNode) + } + assert llm_nodes, "trie graph must have LlmNodes" + ordinals = trie_node_ordinals(llm_nodes) + + client = GraphSegmentUnifiedClient( + base_path=mmap_base_path, benchmark_id=benchmark_id + ).open() + try: + checked = 0 + for node_id, node in llm_nodes.items(): + path = node.metadata["trie"]["prompt_segment_ids"] + expected = pool.materialize(path) + + req = materialize_graph_request_unified( + client, trace.id, ordinals[node_id], "profiling" + ) + assert req is not None, f"node {node_id!r} has no persisted manifest" + # Byte-equal: the worker walked the persisted unified store along + # the interned handle path and rebuilt the exact pool prompt. + assert req["messages"] == expected, node_id + checked += 1 + assert checked == len(llm_nodes) + finally: + client.close() + + +@pytest.mark.asyncio +@pytest.mark.component_integration +async def test_trie_node_prompt_matches_pool_materialize_for_subagent_turn( + trie_dataset_manager: DatasetManager, + mmap_base_path: Path, +) -> None: + """A multi-segment (deepest-path) node round-trips byte-for-byte through persist.""" + dm = trie_dataset_manager + benchmark_id = dm.run.benchmark_id + await dm._configure_graph_workload(_FIX) + + parsed = parse_graph_workload(dm.run, _FIX) + pool = parsed.segment_pool + trace = parsed.traces[0] + llm_nodes = { + nid: n for nid, n in parsed.graph.nodes.items() if isinstance(n, LlmNode) + } + ordinals = trie_node_ordinals(llm_nodes) + + # The node with the longest prompt path exercises the deepest shared-prefix + # chain -- the hardest case for byte-equality across the persist boundary. + deepest_id = max( + llm_nodes, + key=lambda nid: len(llm_nodes[nid].metadata["trie"]["prompt_segment_ids"]), + ) + path = llm_nodes[deepest_id].metadata["trie"]["prompt_segment_ids"] + assert len(path) >= 1 + + client = GraphSegmentUnifiedClient( + base_path=mmap_base_path, benchmark_id=benchmark_id + ).open() + try: + req = materialize_graph_request_unified( + client, trace.id, ordinals[deepest_id], "profiling" + ) + finally: + client.close() + + assert req is not None + assert req["messages"] == pool.materialize(path) diff --git a/tests/component_integration/graph/test_weka_trie_hf_streaming.py b/tests/component_integration/graph/test_weka_trie_hf_streaming.py new file mode 100644 index 0000000000..83ca4bac42 --- /dev/null +++ b/tests/component_integration/graph/test_weka_trie_hf_streaming.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""HF-streaming trie build end-to-end through DatasetManager (Task T8-hf). + +Drives the REAL :meth:`DatasetManager._configure_graph_workload` HF-streaming +build path -- the ``--public-dataset`` path that streams per row instead of +parsing the whole corpus eagerly -- with the segment-trie IR enabled. The HF +``_load_hf_rows`` is monkeypatched to return synthetic weka rows so the path +runs hermetically (no network), exactly as a real HF corpus would feed it. + +Asserts the streaming build produces the ONE interned unified store +(``aiperf_graph_segments_/``), and that worker +``materialize_graph_request_unified`` over the streamed store reproduces the +SAME prompt as the in-memory pool ground truth (the failure mode this task +closes: the streaming path previously dropped the pool and served empty +prompts). +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from aiperf.config.flags.cli_config import CLIConfig +from aiperf.dataset.dataset_manager import DatasetManager +from aiperf.dataset.graph.models import LlmNode +from aiperf.dataset.graph.segment_ir.store_builder import trie_node_ordinals +from aiperf.dataset.graph.workload_detect import parse_graph_workload +from aiperf.dataset.graph_segment_unified_store import GraphSegmentUnifiedClient +from aiperf.graph.worker_materialize import materialize_graph_request_unified +from aiperf.plugin.enums import EndpointType +from tests.unit.conftest import make_run_from_cli + +# An HF-shaped id carrying the weka marker so ``_looks_like_hf_dataset_id`` +# routes the build through the STREAMING path. It is not a real path on disk. +_HF_ID = "synthetic/cc-traces-weka-streamtest" + +_ROWS = [ + { + "id": "trace_alpha", + "models": ["claude-opus-4-5-20251101"], + "block_size": 64, + "hash_id_scope": "local", + "requests": [ + { + "t": 0.0, + "type": "n", + "model": "claude-opus-4-5-20251101", + "in": 180, + "out": 25, + "hash_ids": [1, 2], + "input_types": ["text"], + "output_types": ["text"], + "stop": "end_turn", + "api_time": 0.8, + "think_time": 0.0, + }, # noqa: E501 + { + "t": 1.0, + "type": "n", + "model": "claude-opus-4-5-20251101", + "in": 240, + "out": 30, + "hash_ids": [1, 2, 3], + "input_types": ["text"], + "output_types": ["text"], + "stop": "end_turn", + "api_time": 0.9, + "think_time": 0.1, + }, # noqa: E501 + ], + }, + { + "id": "trace_beta", + "models": ["claude-opus-4-5-20251101"], + "block_size": 64, + "hash_id_scope": "local", + "requests": [ + { + "t": 0.0, + "type": "n", + "model": "claude-opus-4-5-20251101", + "in": 180, + "out": 25, + "hash_ids": [1, 2], + "input_types": ["text"], + "output_types": ["text"], + "stop": "tool_use", + "api_time": 0.8, + "think_time": 0.0, + }, # noqa: E501 + { + "t": 1.0, + "type": "subagent", + "agent_id": "agent_001", + "subagent_type": "Explore", + "duration_ms": 4000, + "total_tokens": 600, + "tool_use_count": 1, + "status": "completed", + "models": ["claude-opus-4-5-20251101"], + "tool_tokens": 0, + "system_tokens": 0, + "requests": [ + { + "t": 1.2, + "type": "n", + "model": "claude-opus-4-5-20251101", + "in": 200, + "out": 30, + "hash_ids": [10, 11], + "input_types": ["text"], + "output_types": ["text"], + "stop": "end_turn", + "api_time": 0.9, + "think_time": 0.0, + }, # noqa: E501 + ], + }, + { + "t": 6.0, + "type": "n", + "model": "claude-opus-4-5-20251101", + "in": 280, + "out": 45, + "hash_ids": [1, 2, 3, 4], + "input_types": ["text"], + "output_types": ["text"], + "stop": "end_turn", + "api_time": 1.3, + "think_time": 0.5, + }, # noqa: E501 + ], + }, +] + + +@pytest.fixture +def hf_rows(monkeypatch: pytest.MonkeyPatch) -> None: + """Patch the HF row loader to return synthetic weka rows (no network).""" + import aiperf.dataset.graph.adapters.weka.trace as weka_trace + + def _fake_load(repo_id, *, split, revision): # noqa: ANN001, ANN202, ARG001 + for row in _ROWS: + yield dict(row) + + monkeypatch.setattr(weka_trace, "_load_hf_rows", _fake_load) + + +@pytest.fixture +def trie_streaming_dm( + mmap_base_path: Path, # noqa: ARG001 # side-effect: patches MMAP_BASE_PATH + hf_rows: None, # noqa: ARG001 # side-effect: stubs the HF loader +) -> DatasetManager: + """A DatasetManager pointed at the synthetic HF id.""" + cli_config = CLIConfig( + model_names=["test-model"], + endpoint_type=EndpointType.CHAT, + streaming=False, + url="http://localhost:8000", + input_file=_HF_ID, + ) + run = make_run_from_cli(cli_config) + return DatasetManager(run=run, service_id="test") + + +@pytest.mark.asyncio +@pytest.mark.component_integration +async def test_hf_streaming_trie_build_persists_and_materializes( + trie_streaming_dm: DatasetManager, + mmap_base_path: Path, +) -> None: + """HF streaming build writes the unified store; worker prompts == pool truth.""" + dm = trie_streaming_dm + benchmark_id = dm.run.benchmark_id + + convs = await dm._configure_graph_workload(Path(_HF_ID)) + assert convs.trace_ids, "streaming build must yield graph traces" + + # The unified store artifacts MUST exist (the gap this task closes). + unified_dir = mmap_base_path / f"aiperf_graph_segments_{benchmark_id}" + for name in ("content.blob", "content.idx", "nodes.blob", "nodes.idx"): + assert (unified_dir / name).exists(), ( + f"streaming trie build wrote no unified store {name}" + ) + + parsed = parse_graph_workload(dm.run, Path(_HF_ID)) + assert parsed.segment_pool is not None + pool = parsed.segment_pool + + client = GraphSegmentUnifiedClient( + base_path=mmap_base_path, benchmark_id=benchmark_id + ).open() + try: + checked = 0 + for trace in parsed.traces: + llm_nodes = { + nid: n + for nid, n in parsed.graphs[trace.graph_ref].nodes.items() + if isinstance(n, LlmNode) + } + ordinals = trie_node_ordinals(llm_nodes) + for node_id, node in llm_nodes.items(): + path = node.metadata["trie"]["prompt_segment_ids"] + req = materialize_graph_request_unified( + client, trace.id, ordinals[node_id], "profiling" + ) + assert req is not None, f"node {node_id!r} has no persisted manifest" + assert req["messages"], f"node {node_id!r} materialized EMPTY prompt" + assert req["messages"] == pool.materialize(path), node_id + checked += 1 + assert checked > 0 + finally: + client.close() diff --git a/tests/component_integration/graph/test_wrap_guard.py b/tests/component_integration/graph/test_wrap_guard.py new file mode 100644 index 0000000000..1b8e9e981c --- /dev/null +++ b/tests/component_integration/graph/test_wrap_guard.py @@ -0,0 +1,336 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fail-loud wrap-guard for EXPLICIT concurrent over-subscription. + +Component-level: no real worker / ZMQ / mock server. A fake ``CreditIssuer`` +echoes each ``issue_graph_credit`` back to the strategy's installed graph-return +observer, so ``execute_phase`` drives the REAL dispatch/setup path (the guard +fires early in ``execute_phase``, before any lane starts). + +Asserts the contract: resolved ``concurrency`` exceeding the +distinct loaded traces is a HARD configuration error UNLESS dataset wrapping is +allowed (explicit ``--allow-dataset-wrap`` or, unset, cache-bust ON) -- never a +silent clone-to-fill. + +- (a) concurrency 100 over 15 distinct, cache-bust OFF, wrap unset -> RAISES; + message carries the cache-bust suggestion. +- (b) same but cache-bust ON -> wraps (no raise). +- (c) explicit ``allow_dataset_wrap=False`` + cache-bust ON -> RAISES; message + OMITS the cache-bust suggestion (already on) and notes the explicit disable. +- (d) default concurrency 1 over any corpus -> NEVER raises. +- (e) a selection knob (num_dataset_entries) capping the corpus -> the message + phrases the shortfall as CAPPED, not EXHAUSTED. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import msgspec +import pytest + +from aiperf.common.enums import CacheBustTarget, CreditPhase +from aiperf.common.exceptions import ConfigurationError + +pytestmark = [pytest.mark.component_integration, pytest.mark.asyncio] + +_FIX_DIR = Path(__file__).parents[2] / "unit" / "graph" / "fixtures" +_MIN = _FIX_DIR / "weka_min.json" + + +@dataclass +class _EchoIssuer: + """Fake CreditIssuer echoing each issued credit to the return observer.""" + + observer: Any = None + issued: int = 0 + returned: int = 0 + sending_complete_calls: int = 0 + all_returned_event_set: bool = False + sent: list[Any] = field(default_factory=list) + + async def issue_graph_credit(self, turn: Any) -> bool: + self.issued += 1 + self.sent.append(turn) + asyncio.get_running_loop().call_soon(self._echo, turn) + return True + + def _echo(self, turn: Any) -> None: + self.returned += 1 + if self.observer is not None: + self.observer(turn, None, False) + + def mark_graph_sending_complete(self) -> None: + self.sending_complete_calls += 1 + + def graph_all_returned(self) -> bool: + return self.returned >= self.issued + + def set_graph_all_returned_event(self) -> None: + self.all_returned_event_set = True + + +class _PhaseCfg: + """Per-phase config stub carrying only the fields the guard/strategy read. + + ``num_dataset_entries`` / ``max_context_length`` mirror the graph-plane + selection knobs the guard reads off the config (via ``getattr``) to phrase + the shortfall as CAPPED rather than EXHAUSTED. + """ + + def __init__( + self, + *, + concurrency: int | None = None, + num_dataset_entries: int | None = None, + max_context_length: int | None = None, + expected_num_sessions: int | None = None, + ) -> None: + self.phase = CreditPhase.PROFILING + self.concurrency = concurrency + self.expected_num_sessions = expected_num_sessions + self.total_expected_requests = None + self.expected_duration_sec = None + self.num_dataset_entries = num_dataset_entries + self.max_context_length = max_context_length + + +def _corpus(n: int, *, gap_free: bool = False): + """Return a ParsedGraph whose ``traces`` holds ``n`` distinct-id instances. + + Replicates the single ``weka_min`` template with ``#N`` suffixes; every clone + resolves to the base graph (single-graph fallback). ``gap_free`` zeros edge + delays so a dispatching (non-raising) run replays instantly instead of parking + on ``weka_min``'s recorded ~1s idle gaps (real sleeps in component tests). + """ + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + + base = from_weka_trace(str(_MIN)) + if gap_free: + graph = base.graph + zeroed = [ + msgspec.structs.replace( + e, + **{ + f: 0.0 + for f in ("delay_after_predecessor_us", "min_start_delay_us") + if getattr(e, f, None) is not None + }, + ) + for e in graph.edges + ] + base = msgspec.structs.replace( + base, graph=msgspec.structs.replace(graph, edges=zeroed) + ) + t0 = base.traces[0] + clones = [t0] + clones.extend(msgspec.structs.replace(t0, id=f"{t0.id}#{i}") for i in range(1, n)) + return msgspec.structs.replace(base, traces=clones) + + +def _make_strategy(parsed, issuer: _EchoIssuer, **overrides): + from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy + + overrides.setdefault("start_min_ratio", 0.0) + overrides.setdefault("start_max_ratio", 0.0) + return GraphIRReplayStrategy( + config=overrides.pop("config", None), + conversation_source=None, + scheduler=None, + stop_checker=None, + credit_issuer=issuer, + lifecycle=overrides.pop("lifecycle", None), + parsed_graph=parsed, + register_observer=lambda obs: setattr(issuer, "observer", obs), + **overrides, + ) + + +async def test_oversubscription_cache_bust_off_wrap_unset_raises(): + """(a) concurrency 100 over 15, cache-bust OFF, wrap unset -> RAISES.""" + parsed = _corpus(15) + strategy = _make_strategy( + parsed, + _EchoIssuer(), + max_concurrent_traces=100, + cache_bust=CacheBustTarget.NONE, + allow_dataset_wrap=None, + ) + # The guard raises on the AWAITED setup_phase path (PhaseRunner launches + # execute_phase fire-and-forget, so a raise there would be swallowed). + with pytest.raises(ConfigurationError) as exc: + await strategy.setup_phase() + + msg = str(exc.value) + assert "concurrency 100 exceeds 15 distinct loaded traces" in msg + # EXHAUSTED phrasing (no selection knob was set). + assert "only 15 distinct traces available" in msg + # Cache-bust suggestion present ONLY because cache-bust is OFF. + assert "--cache-bust first_turn_prefix" in msg + assert "clones do not collide on identical prefixes" in msg + + +async def test_oversubscription_cache_bust_on_wraps_no_raise(): + """(b) same over-subscription but cache-bust ON -> wraps (no raise).""" + parsed = _corpus(15, gap_free=True) + issuer = _EchoIssuer() + strategy = _make_strategy( + parsed, + issuer, + max_concurrent_traces=100, + cache_bust=CacheBustTarget.FIRST_TURN_PREFIX, + allow_dataset_wrap=None, + ) + await strategy.setup_phase() + await asyncio.wait_for(strategy.execute_phase(), timeout=15.0) + + # No raise: the run wraps/covers the corpus and completes. + assert strategy.completed_traces == 15 + assert issuer.sending_complete_calls >= 1 + + +async def test_explicit_wrap_false_raises_and_omits_cache_bust_suggestion(): + """(c) explicit allow_dataset_wrap=False + cache-bust ON -> RAISES, no cache-bust hint.""" + parsed = _corpus(15) + strategy = _make_strategy( + parsed, + _EchoIssuer(), + max_concurrent_traces=100, + cache_bust=CacheBustTarget.FIRST_TURN_PREFIX, + allow_dataset_wrap=False, + ) + # The guard raises on the AWAITED setup_phase path (PhaseRunner launches + # execute_phase fire-and-forget, so a raise there would be swallowed). + with pytest.raises(ConfigurationError) as exc: + await strategy.setup_phase() + + msg = str(exc.value) + assert "concurrency 100 exceeds 15 distinct loaded traces" in msg + # The real distinguisher between (a) and (c): cache-bust already ON -> the + # cache-bust suggestion must be OMITTED (no cache-bust mention at all). + assert "--cache-bust" not in msg + assert "clones do not collide" not in msg + # The disabled-wrapping note is present but NEUTRAL -- it must never claim the + # user passed --allow-dataset-wrap false (``GraphDispatchResolver`` derives that + # False from an UNSET flag on the default path). + assert "Dataset wrapping is disabled" in msg + assert "--allow-dataset-wrap false" not in msg + + +async def test_default_concurrency_one_never_raises(): + """(d) default concurrency 1 over a corpus -> never raises (1 <= corpus).""" + parsed = _corpus(3, gap_free=True) + issuer = _EchoIssuer() + strategy = _make_strategy( + parsed, + issuer, + config=_PhaseCfg(concurrency=None), + max_concurrent_traces=None, # resolves to the aiperf default of 1 + cache_bust=CacheBustTarget.NONE, + allow_dataset_wrap=False, + ) + assert strategy._max_concurrent == 1 + await strategy.setup_phase() + await asyncio.wait_for(strategy.execute_phase(), timeout=15.0) + + assert strategy.completed_traces == 3 + + +async def test_session_budget_within_corpus_stands_down(): + """(f) an explicit session budget <= distinct traces -> NO raise even with + concurrency over-provisioned: total instances are bounded below any + cloning need, so concurrency is a mere ceiling (e.g. the common + ``--num-conversations 1 --concurrency 4`` single-pass shape).""" + parsed = _corpus(3, gap_free=True) + issuer = _EchoIssuer() + strategy = _make_strategy( + parsed, + issuer, + config=_PhaseCfg(concurrency=100, expected_num_sessions=3), + max_concurrent_traces=100, + cache_bust=CacheBustTarget.NONE, + allow_dataset_wrap=False, + ) + await strategy.setup_phase() + await asyncio.wait_for(strategy.execute_phase(), timeout=15.0) + + assert strategy.completed_traces == 3 + assert issuer.sending_complete_calls >= 1 + + +async def test_session_budget_beyond_corpus_still_raises(): + """(f') a session budget EXCEEDING the distinct corpus still needs clones, + so the guard keeps failing loud without wrapping.""" + parsed = _corpus(3) + strategy = _make_strategy( + parsed, + _EchoIssuer(), + config=_PhaseCfg(concurrency=100, expected_num_sessions=6), + max_concurrent_traces=100, + cache_bust=CacheBustTarget.NONE, + allow_dataset_wrap=None, + ) + with pytest.raises(ConfigurationError) as exc: + await strategy.setup_phase() + assert "concurrency 100 exceeds 3 distinct loaded traces" in str(exc.value) + + +async def test_capped_corpus_phrases_message_as_capped(): + """(e) a selection knob capping the corpus -> CAPPED phrasing, not EXHAUSTED.""" + parsed = _corpus(15) + strategy = _make_strategy( + parsed, + _EchoIssuer(), + config=_PhaseCfg(concurrency=100, num_dataset_entries=15), + max_concurrent_traces=100, + cache_bust=CacheBustTarget.NONE, + allow_dataset_wrap=None, + ) + # The guard raises on the AWAITED setup_phase path (PhaseRunner launches + # execute_phase fire-and-forget, so a raise there would be swallowed). + with pytest.raises(ConfigurationError) as exc: + await strategy.setup_phase() + + msg = str(exc.value) + assert "capped to 15 by" in msg + assert "--num-dataset-entries" in msg + assert "only 15 distinct traces available" not in msg + + +async def test_from_run_threads_num_dataset_entries_onto_phase_config(): + """from_run seam: --num-dataset-entries N lands on the graph CreditPhaseConfig. + + Closes the three-touch wiring trap the CAPPED branch depends on. The guard + reads ``num_dataset_entries`` straight off ``self._config`` (the + CreditPhaseConfig the runner hands the strategy), so if ``from_run`` does + not thread the resolved cap onto that config the CAPPED phrasing silently + dies and every real run says EXHAUSTED even when the user capped the corpus. + Mirrors ``test_duplication_report.test_cache_bust_seam_reaches_strategy_end_to_end``. + """ + from aiperf.config.flags.cli_config import CLIConfig + from aiperf.plugin.enums import TimingMode + from aiperf.timing.config import TimingConfig + from tests.unit.conftest import make_run_from_cli + + cfg = CLIConfig( + model_names=["test-model"], + input_file=str(_MIN), + request_count=3, + conversation_num_dataset_entries=7, + ) + run = make_run_from_cli(cfg) + + tc = TimingConfig.from_run(run) + profiling = [p for p in tc.phase_configs if p.phase == CreditPhase.PROFILING] + assert profiling, "expected a graph profiling phase" + phase_cfg = profiling[0] + assert phase_cfg.timing_mode == TimingMode.GRAPH_IR + # from_run resolves the explicit --num-dataset-entries cap (the default + # dataset's ``entries``) onto the CreditPhaseConfig, feeding the guard's + # CAPPED branch. + assert phase_cfg.num_dataset_entries == 7 + # No synthesis.max_context_length set -> the sibling cap resolves to None. + assert phase_cfg.max_context_length is None diff --git a/tests/component_integration/timing/test_dag_v1_adversarial.py b/tests/component_integration/timing/test_dag_v1_adversarial.py index 2a58c93458..4617773e09 100644 --- a/tests/component_integration/timing/test_dag_v1_adversarial.py +++ b/tests/component_integration/timing/test_dag_v1_adversarial.py @@ -3,9 +3,9 @@ """Adversarial component-integration tests for DAG prereq gating under validate_for_orchestrator_v1. Covers the full DagJsonlLoader -> DatasetMetadata -> validate_for_orchestrator_v1 -pipeline, plus the two post-fix invariants: -- Task 7 fix: forward / same-turn prereq branch references are rejected. -- Task 8 fix: branches consumed by more than one gated turn are rejected. +pipeline, plus two guard invariants: +- forward / same-turn prereq branch references are rejected. +- branches consumed by more than one gated turn are rejected. """ from __future__ import annotations @@ -190,7 +190,7 @@ def test_two_independent_conversations_validate_separately(tmp_path: Path) -> No def test_forward_prereq_reference_rejected_end_to_end_bug_fix_1() -> None: - """Task 7 fix: a prereq that references a branch declared on a later turn is rejected.""" + """A prereq that references a branch declared on a later turn is rejected.""" md = DatasetMetadata( conversations=[ ConversationMetadata( @@ -223,7 +223,7 @@ def test_forward_prereq_reference_rejected_end_to_end_bug_fix_1() -> None: def test_same_turn_prereq_reference_rejected_end_to_end_bug_fix_1() -> None: - """Task 7 fix: a prereq that references a branch declared on the same turn is rejected.""" + """A prereq that references a branch declared on the same turn is rejected.""" md = DatasetMetadata( conversations=[ ConversationMetadata( diff --git a/tests/fixtures/dag/graph_parity/delayed_turn.dag.jsonl b/tests/fixtures/dag/graph_parity/delayed_turn.dag.jsonl new file mode 100644 index 0000000000..c707d83c2d --- /dev/null +++ b/tests/fixtures/dag/graph_parity/delayed_turn.dag.jsonl @@ -0,0 +1 @@ +{"session_id":"seq","turns":[{"model":"m","messages":[{"role":"user","content":"First turn, no delay."}],"max_tokens":16},{"model":"m","messages":[{"role":"user","content":"Second turn, gated by delay."}],"max_tokens":16,"delay":250}]} diff --git a/tests/fixtures/dag/graph_parity/fork_minimal.dag.jsonl b/tests/fixtures/dag/graph_parity/fork_minimal.dag.jsonl new file mode 100644 index 0000000000..3707ac5d21 --- /dev/null +++ b/tests/fixtures/dag/graph_parity/fork_minimal.dag.jsonl @@ -0,0 +1,3 @@ +{"session_id":"root","turns":[{"model":"Qwen3-0.6B","messages":[{"role":"system","content":"You are a careful assistant."},{"role":"user","content":"Please summarize the attached document."}],"max_tokens":128,"forks":["branch-a","branch-b"]}]} +{"session_id":"branch-a","turns":[{"model":"Qwen3-0.6B","messages":[{"role":"user","content":"Expand on the first section in more detail."},{"role":"user","content":"Add a brief counter-argument."}],"max_tokens":96},{"model":"Qwen3-0.6B","messages":[{"role":"user","content":"Now tighten the expansion."},{"role":"user","content":"Keep the counter-argument intact."}],"max_tokens":64}]} +{"session_id":"branch-b","turns":[{"model":"Qwen3-0.6B","messages":[{"role":"user","content":"Point out weaknesses in the summary."}],"max_tokens":128},{"model":"Qwen3-0.6B","messages":[{"role":"user","content":"Fold the critique into a revised summary."}],"max_tokens":96}]} diff --git a/tests/fixtures/dag/graph_parity/fork_toolcalls.dag.jsonl b/tests/fixtures/dag/graph_parity/fork_toolcalls.dag.jsonl new file mode 100644 index 0000000000..79ab5642bf --- /dev/null +++ b/tests/fixtures/dag/graph_parity/fork_toolcalls.dag.jsonl @@ -0,0 +1,2 @@ +{"session_id":"root","turns":[{"model":"Qwen3-0.6B","messages":[{"role":"system","content":"You are a careful planning assistant."},{"role":"user","content":"Plan the rollout and call the planner [tool]."}],"max_tokens":128,"forks":["branch-a"]}]} +{"session_id":"branch-a","turns":[{"model":"Qwen3-0.6B","messages":[{"role":"user","content":"Just invoke the checker, no prose [tool-only]."}],"max_tokens":96},{"model":"Qwen3-0.6B","messages":[{"role":"user","content":"Now summarize what the tools returned."}],"max_tokens":64}]} diff --git a/tests/fixtures/dag/graph_parity/mixed_full.dag.jsonl b/tests/fixtures/dag/graph_parity/mixed_full.dag.jsonl new file mode 100644 index 0000000000..581460a5ba --- /dev/null +++ b/tests/fixtures/dag/graph_parity/mixed_full.dag.jsonl @@ -0,0 +1,3 @@ +{"session_id":"root","turns":[{"model":"mixed-model","messages":[{"role":"system","content":"You are a planner."},{"role":"user","content":"Draft the migration plan."}],"max_tokens":32,"tools":[{"type":"function","function":{"name":"lookup_docs","description":"Look up internal docs.","parameters":{"type":"object","properties":{"query":{"type":"string"}},"required":["query"]}}}],"extra":{"temperature":0.5,"top_p":0.9,"ignore_eos":true},"spawns":[{"children":["helper"],"join_at":2}]},{"messages":[{"role":"user","content":"Refine the plan for rollback safety."}],"max_tokens":24,"delay":20},{"model":"mixed-model","messages":[{"role":"user","content":"Incorporate the helper results and finalize."}],"max_tokens":16,"forks":["finisher"]}]} +{"session_id":"helper","turns":[{"messages":[{"role":"user","content":"Inventory the affected services."}],"max_tokens":12},{"model":"helper-model","messages":[{"role":"user","content":"Rank them by risk."}],"max_tokens":8,"delay":10}]} +{"session_id":"finisher","turns":[{"model":"finisher-model","messages":[{"role":"user","content":"Write the executive summary."}],"max_tokens":8,"extra":{"temperature":0.1}}]} diff --git a/tests/fixtures/dag/graph_parity/prespawn.dag.jsonl b/tests/fixtures/dag/graph_parity/prespawn.dag.jsonl new file mode 100644 index 0000000000..0dd6148e96 --- /dev/null +++ b/tests/fixtures/dag/graph_parity/prespawn.dag.jsonl @@ -0,0 +1,2 @@ +{"session_id":"root","pre_session_spawns":["bg-child"],"turns":[{"model":"m","messages":[{"role":"user","content":"Owner turn 0 kickoff."}],"max_tokens":16},{"model":"m","messages":[{"role":"user","content":"Owner turn 1 follow-up."}],"max_tokens":16}]} +{"session_id":"bg-child","turns":[{"model":"m","messages":[{"role":"user","content":"Background child work item."}],"max_tokens":16}]} diff --git a/tests/fixtures/dag/graph_parity/spawn_join.dag.jsonl b/tests/fixtures/dag/graph_parity/spawn_join.dag.jsonl new file mode 100644 index 0000000000..b8fc0b11ab --- /dev/null +++ b/tests/fixtures/dag/graph_parity/spawn_join.dag.jsonl @@ -0,0 +1,2 @@ +{"session_id":"parent","turns":[{"model":"spawn-model","messages":[{"role":"system","content":"You coordinate helpers."},{"role":"user","content":"Kick off the analysis."}],"max_tokens":32,"spawns":[{"children":["worker-a"],"join_at":2}]},{"model":"spawn-model","messages":[{"role":"user","content":"Meanwhile, outline the report."}],"max_tokens":24},{"model":"spawn-model","messages":[{"role":"user","content":"Fold in the helper findings."}],"max_tokens":16}]} +{"session_id":"worker-a","turns":[{"model":"spawn-model","messages":[{"role":"user","content":"Collect the raw metrics."}],"max_tokens":24},{"model":"spawn-model","messages":[{"role":"user","content":"Summarize the metrics."}],"max_tokens":16}]} diff --git a/tests/fixtures/dag/graph_parity/spawn_repeat.dag.jsonl b/tests/fixtures/dag/graph_parity/spawn_repeat.dag.jsonl new file mode 100644 index 0000000000..9b72a1def2 --- /dev/null +++ b/tests/fixtures/dag/graph_parity/spawn_repeat.dag.jsonl @@ -0,0 +1,2 @@ +{"session_id":"root","turns":[{"model":"Qwen3-0.6B","messages":[{"role":"user","content":"plan"}],"max_tokens":32,"spawns":["kid"]},{"model":"Qwen3-0.6B","messages":[{"role":"user","content":"revise"}],"max_tokens":32,"spawns":["kid"]}]} +{"session_id":"kid","turns":[{"model":"Qwen3-0.6B","messages":[{"role":"user","content":"work"}],"max_tokens":16}]} diff --git a/tests/fixtures/dag/two_spawn_trees.dag.jsonl b/tests/fixtures/dag/two_spawn_trees.dag.jsonl new file mode 100644 index 0000000000..5169c1f03f --- /dev/null +++ b/tests/fixtures/dag/two_spawn_trees.dag.jsonl @@ -0,0 +1,4 @@ +{"session_id":"alpha","turns":[{"model":"spawn-model","messages":[{"role":"system","content":"You coordinate helpers."},{"role":"user","content":"Kick off the alpha analysis."}],"max_tokens":32,"spawns":[{"children":["alpha-helper"],"join_at":2}]},{"model":"spawn-model","messages":[{"role":"user","content":"Meanwhile, outline the alpha report."}],"max_tokens":24},{"model":"spawn-model","messages":[{"role":"user","content":"Fold in the alpha helper findings."}],"max_tokens":16}]} +{"session_id":"alpha-helper","turns":[{"model":"spawn-model","messages":[{"role":"user","content":"Collect the alpha raw metrics."}],"max_tokens":24},{"model":"spawn-model","messages":[{"role":"user","content":"Summarize the alpha metrics."}],"max_tokens":16}]} +{"session_id":"beta","turns":[{"model":"spawn-model","messages":[{"role":"system","content":"You coordinate helpers."},{"role":"user","content":"Kick off the beta analysis."}],"max_tokens":32,"spawns":[{"children":["beta-helper"],"join_at":2}]},{"model":"spawn-model","messages":[{"role":"user","content":"Meanwhile, outline the beta report."}],"max_tokens":24},{"model":"spawn-model","messages":[{"role":"user","content":"Fold in the beta helper findings."}],"max_tokens":16}]} +{"session_id":"beta-helper","turns":[{"model":"spawn-model","messages":[{"role":"user","content":"Collect the beta raw metrics."}],"max_tokens":24},{"model":"spawn-model","messages":[{"role":"user","content":"Summarize the beta metrics."}],"max_tokens":16}]} diff --git a/tests/harness/dynamo_synth_corpus.py b/tests/harness/dynamo_synth_corpus.py new file mode 100644 index 0000000000..fe18cf66c6 --- /dev/null +++ b/tests/harness/dynamo_synth_corpus.py @@ -0,0 +1,216 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Deterministic synthetic ``dynamo.request.trace.v1`` capture generator. + +The build-plane memory measurement (``tests/unit/dataset/graph/adapters/ +test_dynamo_corpus_scale_memory.py``) needs a corpus-scale dynamo capture whose +shape reproduces the two amplifiers real captures exhibit, WITHOUT committing a +multi-GB fixture: + +1. **Hash-id-slot amplification.** Dynamo records the FULL prompt block-hash + list on every ``request_end`` -- turn ``k`` re-lists all of turn ``k-1``'s + blocks plus its own fresh ones. A chain of ``T`` turns adding ``G`` fresh + blocks per turn therefore lists ``G * T(T+1)/2`` hash-id slots even though it + only introduces ``T * G`` unique blocks. After JSON decode every slot starts + as a distinct Python ``int`` object (u64 values, above the small-int cache), + which without interning would make the persistent ``TrieRequest.hash_ids`` + lists the single + largest build tier at corpus scale (~24.8 GB @1M nodes). + ``_collect_records`` interns those recorded ints read-time, so + the resident slot cost is one list pointer plus an amortized single canonical + ``int`` per UNIQUE value; this generator's ``H``-vs-``U`` amplification + (``H = G*T(T+1)/2`` slots over ``U = T*G`` unique blocks) is exactly what + exercises that intern table. +2. **Unique-block growth.** The ``T * G`` fresh blocks per session are globally + distinct (seeded 64-bit ids never recur across turns or sessions), so the + resolution automaton built by + :func:`~aiperf.dataset.graph.segment_ir.trie_content.resolve_content_parents` + (one int state per unique position) and the dynamo decode cache (one entry + per unique hash) both grow with the unique-block count. + +The generated capture is byte-deterministic given ``seed`` and every record is +block-consistent per +:func:`~aiperf.dataset.graph.adapters.dynamo.trie_lowering._assert_block_aligned` +(``(n-1)*bs < input_length <= n*bs``): ``input_length`` is set to +``n_hashes * block_size`` so every recorded block is fully covered and the +covered-count ISL gate reconstructs exactly ``input_length`` tokens. Records +parse through :func:`~aiperf.dataset.graph.adapters.dynamo.trace.from_dynamo_trace` +unchanged. +""" + +from __future__ import annotations + +import random +from dataclasses import dataclass +from pathlib import Path + +import orjson + +# Epoch-ish base so timestamps look like real Unix-ms captures; the absolute +# value is irrelevant (the trie lowering works trace-relative), only per-session +# monotonicity matters. +_BASE_UNIX_MS = 1_700_000_000_000 +# Per-turn start stride (a turn begins this many ms after the previous turn of +# its session). Sessions are offset by a 1 ms per-session jitter so no two +# records share an ``event_time`` (the interval-order rank tie-breaks on +# ``node_id`` regardless, but distinct times keep the sort stable). +_TURN_STRIDE_MS = 1_000 +# Every request's recorded interval is made to SPAN the whole capture window +# (:func:`_capture_duration_ms`), so all requests overlap pairwise and NOTHING +# finishes before anything else starts. The interval-order edge pass +# (``build_interval_edges``, O(candidates^2) per node) then sees an empty +# finished-before frontier at every node -- each node roots at START or +# start-anchors to its in-flight causal parent -- keeping edge derivation off +# the memory measurement's critical path. Edge STRUCTURE is irrelevant to the +# four memory tiers (hash-id ints, resolution trie, decode cache, content pool), +# which is all this capture is built to size. +_OVERLAP_MULTIPLIER = 4 +# Small fixed recorded completion; the response segment synthesizes this many +# tokens per node (a cheap partial-tail slice), so it is kept low to keep the +# generator's amplification concentrated in the prompt hash lists. +_OUTPUT_TOKENS = 8 + + +def _capture_duration_ms(sessions: int, turns_per_session: int) -> int: + """Recorded per-request duration that makes every interval overlap. + + Larger than the whole capture's start span (``turns * stride + sessions``), + so the earliest request only ends well after the latest request starts -- + the full-overlap condition the interval-order pass needs to stay cheap. + """ + span = turns_per_session * _TURN_STRIDE_MS + sessions + return span * _OVERLAP_MULTIPLIER + 1_000 + + +@dataclass(frozen=True, slots=True) +class SyntheticCorpusShape: + """Derived element counts for one synthetic capture's parameters. + + All three are pure functions of the generator parameters, so the memory + measurement's analytic per-tier model is computed from THESE (never from + hardcoded byte totals): + + * ``n_nodes`` -- one trie node per ``request_end`` (``sessions * turns``). + * ``n_hash_slots`` -- total ``TrieRequest.hash_ids`` entries across all + nodes = ``sessions * new_blocks_per_turn * T(T+1)/2``; the hash-id-int + tier scales with this. + * ``n_unique_blocks`` -- globally distinct block hashes = + ``sessions * turns * new_blocks_per_turn``; the resolution-trie and + decode-cache tiers scale with this. + * ``n_unique_segments`` -- globally distinct content-addressed message/response + segments = ``sessions * (2*turns - 1) + n_nodes`` (one prompt + one response + segment per turn, minus the shared per-session prompt root, plus the response + leaves); the content pool (``Segment`` objects) and the canonical + sid-string addressing tier scale with this. + """ + + n_nodes: int + n_hash_slots: int + n_unique_blocks: int + n_unique_segments: int + + +def synthetic_corpus_shape( + *, sessions: int, turns_per_session: int, new_blocks_per_turn: int +) -> SyntheticCorpusShape: + """Element counts a capture with these parameters will produce.""" + t = turns_per_session + return SyntheticCorpusShape( + n_nodes=sessions * t, + n_hash_slots=sessions * new_blocks_per_turn * (t * (t + 1) // 2), + n_unique_blocks=sessions * t * new_blocks_per_turn, + n_unique_segments=sessions * (2 * t - 1) + sessions * t, + ) + + +def _fresh_hash(rng: random.Random) -> int: + """One positive 64-bit block hash (u64; never 0, never negative). + + Dynamo records ``input_sequence_hashes`` as Rust ``u64`` and the reader + rejects negatives (they collide with the virtual negative-id namespace), so + the generator draws from ``[1, 2**64)``. + """ + return rng.getrandbits(64) or 1 + + +def write_synthetic_dynamo_capture( + path: str | Path, + *, + sessions: int, + turns_per_session: int, + new_blocks_per_turn: int, + block_size: int = 16, + seed: int, +) -> Path: + """Write a deterministic chain-heavy ``dynamo.request.trace.v1`` JSONL file. + + Emits ``sessions * turns_per_session`` ``request_end`` records shaped like + ``tests/unit/dataset/test_dynamo_streaming_store_parity.py``'s + ``_dynamo_record``. Within a session, turn ``k`` (1-indexed) carries the + full previous prefix plus ``new_blocks_per_turn`` fresh 64-bit hashes, so + ``input_sequence_hashes`` grows by ``new_blocks_per_turn`` each turn and + every earlier block is re-listed -- the hash-id-slot amplification. Fresh + hashes are globally unique (one seeded RNG stream, no recurrence), so the + unique-block count is exactly ``sessions * turns_per_session * + new_blocks_per_turn``. + + ``input_length`` is ``len(hashes) * block_size`` (fully block-covered), which + satisfies the block-alignment gate and makes the covered-count ISL target + equal to ``input_length``. Timestamps are monotone within each session. + + Records are streamed to disk one JSON line at a time so a corpus-scale + capture (millions of records) never materializes in memory. + + Returns the written ``path``. + """ + out = Path(path) + rng = random.Random(seed) + duration_ms = _capture_duration_ms(sessions, turns_per_session) + with out.open("wb") as f: + first = True + for s in range(sessions): + session_id = f"synth-{seed}-{s:08d}" + prefix: list[int] = [] + for k in range(1, turns_per_session + 1): + for _ in range(new_blocks_per_turn): + prefix.append(_fresh_hash(rng)) + n = len(prefix) + input_length = n * block_size + # +s jitter keeps event_times distinct across sessions; the + # (k-1) stride keeps a session's turns strictly increasing. + received_ms = _BASE_UNIX_MS + (k - 1) * _TURN_STRIDE_MS + s + record = { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": received_ms + duration_ms, + "event_source": "dynamo", + "agent_context": {"session_id": session_id}, + "request": { + "request_id": f"{session_id}-r{k}", + "model": "synthetic-model", + "input_tokens": input_length, + "output_tokens": _OUTPUT_TOKENS, + "cached_tokens": 0, + "request_received_ms": received_ms, + "total_time_ms": duration_ms, + "replay": { + "trace_block_size": block_size, + "input_length": input_length, + # copy() so each record owns its own list snapshot at + # this turn's length (the prefix keeps growing). + "input_sequence_hashes": prefix.copy(), + }, + }, + } + if not first: + f.write(b"\n") + f.write(orjson.dumps(record)) + first = False + return out + + +__all__ = [ + "SyntheticCorpusShape", + "synthetic_corpus_shape", + "write_synthetic_dynamo_capture", +] diff --git a/tests/integration/graph/fixtures/dynamo_dir/_build_fixture.py b/tests/integration/graph/fixtures/dynamo_dir/_build_fixture.py new file mode 100644 index 0000000000..28fba86904 --- /dev/null +++ b/tests/integration/graph/fixtures/dynamo_dir/_build_fixture.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Reproducibly emit the integration dynamo e2e fixture (``trace.jsonl.gz``). + +Run to regenerate:: + + uv run python tests/integration/graph/fixtures/dynamo_dir/_build_fixture.py + +Shape: ONE root session (``parent``, three clean linear ``request_end`` turns +with strictly-nested replay hashes) + ONE subagent (``child``, two nested-hash +turns) spliced at the parent's turn 3 via a ``:subagent:child:invoke`` tool call. +Exactly one root -> avoids the Task-8 multi-root gate. + +Each replay hash covers one 16-token block, so ``input_length == 16 * len(hashes)`` +stays block-aligned for the ``_assert_block_aligned_isl`` gate at +``trace_block_size=16`` (``(n-1)*16 < input_length <= n*16``). This mirrors the +component-integration fixture in +``tests/component_integration/graph/test_dynamo_e2e_materialize.py``. +""" + +from __future__ import annotations + +import gzip +from pathlib import Path +from typing import Any + +import orjson + +FIXTURE_DIR = Path(__file__).resolve().parent +MODEL = "m" +_BLOCK_SIZE = 16 + + +def _request_end( + *, + ts: int, + session_id: str, + hashes: list[int], + parent_session_id: str | None = None, +) -> dict[str, Any]: + """One ``request_end`` carrying replay hashes (one 16-token block each).""" + ctx: dict[str, Any] = {"session_id": session_id} + if parent_session_id is not None: + ctx["parent_session_id"] = parent_session_id + input_length = _BLOCK_SIZE * len(hashes) + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": ctx, + "request": { + "request_id": f"{session_id}-{ts}", + "model": MODEL, + "input_tokens": input_length, + "output_tokens": 16, + "cached_tokens": 0, + "replay": { + "trace_block_size": _BLOCK_SIZE, + "input_length": input_length, + "input_sequence_hashes": hashes, + }, + }, + } + + +def _tool_event( + *, ts: int, session_id: str, event_type: str, tool_call_id: str +) -> dict[str, Any]: + tool: dict[str, Any] = {"tool_call_id": tool_call_id, "tool_class": "search"} + if event_type == "tool_end": + tool["duration_ms"] = 40.0 + tool["status"] = "succeeded" + return { + "schema": "dynamo.request.trace.v1", + "event_type": event_type, + "event_time_unix_ms": ts, + "event_source": "harness", + "agent_context": {"session_id": session_id}, + "tool": tool, + } + + +def build_records() -> list[dict[str, Any]]: + return [ + _request_end(ts=1000, session_id="parent", hashes=[111, 222]), + _request_end(ts=1100, session_id="parent", hashes=[111, 222, 333]), + _request_end(ts=1200, session_id="parent", hashes=[111, 222, 333, 444]), + _tool_event( + ts=1220, + session_id="parent", + event_type="tool_start", + tool_call_id=":subagent:child:invoke", + ), + _tool_event( + ts=1260, + session_id="parent", + event_type="tool_end", + tool_call_id=":subagent:child:invoke", + ), + _request_end( + ts=1280, + session_id="child", + parent_session_id="parent", + hashes=[900, 901], + ), + _request_end( + ts=1300, + session_id="child", + parent_session_id="parent", + hashes=[900, 901, 902], + ), + _request_end(ts=1400, session_id="parent", hashes=[111, 222, 333, 444, 555]), + ] + + +def write_fixture() -> Path: + # Segment naming (`.NNNNNN.jsonl.gz`) so `DynamoTraceAdapter.can_load` + # autoroutes the DIRECTORY input (the plain `*.jsonl.gz` dir glob is only + # accepted by `discover_segments`, not by the dir can_load sniff). + out = FIXTURE_DIR / "trace.000000.jsonl.gz" + records = sorted(build_records(), key=lambda r: r["event_time_unix_ms"]) + with gzip.open(out, "wb") as f: + for r in records: + f.write(orjson.dumps(r)) + f.write(b"\n") + return out + + +if __name__ == "__main__": + print(write_fixture()) diff --git a/tests/integration/graph/fixtures/dynamo_dir/trace.000000.jsonl.gz b/tests/integration/graph/fixtures/dynamo_dir/trace.000000.jsonl.gz new file mode 100644 index 0000000000000000000000000000000000000000..dd0fc4d2d09adbb317099ef09feee30ff8a56853 GIT binary patch literal 454 zcmV;%0XhC3iwFpPzeH*R|8#O;V`VNdFfcGMFfM9yZ*FV=&6Qzqf-n$-*k4(B6SKSwbyHLDb%e+D=<|58(XN zw?e{n%LpZ~1>?EjH-KrUbq+>-Embw}V5%Lj9Vu9nB^E|&3u(wuo=SAy(EXn;1{`ULKpQSxY zRA_~d>anssnPO%5X?*CCf|LLX6dCOxE$%)@U3Zw(1L1!>WZmvb^nbR=u*AD%nXvIL wuPpgr=GpTNjT3WxMx4m`BZfvk05_hu5#UDI8;1s+&1N0sPYnZw4p0gJ0OG&azyJUM literal 0 HcmV?d00001 diff --git a/tests/integration/graph/fixtures/native_dynamic/accumulate_chain.yaml b/tests/integration/graph/fixtures/native_dynamic/accumulate_chain.yaml new file mode 100644 index 0000000000..c5e4cc43bb --- /dev/null +++ b/tests/integration/graph/fixtures/native_dynamic/accumulate_chain.yaml @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Native accumulate-chain graph for the conversation-channel E2E: each turn +# reads @hist and writes hist, so a downstream turn's prompt reconstructs the +# FULL user/assistant alternation of the prior turns (each user turn authored +# once; the assistant replies spliced live from the dynamic pool). The root +# `turn1` reads @hist (= the seeded system message) to satisfy Gate B. Nodes +# are non-streaming so the mock returns a single assembled body per request. +graph: + state: + hist: {type: messages, reducer: add_messages} + nodes: + turn1: + node_type: llm + prompt: + - "@hist" + - {role: user, content: "Name a primary color."} + output: hist + streaming: false + max_tokens: 12 + turn2: + node_type: llm + prompt: + - "@hist" + - {role: user, content: "Name a farm animal."} + output: hist + streaming: false + max_tokens: 12 + turn3: + node_type: llm + prompt: + - "@hist" + - {role: user, content: "Combine your two answers."} + output: t3_out + streaming: false + max_tokens: 12 + edges: + - {source: START, target: turn1} + - {source: turn1, target: turn2} + - {source: turn2, target: turn3} + - {source: turn3, target: END} +traces: + - id: conv-chain-1 + initial_state: + hist: + - {role: system, content: "You are terse."} diff --git a/tests/integration/graph/fixtures/native_dynamic/planner_reviewer.yaml b/tests/integration/graph/fixtures/native_dynamic/planner_reviewer.yaml new file mode 100644 index 0000000000..80123faeaf --- /dev/null +++ b/tests/integration/graph/fixtures/native_dynamic/planner_reviewer.yaml @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Native dynamic-content graph for the mock-server E2E: `review` splices +# `plan`'s ACTUAL response (block-level `@plan_out` ref) into its own prompt. +# Both nodes are non-streaming so the mock returns a single assembled body the +# test can extract verbatim. A small output cap keeps the spliced content short. +graph: + nodes: + plan: + node_type: llm + prompt: + - role: user + content: "Make a plan for testing a cache." + output: plan_out + streaming: false + max_tokens: 16 + review: + node_type: llm + prompt: + - role: user + content: + - "Review this plan: " + - "@plan_out" + output: review_out + streaming: false + max_tokens: 16 + edges: + - {source: START, target: plan} + - {source: plan, target: review} + - {source: review, target: END} +traces: + - id: conv-dynamic-1 diff --git a/tests/integration/graph/fixtures/weka_multigraph_dir/00_subagent.json b/tests/integration/graph/fixtures/weka_multigraph_dir/00_subagent.json new file mode 100644 index 0000000000..751f7d4ec8 --- /dev/null +++ b/tests/integration/graph/fixtures/weka_multigraph_dir/00_subagent.json @@ -0,0 +1,27 @@ +{ + "id": "trace_mg_subagent", + "models": ["claude-opus-4-5-20251101"], + "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "n", "model": "claude-opus-4-5-20251101", "in": 180, "out": 25, "hash_ids": [1, 2], "input_types": ["text"], "output_types": ["text"], "stop": "tool_use", "api_time": 0.8, "think_time": 0.0}, + { + "t": 1.0, + "type": "subagent", + "agent_id": "agent_001", + "subagent_type": "Explore", + "duration_ms": 4000, + "total_tokens": 600, + "tool_use_count": 1, + "status": "completed", + "models": ["claude-opus-4-5-20251101"], + "tool_tokens": 0, + "system_tokens": 0, + "requests": [ + {"t": 1.2, "type": "n", "model": "claude-opus-4-5-20251101", "in": 200, "out": 30, "hash_ids": [10, 11], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 0.9, "think_time": 0.0}, + {"t": 2.4, "type": "n", "model": "claude-opus-4-5-20251101", "in": 260, "out": 40, "hash_ids": [10, 11, 12], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 1.1, "think_time": 0.2} + ] + }, + {"t": 6.0, "type": "n", "model": "claude-opus-4-5-20251101", "in": 280, "out": 45, "hash_ids": [1, 2, 3, 4], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 1.3, "think_time": 0.5} + ] +} diff --git a/tests/integration/graph/fixtures/weka_multigraph_dir/01_plain.json b/tests/integration/graph/fixtures/weka_multigraph_dir/01_plain.json new file mode 100644 index 0000000000..13137e680e --- /dev/null +++ b/tests/integration/graph/fixtures/weka_multigraph_dir/01_plain.json @@ -0,0 +1,11 @@ +{ + "id": "trace_mg_plain", + "models": ["claude-opus-4-5-20251101"], + "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "n", "model": "claude-opus-4-5-20251101", "in": 180, "out": 25, "hash_ids": [1, 2], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 0.8, "think_time": 0.0}, + {"t": 1.5, "type": "n", "model": "claude-opus-4-5-20251101", "in": 230, "out": 35, "hash_ids": [1, 2, 3], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 1.0, "think_time": 0.5}, + {"t": 3.0, "type": "n", "model": "claude-opus-4-5-20251101", "in": 280, "out": 45, "hash_ids": [1, 2, 3, 4], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 1.3, "think_time": 0.5} + ] +} diff --git a/tests/integration/graph/test_dag_jsonl_fidelity_real.py b/tests/integration/graph/test_dag_jsonl_fidelity_real.py new file mode 100644 index 0000000000..84e4ccd6b7 --- /dev/null +++ b/tests/integration/graph/test_dag_jsonl_fidelity_real.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""REAL empirical proof: legacy vs graph dag_jsonl raw exports are wire-parity-equal. + +This is the Task-7 acceptance gate's real-run half (the EXTERNAL proof). The +in-process wire-parity golden gate +(``tests/component_integration/graph/test_dag_jsonl_byte_parity.py``) proves the +graph adapter emits payload-identical wire bodies (canonical order-insensitive +comparison) to the legacy path at the transport seam; the hermetic half of this gate +(``tests/component_integration/graph/test_dag_jsonl_fidelity.py``) drives +:func:`tools.dag_jsonl_fidelity.prove_parity` on hand-crafted synthetic export +fixtures. This file runs REAL ``aiperf profile --export-level raw`` subprocesses +and diffs the produced raw exports. + +It starts ONE shared ``aiperf-mock-server`` process (its per-prompt response +generation is deterministic WITHIN a process -- ``aiperf_mock_server/tokens.py`` +seeds off ``hash(...)`` -- so BOTH runs must hit the SAME server for fork-child +prompts, which embed the parent's live reply, to agree), then runs legacy +(``--custom-dataset-type dag_jsonl``) and graph (``--graph-format dag_jsonl``) +against it with identical ``--model``/``--url``/``--random-seed 42`` and +STREAMING ON (the first streaming parity exercise -- SSE responses flow through +``build_assistant_turn`` on both planes), a single worker, single pass. +``prove_parity`` on the two ``profile_export_raw.jsonl`` files must PASS every +criterion. + +It lives in the INTEGRATION lane (not component) on purpose: it spawns real +subprocesses against a live mock server, so the component-integration package's +in-process patches (e.g. the FakeTokenizer) do NOT cross the subprocess +boundary. Mirrors ``tests/integration/graph/test_weka_trace_fidelity_real.py`` +(server-startup helpers, ``.venv/bin`` skip guard, NO_PROXY/HF_HUB_OFFLINE env, +300s subprocess timeout). +""" + +from __future__ import annotations + +import os +import socket +import subprocess +import time +from pathlib import Path + +import pytest +from pytest import param + +from tools.dag_jsonl_fidelity import prove_parity + +_REPO = Path(__file__).resolve().parents[3] +_FIXTURE_DIR = _REPO / "tests" / "fixtures" / "dag" / "graph_parity" +_MODEL = "test-model" + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _wait_for_port(port: int, timeout_s: float = 15.0) -> bool: + deadline = time.time() + timeout_s + while time.time() < deadline: + with socket.socket() as s: + if s.connect_ex(("127.0.0.1", port)) == 0: + return True + time.sleep(0.1) + return False + + +def _run_profile( + aiperf_bin: Path, + *, + fixture: Path, + url: str, + format_flag: list[str], + artifact_dir: Path, + env: dict[str, str], + extra: list[str], +) -> subprocess.CompletedProcess[str]: + """Run one single-pass, streaming ``aiperf profile`` pass with warmup off. + + Warmup is OFF because no ``--warmup-*`` trigger flag is passed: a warmup + phase is only built when ``--warmup-request-count`` / ``--warmup-num-sessions`` + / ``--warmup-duration`` is explicitly set (see + ``src/aiperf/config/flags/_converter_warmup.py``), and those fields are all + ``gt=0`` so there is no "zero to disable" value -- omitting the trigger IS + the disable. ``prove_parity`` additionally filters to the PROFILING phase, + so even a stray warmup row on either plane cannot inflate the comparison. + """ + cmd = [ + str(aiperf_bin), + "profile", + "--input-file", + str(fixture), + "--url", + url, + "--endpoint-type", + "chat", + "--model", + _MODEL, + "--tokenizer", + "gpt2", + *format_flag, + "--num-conversations", + "1", + "--workers-max", + "1", + "--record-processor-service-count", + "1", + "--streaming", + "--export-level", + "raw", + "--artifact-dir", + str(artifact_dir), + "--random-seed", + "42", + "--ui", + "simple", + *extra, + ] + return subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=300) + + +@pytest.mark.integration +@pytest.mark.parametrize( + "fixture_name", + [ + param("fork_minimal.dag.jsonl", id="fork-minimal"), + param("mixed_full.dag.jsonl", id="mixed-full"), + ], +) # fmt: skip +def test_dag_jsonl_graph_raw_export_matches_legacy( + tmp_path: Path, fixture_name: str +) -> None: + """Legacy vs graph dag_jsonl raw exports are byte-parity-equal (streaming ON). + + ONE shared mock server answers both runs so the deterministic per-prompt + replies (and thus every FORK/spawn child prompt that embeds a parent reply) + are identical across the two planes. + """ + venv = _REPO / ".venv" / "bin" + mock_bin = venv / "aiperf-mock-server" + aiperf_bin = venv / "aiperf" + if not mock_bin.exists() or not aiperf_bin.exists(): + pytest.skip("aiperf / aiperf-mock-server not installed in .venv") + + fixture = _FIXTURE_DIR / fixture_name + assert fixture.is_file(), f"missing fixture {fixture}" + + port = _free_port() + url = f"http://127.0.0.1:{port}" + env = { + **os.environ, + "NO_PROXY": "127.0.0.1,localhost", + "HF_HUB_OFFLINE": "1", + "PYTHONUNBUFFERED": "1", + } + legacy_dir = tmp_path / "legacy" + graph_dir = tmp_path / "graph" + + mock = subprocess.Popen( + [str(mock_bin), "--port", str(port), "--ttft", "0", "--itl", "0"], + env=env, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + try: + assert _wait_for_port(port), "mock server did not come up" + legacy_proc = _run_profile( + aiperf_bin, + fixture=fixture, + url=url, + format_flag=["--custom-dataset-type", "dag_jsonl"], + artifact_dir=legacy_dir, + env=env, + # Legacy fork fanout dispatches children as separate sessions, so it + # needs >=2 session slots or a parent and its child serialize and the + # fork seeding stalls. The graph run below omits it because its replay + # lanes self-schedule fork/spawn children off the parent's completion. + extra=["--concurrency", "2"], + ) + assert legacy_proc.returncode == 0, ( + f"legacy profile exited {legacy_proc.returncode}\n" + f"STDERR tail:\n{legacy_proc.stderr[-3000:]}" + ) + graph_proc = _run_profile( + aiperf_bin, + fixture=fixture, + url=url, + format_flag=["--graph-format", "dag_jsonl"], + artifact_dir=graph_dir, + env=env, + extra=[], + ) + assert graph_proc.returncode == 0, ( + f"graph profile exited {graph_proc.returncode}\n" + f"STDERR tail:\n{graph_proc.stderr[-3000:]}" + ) + finally: + mock.terminate() + try: + mock.wait(timeout=10) + except subprocess.TimeoutExpired: + mock.kill() + + legacy_raw = next(legacy_dir.rglob("profile_export_raw.jsonl"), None) + graph_raw = next(graph_dir.rglob("profile_export_raw.jsonl"), None) + assert legacy_raw is not None, "legacy run produced no profile_export_raw.jsonl" + assert graph_raw is not None, "graph run produced no profile_export_raw.jsonl" + + report = prove_parity(legacy_raw, graph_raw) + assert report.passed, report.render() + # Non-vacuous: at least one request was actually compared on both planes. + assert report.messages.checked >= 1, report.render() diff --git a/tests/integration/graph/test_dag_multiworker_placement.py b/tests/integration/graph/test_dag_multiworker_placement.py new file mode 100644 index 0000000000..9c5eaf104a --- /dev/null +++ b/tests/integration/graph/test_dag_multiworker_placement.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DAG instance co-placement over the full multiprocess stack. + +The live A/B for graph INSTANCE co-placement: a real ``aiperf profile`` run +of TWO independent spawn+join dag trees at concurrency 2 with TWO workers +must place every trajectory (session) of one dag tree instance on ONE +worker. The spawned child sessions mint their own ``x_correlation_id``s, so +without the router keying graph sessions on the instance ``trace_id`` they +would balance least-loaded onto the OTHER tree's worker -- and the join +turn's splice would miss the worker-local dynamic pool (``pool_missing`` +trace-stop), which single-worker component tests can never observe. + +Tree instances are reconstructed from the records themselves: every child +record carries ``parent_correlation_id`` (the parent session's corr), so a +dag tree is a connected component over the corr <- parent-corr edges. The +decisive assertions: each component's records span exactly ONE worker, and +at least one component really contains multiple sessions (anti-vacuous). +""" + +from __future__ import annotations + +from collections import defaultdict +from pathlib import Path + +import pytest + +from tests.harness.utils import AIPerfCLI, AIPerfMockServer + +_FIXTURE = ( + Path(__file__).resolve().parents[2] + / "fixtures" + / "dag" + / "two_spawn_trees.dag.jsonl" +) + + +@pytest.mark.integration +@pytest.mark.asyncio +class TestDagMultiworkerPlacement: + async def test_each_dag_tree_pins_to_one_worker( + self, + cli: AIPerfCLI, + aiperf_mock_server: AIPerfMockServer, + ) -> None: + result = await cli.run( + f""" + aiperf profile \ + --model spawn-model \ + --url {aiperf_mock_server.url} \ + --input-file {_FIXTURE} \ + --graph-format dag_jsonl \ + --endpoint-type chat \ + --tokenizer builtin \ + --random-seed 1234 \ + --num-conversations 2 \ + --concurrency 2 \ + --workers-max 2 \ + --export-level raw \ + --ui simple + """, + timeout=300.0, + assert_success=True, + ) + assert result.exit_code == 0, result.stderr[-2000:] + + raw = result.raw_records + assert raw is not None and raw, "no raw records exported" + + # Union-find over corr <- parent-corr edges: one component per dag + # tree instance (spawned children carry parent_correlation_id). + root: dict[str, str] = {} + + def find(x: str) -> str: + root.setdefault(x, x) + while root[x] != x: + root[x] = root[root[x]] + x = root[x] + return x + + def union(a: str, b: str) -> None: + root[find(a)] = find(b) + + for rec in raw: + corr = rec.metadata.x_correlation_id + assert corr, f"record missing x_correlation_id: {rec.metadata}" + find(corr) + if rec.metadata.parent_correlation_id: + union(corr, rec.metadata.parent_correlation_id) + + workers_by_tree: dict[str, set[str]] = defaultdict(set) + sessions_by_tree: dict[str, set[str]] = defaultdict(set) + for rec in raw: + worker = rec.metadata.worker_id + assert worker is not None, f"record missing worker_id: {rec.metadata}" + tree = find(rec.metadata.x_correlation_id) + workers_by_tree[tree].add(worker) + sessions_by_tree[tree].add(rec.metadata.x_correlation_id) + + # THE co-placement invariant: no dag tree instance spans workers. + scattered = {t: w for t, w in workers_by_tree.items() if len(w) > 1} + assert not scattered, f"dag trees scattered across workers: {scattered}" + + # Anti-vacuous: at least one tree really had multiple sessions (the + # spawn actually produced child trajectories with their own corrs). + assert any(len(s) > 1 for s in sessions_by_tree.values()), ( + f"expected a multi-session dag tree, got per-tree session counts " + f"{[len(s) for s in sessions_by_tree.values()]}" + ) diff --git a/tests/integration/graph/test_dynamic_slots_mock_e2e.py b/tests/integration/graph/test_dynamic_slots_mock_e2e.py new file mode 100644 index 0000000000..99c8254b6b --- /dev/null +++ b/tests/integration/graph/test_dynamic_slots_mock_e2e.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dynamic content slots over the full multiprocess stack. + +The live E2E counterpart to `tests/unit/graph/test_dynamic_slots_e2e.py`: +a native two-node graph (`plan` -> `review`, where `review` splices `plan`'s +ACTUAL response via `@plan_out`) runs through real `aiperf profile` against the +in-repo mock server, with two workers. It proves the dynamic-content pipeline +wires end-to-end over real ZMQ / multiprocess -- capture on one credit, sticky +routing to the same worker, and splice on the next -- not just in-process. + +Asserted on the raw export: + +* the `review` request's wire payload contains `plan`'s exact response text + (the capture -> pool -> splice round-trip crossed process boundaries); +* both requests were handled by the SAME worker (per-trace sticky routing -- + dynamic content depends on it, so a break shows here first). +""" + +from __future__ import annotations + +from pathlib import Path + +import orjson +import pytest + +from tests.harness.utils import AIPerfCLI, AIPerfMockServer + +_FIX_DIR = Path(__file__).parent / "fixtures" / "native_dynamic" +FIXTURE = _FIX_DIR / "planner_reviewer.yaml" +CHAIN_FIXTURE = _FIX_DIR / "accumulate_chain.yaml" + +_PLAN_PROMPT = "Make a plan for testing a cache." +_REVIEW_PREFIX = "Review this plan: " + + +def _user_content(record) -> str: + messages = record.payload.get("messages", []) + assert messages, f"record has no messages: {record.payload}" + content = messages[-1].get("content") + assert isinstance(content, str), f"expected string content, got {content!r}" + return content + + +def _assistant_content(record) -> str: + """Extract the assistant text the mock returned (non-streaming chat body).""" + assert record.responses, f"record has no responses: {record.metadata}" + body = orjson.loads(record.responses[0].text) + return body["choices"][0]["message"]["content"] + + +@pytest.mark.integration +@pytest.mark.asyncio +class TestDynamicSlotsMockE2E: + async def test_review_prompt_contains_plan_response( + self, + cli: AIPerfCLI, + aiperf_mock_server: AIPerfMockServer, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + # Dynamic slots are incompatible with the t* snapshot window; keep it off. + # 0.0/0.0 is now the default for every run; pinned explicitly here so this + # test never depends on the ambient default. + + result = await cli.run( + f""" + aiperf profile \ + --model mock-model \ + --url {aiperf_mock_server.url} \ + --input-file {FIXTURE} \ + --graph-format native \ + --endpoint-type chat \ + --tokenizer builtin \ + --random-seed 7 \ + --num-conversations 1 \ + --concurrency 1 \ + --workers-max 2 \ + --export-level raw \ + --ui simple + """, + timeout=300.0, + assert_success=True, + ) + assert result.exit_code == 0, result.stderr[-2000:] + + raw = result.raw_records + assert raw is not None and len(raw) == 2, ( + f"expected 2 raw records (plan + review), got " + f"{None if raw is None else len(raw)}" + ) + + plan = next(r for r in raw if _PLAN_PROMPT in _user_content(r)) + review = next(r for r in raw if _user_content(r).startswith(_REVIEW_PREFIX)) + assert plan is not review + + # The producer actually returned non-empty content... + plan_response = _assistant_content(plan) + assert plan_response, "mock returned empty plan content" + + # ...and the consumer's wire prompt spliced it verbatim after the static + # instruction: capture -> pool -> splice survived the process boundary. + assert _user_content(review) == f"{_REVIEW_PREFIX}{plan_response}" + + # Dynamic content requires per-trace sticky routing: both credits of the + # one trace instance must have hit the same worker (else the pool value + # would have been missing and the trace would have errored). + assert plan.metadata.worker_id is not None + assert plan.metadata.worker_id == review.metadata.worker_id + + # No trace errored (a pool miss would surface as a record error). + assert all(r.error is None for r in raw), [ + r.error for r in raw if r.error is not None + ] + + async def test_accumulate_chain_reconstructs_full_alternation( + self, + cli: AIPerfCLI, + aiperf_mock_server: AIPerfMockServer, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """A 3-turn accumulate chain: turn3's wire prompt reconstructs the full + user/assistant alternation of turns 1-2, with the earlier assistant + messages being their ACTUAL live replies -- over the real stack.""" + + result = await cli.run( + f""" + aiperf profile \ + --model mock-model \ + --url {aiperf_mock_server.url} \ + --input-file {CHAIN_FIXTURE} \ + --graph-format native \ + --endpoint-type chat \ + --tokenizer builtin \ + --random-seed 11 \ + --num-conversations 1 \ + --concurrency 1 \ + --workers-max 2 \ + --export-level raw \ + --ui simple + """, + timeout=300.0, + assert_success=True, + ) + assert result.exit_code == 0, result.stderr[-2000:] + + raw = result.raw_records + assert raw is not None and len(raw) == 3, ( + f"expected 3 records, got {None if raw is None else len(raw)}" + ) + by_last_user = {_user_content(r): r for r in raw} + turn1 = by_last_user["Name a primary color."] + turn2 = by_last_user["Name a farm animal."] + turn3 = by_last_user["Combine your two answers."] + + r1 = _assistant_content(turn1) + r2 = _assistant_content(turn2) + assert r1 and r2, "mock returned empty content" + + # turn3's wire prompt is the full reconstructed conversation: the seed, + # each authored user turn, and each prior turn's LIVE reply, in order. + assert turn3.payload["messages"] == [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": "Name a primary color."}, + {"role": "assistant", "content": r1}, + {"role": "user", "content": "Name a farm animal."}, + {"role": "assistant", "content": r2}, + {"role": "user", "content": "Combine your two answers."}, + ] + + # All three credits pinned to one worker (per-trace sticky), no errors. + workers = {r.metadata.worker_id for r in raw} + assert len(workers) == 1 and None not in workers + assert all(r.error is None for r in raw) diff --git a/tests/integration/graph/test_dynamo_e2e.py b/tests/integration/graph/test_dynamo_e2e.py new file mode 100644 index 0000000000..1e3914307e --- /dev/null +++ b/tests/integration/graph/test_dynamo_e2e.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Real-mp E2E: a Dynamo agent-trace runs end-to-end and produces records. + +Runs the full multiprocess ``aiperf`` stack (subprocess ``python -m aiperf +profile`` + live mock server + real Worker + HTTP) over a current-schema +``dynamo.request.trace.v1`` directory: ONE root session (three clean linear +recorded-hash turns) + ONE subagent child (two recorded-hash turns) spliced at +the parent's turn 3. This is the FIRST proof dynamo produces records through +the real build+schedule pipeline -- if the build-plane and schedule-plane +ordinals disagree, every dispatch dies with ``GraphEnvelopeMissing`` and no +record lands. + +The fixture records carry replay metadata, so the adapter uses the recorded +``input_sequence_hashes`` (recorded-when-present, no mode selection). The +determinism env (start-ratio pin) mirrors ``test_weka_unified_store_ab.py``. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from tests.harness.utils import AIPerfCLI, AIPerfMockServer + +DYN_DIR = Path(__file__).parent / "fixtures" / "dynamo_dir" + +# Pinned seed so the synthesized recorded-hash content is deterministic. +_SEED = "1234" + + +async def _run( + cli: AIPerfCLI, + mock: AIPerfMockServer, + monkeypatch: pytest.MonkeyPatch, +) -> object: + """Run one full profile pass over the dynamo directory.""" + return await cli.run( + f""" + aiperf profile \ + --model m \ + --url {mock.url} \ + --endpoint-type chat \ + --input-file {DYN_DIR} \ + --tokenizer builtin \ + --random-seed {_SEED} \ + --num-conversations 1 \ + --concurrency 2 \ + --workers-max 2 \ + --export-level raw \ + --ui simple + """, + timeout=300.0, + assert_success=True, + ) + + +@pytest.mark.integration +@pytest.mark.asyncio +class TestDynamoE2E: + async def test_dynamo_produces_records( + self, + cli: AIPerfCLI, + aiperf_mock_server: AIPerfMockServer, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Dynamo dispatches and produces records with valid ISL (no GraphEnvelopeMissing). + + The flat trie lowering emits every dispatch node with a per-node scratch + ``{node_id}_out`` output channel (TEXT / overwrite, scalar-tolerant), so + the schedule-plane executor's scalar placeholder write never hits a list + reducer. The prompt content comes from the segment store via + ``prompt_segment_ids`` stamped at parse time by the shared LCP + segment-trie core. + """ + r = await _run(cli, aiperf_mock_server, monkeypatch) + assert r.exit_code == 0, r.stderr[-2000:] + + assert r.request_count > 0, "dynamo produced no profiling records" + assert r.jsonl, "no per-record metrics captured" + + isls: list[int] = [] + for rec in r.jsonl: + isl = rec.metrics.get("input_sequence_length") + assert isl is not None, f"record missing ISL: {rec.metadata}" + assert isl.value > 0, f"record has non-positive ISL: {isl.value}" + isls.append(int(isl.value)) + + # Content-parity guard: the fix is dispatch-completion only -- prompts are + # still materialized from the segment store (``prompt_segment_ids``), NOT + # the runtime channel. The fixture's root session accumulates strictly- + # nested replay hashes ([111,222] < [..,333] < [..,444] < [..,555]), + # so store-sourced prompts grow turn over turn: the recorded ISLs MUST be + # multiple distinct, growing values. Had the scalar placeholder ("") leaked + # into the prompt content (a content regression, not a dispatch-completion + # fix), every ISL would collapse to a single tiny constant. Requiring >=3 + # distinct ISLs spanning a real range proves the accumulating store content + # reached the wire unchanged. + distinct = sorted(set(isls)) + assert len(distinct) >= 3, ( + "expected multiple distinct store-sourced ISLs (accumulating prompts); " + f"got {distinct} -- a single/tiny constant would mean placeholder content " + "leaked past dispatch" + ) + assert max(isls) > min(isls), ( + f"ISLs must grow with the accumulating conversation; got {distinct}" + ) diff --git a/tests/integration/graph/test_first_token_live_timing.py b/tests/integration/graph/test_first_token_live_timing.py new file mode 100644 index 0000000000..6754f679ec --- /dev/null +++ b/tests/integration/graph/test_first_token_live_timing.py @@ -0,0 +1,584 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""LIVE proof that a post-TTFT overlapped child gates on the parent's OBSERVED +first token, while a pre-TTFT sibling stays dispatch-anchored. + +Drives ACTUAL ``aiperf profile --export-level raw`` runs of a purpose-built weka +trace (``weka_first_token_anchor.json``) whose STREAMING parent (``type: "s"``, +``out=200``) overlaps two async children, keyed by the recorded ``out`` value +that survives to the wire as ``payload.max_tokens``: + +* ``trace_first_token_anchor:0`` -- the long STREAMING parent P (``out=200``); + ``ttft + 200 x itl``. +* ``agent_pre:0`` -- the PRE-TTFT child (``out=45``), recorded at ``t=1.0`` + (before P's recorded first token at ``ttft=2.0s``): a pure DISPATCH anchor, + D=1.0s. +* ``agent_post:0`` -- the POST-TTFT child (``out=30``), recorded at ``t=4.0`` + (at/after P's recorded first token): a first-token-refined anchor, + D'=D-ttft=2.0s. +* ``trace_first_token_anchor:1`` -- the END-anchored tail R (``out=60``), fired + after P COMPLETES. + +The proof: two runs identical EXCEPT ``--ttft`` (1000ms vs 5000ms) move the +post-TTFT child by EXACTLY the controlled 4.0s TTFT delta -- because it gates at +the parent's OBSERVED first token (``~ttft``) + D' -- while the pre-TTFT control +holds at ~1.0s because it gates at the parent's DISPATCH + D, invariant to server +speed. That the delta transfers to the post-TTFT child and NOT to the pre-TTFT +child IS the feature. Both children fire while the parent is verifiably in flight +(``ttft + 200 x 60ms`` = 13s low / 17s high). + +Two operational levers make the nodes measurable (identical to +``test_start_anchor_live_timing``): + +* The t* snapshot window is off by default (``--trajectory-start-min/max-ratio`` + unset => t*=0), so every node dispatches LIVE in profiling (an engaged window + could snapshot the pre-t* nodes into warmup HISTORY). +* ``--extra-inputs ignore_eos:true`` makes the mock emit EXACTLY ``max_tokens`` + output tokens so P's duration is the intended ``ttft + out x itl``. + +The parent is STREAMING so the worker emits a ``FirstToken`` at the mock's first +SSE chunk (``~ttft``): the graph first-token observer stamps P's observed first +token and releases the post-TTFT child at that wall + D'. A pre-TTFT child never +subscribes to the first-token latch; it rides the dispatch anchor. + +Lives in the INTEGRATION lane: the structural precondition rebuilds the trie +graph in-process, and the fidelity of the live run depends on the real ``gpt2`` +tokenizer the subprocess uses. +""" + +from __future__ import annotations + +import os +import socket +import subprocess +import time +from collections import defaultdict +from pathlib import Path + +import orjson +import pytest + +from aiperf.dataset.graph.adapters.weka.trace_models import WekaTrace +from aiperf.dataset.graph.adapters.weka.trie_build import ( + ReconCallbacks, + build_trie_graph, +) + +_REPO = Path(__file__).resolve().parents[3] +_FIX = _REPO / "tests" / "unit" / "graph" / "fixtures" / "weka_first_token_anchor.json" +_MODEL = "claude-opus-4-5-20251101" + +# Record -> node identity key: each node's recorded ``out`` reaches the wire as +# ``payload.max_tokens`` (``dispatch_overrides.max_tokens == out``). +_PARENT_OUT = 200 # trace_first_token_anchor:0 -- the long STREAMING parent P +_PRE_OUT = 45 # agent_pre:0 -- PRE-TTFT child, pure dispatch anchor (D=1.0s) +_POST_OUT = 30 # agent_post:0 -- POST-TTFT child, first-token anchor (D'=2.0s) +_TAIL_OUT = 60 # trace_first_token_anchor:1 -- END-anchored tail, after P completes + +# D' (``delay_after_predecessor_first_token_us``) for the post-TTFT child, in +# seconds: gate = observed_first_token + D'. Recorded start 4.0s - ttft 2.0s. +_D_PRIME_S = 2.0 +# D (``delay_after_predecessor_start_us``) for the pre-TTFT child, in seconds: +# gate = parent_dispatch + D, invariant to server speed. +_D_PRE_S = 1.0 +_TOL_S = 0.25 + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _wait_for_port(port: int, timeout_s: float = 15.0) -> bool: + deadline = time.time() + timeout_s + while time.time() < deadline: + with socket.socket() as s: + if s.connect_ex(("127.0.0.1", port)) == 0: + return True + time.sleep(0.1) + return False + + +def _run_profile( + tmp_path: Path, *, ttft: int, itl: int, streaming: bool = True +) -> Path: + """Drive one live ``aiperf profile`` of the first-token fixture; return the raw export. + + ``ttft`` / ``itl`` set the mock latency model (ms); ``ignore_eos`` is always + forced so the streaming parent emits exactly ``max_tokens`` output tokens and + its duration is the intended ``ttft + out x itl``. + + ``streaming`` toggles the GLOBAL ``--streaming`` run flag. With the per-node + wire override the recorded ``"s"`` parent streams from its own + ``dispatch_overrides.stream`` even when the global flag is OFF (recorded mode + wins for graph credits), so the ``streaming=False`` run proves the parent + still emits a mid-flight ``FirstToken`` without the run itself being + streaming, while the recorded ``"n"`` children stay non-streaming in BOTH + modes. + """ + venv = _REPO / ".venv" / "bin" + mock_bin = venv / "aiperf-mock-server" + aiperf_bin = venv / "aiperf" + if not mock_bin.exists() or not aiperf_bin.exists(): + pytest.skip("aiperf / aiperf-mock-server not installed in .venv") + + port = _free_port() + tag = "stream" if streaming else "nostream" + artifact_dir = tmp_path / f"artifacts_ttft{ttft}_{tag}" + env = { + **os.environ, + "NO_PROXY": "127.0.0.1,localhost", + "HF_HUB_OFFLINE": "1", + "PYTHONUNBUFFERED": "1", + # t*=0 => full native replay: EVERY node dispatches live in profiling. + } + + mock = subprocess.Popen( + [str(mock_bin), "--port", str(port), "--ttft", str(ttft), "--itl", str(itl)], + env=env, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + try: + assert _wait_for_port(port), "mock server did not come up" + cmd = [ + str(aiperf_bin), + "profile", + "--input-file", + str(_FIX), + "--url", + f"http://127.0.0.1:{port}", + "--endpoint-type", + "chat", + "--model", + _MODEL, + "--tokenizer", + "gpt2", + "--num-conversations", + "1", + "--concurrency", + "4", + "--benchmark-duration", + "60", + "--export-level", + "raw", + "--artifact-dir", + str(artifact_dir), + "--random-seed", + "42", + "--extra-inputs", + "ignore_eos:true", + ] + if streaming: + # The GLOBAL ``--streaming`` flag is now only a FALLBACK: the recorded + # per-node ``dispatch_overrides.stream`` (``apply_run_level_payload_options``) + # stamps the wire-body ``stream`` per credit and the transport reads + # SSE vs JSON per request. The ``streaming=False`` run below proves the + # recorded ``"s"`` parent streams (and emits a mid-flight ``FirstToken``) + # WITHOUT this flag; here it exercises the mixed global-on case where + # the recorded ``"n"`` children must STILL stay non-streaming. + cmd.append("--streaming") + proc = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=300) + finally: + mock.terminate() + try: + mock.wait(timeout=10) + except subprocess.TimeoutExpired: + mock.kill() + + assert proc.returncode == 0, ( + f"aiperf profile (first-token, ttft={ttft}) exited {proc.returncode}\n" + f"STDERR tail:\n{proc.stderr[-3000:]}" + ) + raw = next(artifact_dir.rglob("profile_export_raw.jsonl"), None) + assert raw is not None, "no profile_export_raw.jsonl produced" + return raw + + +def _profiling_starts(raw: Path) -> dict[int, int]: + """Map ``payload.max_tokens`` -> earliest profiling ``request_start_ns`` (ns). + + ``max_tokens == out`` is the record->node identity key. Full-replay with + ``--num-conversations 1`` dispatches each node once; ``min`` is defensive + against any recycle re-fire. + """ + starts: dict[int, int] = {} + for line in raw.read_text().splitlines(): + if not line.strip(): + continue + obj = orjson.loads(line) + meta = obj.get("metadata", {}) or {} + if meta.get("benchmark_phase") != "profiling": + continue + start = meta.get("request_start_ns") + payload = obj.get("payload", {}) or {} + # The native LlmNode.max_tokens cap is endpoint-mapped to the wire + # token field (max_completion_tokens for chat). + max_tokens = payload.get("max_tokens") or payload.get("max_completion_tokens") + if start is None or max_tokens is None: + continue + prev = starts.get(max_tokens) + if prev is None or start < prev: + starts[max_tokens] = start + return starts + + +def _parent_finish_offset(raw: Path, p0: int) -> float: + """Parent P's completion offset (s) from its own dispatch -- flight-window proof.""" + for line in raw.read_text().splitlines(): + if not line.strip(): + continue + obj = orjson.loads(line) + meta = obj.get("metadata", {}) or {} + if meta.get("benchmark_phase") != "profiling": + continue + payload = obj.get("payload", {}) or {} + cap = payload.get("max_tokens") or payload.get("max_completion_tokens") + if cap == _PARENT_OUT: + end = meta.get("request_end_ns") + if end is not None: + return (end - p0) / 1e9 + return float("nan") + + +def _offsets(raw: Path) -> tuple[float, float, float]: + """Extract ``(d_pre, d_post, parent_finish)`` (all seconds from P's dispatch).""" + starts = _profiling_starts(raw) + for out in (_PARENT_OUT, _PRE_OUT, _POST_OUT, _TAIL_OUT): + assert out in starts, f"no profiling record for out={out}; got {sorted(starts)}" + p0 = starts[_PARENT_OUT] + d_pre = (starts[_PRE_OUT] - p0) / 1e9 + d_post = (starts[_POST_OUT] - p0) / 1e9 + parent_finish = _parent_finish_offset(raw, p0) + return d_pre, d_post, parent_finish + + +def _response_counts(raw: Path) -> dict[int, int]: + """Map ``payload.max_tokens`` -> wire ``responses`` count of its earliest record. + + A STREAMING request yields MANY ``responses`` entries (one ``SSEMessage`` per + SSE chunk + terminators); a non-streaming request yields exactly ONE JSON + body. Keyed by the earliest-start profiling record per node so it lines up + with :func:`_profiling_starts`. + """ + best: dict[int, tuple[int, int]] = {} + for line in raw.read_text().splitlines(): + if not line.strip(): + continue + obj = orjson.loads(line) + meta = obj.get("metadata", {}) or {} + if meta.get("benchmark_phase") != "profiling": + continue + start = meta.get("request_start_ns") + payload = obj.get("payload", {}) or {} + # The native LlmNode.max_tokens cap is endpoint-mapped to the wire + # token field (max_completion_tokens for chat). + max_tokens = payload.get("max_tokens") or payload.get("max_completion_tokens") + if start is None or max_tokens is None: + continue + n_resp = len(obj.get("responses") or []) + prev = best.get(max_tokens) + if prev is None or start < prev[0]: + best[max_tokens] = (start, n_resp) + return {mt: n for mt, (_start, n) in best.items()} + + +def _assert_mixed_mode_wire(raw: Path, label: str) -> tuple[int, int, int, int]: + """Recorded per-node stream mode reached the wire; return the response counts. + + The recorded STREAMING parent P (``out=200``) must have MULTIPLE ``responses`` + (SSE chunks); each recorded ``"n"`` child (``out=45/30/60``) must have EXACTLY + ONE (a single JSON body) -- even in a global ``--streaming`` run, the per-node + override keeps them non-streaming. Returns + ``(parent, pre, post, tail)`` counts for the report. + """ + counts = _response_counts(raw) + for out in (_PARENT_OUT, _PRE_OUT, _POST_OUT, _TAIL_OUT): + assert out in counts, ( + f"[{label}] no response count for out={out}; got {sorted(counts)}" + ) + assert counts[_PARENT_OUT] > 1, ( + f"[{label}] recorded streaming parent (out={_PARENT_OUT}) had " + f"{counts[_PARENT_OUT]} wire responses, expected MULTIPLE SSE chunks" + ) + for out in (_PRE_OUT, _POST_OUT, _TAIL_OUT): + assert counts[out] == 1, ( + f"[{label}] recorded 'n' child (out={out}) had {counts[out]} wire " + "responses, expected EXACTLY ONE non-streaming JSON body" + ) + return counts[_PARENT_OUT], counts[_PRE_OUT], counts[_POST_OUT], counts[_TAIL_OUT] + + +def _summary_metrics(raw: Path) -> dict: + """Parse the JSON summary export (``profile_export_aiperf.json``) beside the raw JSONL. + + The metrics JSON exporter (``exporters/metrics_json_exporter.py``) writes one + top-level key per metric ``tag``; each value is a ``JsonMetricResult`` dict + whose aggregate-counter value lives under ``avg`` and whose TTFT ``avg`` is in + the display unit (ms). Both exports write to the SAME artifacts dir, so the + summary sits beside the raw JSONL. + """ + summary = raw.parent / "profile_export_aiperf.json" + assert summary.exists(), f"no profile_export_aiperf.json beside {raw}" + return orjson.loads(summary.read_bytes()) + + +def _parent_dispatch_count(raw: Path) -> int: + """Count profiling records whose ``max_tokens == _PARENT_OUT`` (one per parent dispatch). + + The recorded ``"s"`` parent (``out=200``) is the ONLY streaming node, so its + dispatch count is the ground-truth streamed-request denominator; with + ``--num-conversations 1`` this is the number of trace instances. + """ + n = 0 + for line in raw.read_text().splitlines(): + if not line.strip(): + continue + obj = orjson.loads(line) + meta = obj.get("metadata", {}) or {} + if meta.get("benchmark_phase") != "profiling": + continue + payload = obj.get("payload", {}) or {} + cap = payload.get("max_tokens") or payload.get("max_completion_tokens") + if cap == _PARENT_OUT: + n += 1 + return n + + +def _assert_streaming_metrics( + raw: Path, label: str, ttft_s: float +) -> tuple[int, int, float]: + """Prove the per-record streaming metrics from the JSON summary export. + + * ``streamed_request_count`` == the parent-dispatch count (only the recorded + ``"s"`` parent streams; the three recorded ``"n"`` children do not); + * ``request_count`` == 4x that (parent + 3 children per trace instance); + * TTFT avg within ``_TOL_S`` of the CONTROLLED ``--ttft`` -- the numerical + proof that the non-streamed children (whose FULL latencies run ~2.3-13s) are + EXCLUDED from the TTFT distribution. A broken per-record guard would fold a + child's single completion timestamp in as its "first token" and drag the + TTFT mean far above the parent's ~``ttft``. + + Returns ``(streamed_request_count, request_count, ttft_avg_s)`` for the report. + """ + summary = _summary_metrics(raw) + n_parent = _parent_dispatch_count(raw) + assert n_parent >= 1, f"[{label}] no parent dispatch in raw export" + + src = summary.get("streamed_request_count") + assert src is not None, f"[{label}] streamed_request_count absent from summary" + assert src.get("avg") == n_parent, ( + f"[{label}] streamed_request_count={src.get('avg')} != " + f"{n_parent} parent dispatch(es)" + ) + + rc = summary.get("request_count") + assert rc is not None, f"[{label}] request_count absent from summary" + assert rc.get("avg") == 4 * n_parent, ( + f"[{label}] request_count={rc.get('avg')} != 4x parent dispatches " + f"({4 * n_parent})" + ) + + ttft = summary.get("time_to_first_token") + assert ttft is not None, ( + f"[{label}] time_to_first_token absent -- the streaming parent must " + "produce a TTFT under the relaxed per-record gate" + ) + ttft_avg_ms = ttft.get("avg") + assert ttft_avg_ms is not None, f"[{label}] time_to_first_token.avg is null" + ttft_avg_s = ttft_avg_ms / 1000.0 + assert abs(ttft_avg_s - ttft_s) <= _TOL_S, ( + f"[{label}] TTFT avg {ttft_avg_s:.3f}s not within {_TOL_S}s of controlled " + f"--ttft {ttft_s:.3f}s -- non-streamed children (~2.3-13s full latencies) " + "must be EXCLUDED; pollution here means the per-record guard is broken" + ) + return int(src["avg"]), int(rc["avg"]), ttft_avg_s + + +# --- structural precondition (cheap; before any live run) ------------------ + + +_BLOCK_SIZE = 64 + + +def _stub_decode_block_tokens(hash_ids: list[int]) -> list[int]: + out: list[int] = [] + for h in hash_ids: + out.extend(range(h * 100, h * 100 + _BLOCK_SIZE)) + return out + + +def _stub_partial_tail_tokens(n_tokens: int, seed: str) -> list[int]: + base = sum(ord(c) for c in seed) * 1000 + return list(range(base, base + n_tokens)) + + +def _stub_decode_tokens_to_text(tokens: list[int]) -> str: + return "|".join(str(t) for t in tokens) + + +_STUB_CALLBACKS = ReconCallbacks( + decode_block_tokens=_stub_decode_block_tokens, + sample_partial_tail_tokens=_stub_partial_tail_tokens, + decode_tokens_to_text=_stub_decode_tokens_to_text, +) + + +@pytest.mark.integration +def test_first_token_fixture_structural_preconditions() -> None: + """The trie build lowers the fixture to the exact post-TTFT edge geometry. + + Guards the live proof: if the fixture ever stops producing a pre-TTFT pure + dispatch anchor + a post-TTFT first-token anchor + an end-anchored tail, the + live timings below would be measuring the wrong thing. + """ + parsed, _pool = build_trie_graph( + WekaTrace.model_validate(orjson.loads(_FIX.read_bytes())), + callbacks=_STUB_CALLBACKS, + ) + incoming: dict[str, list] = defaultdict(list) + for edge in parsed.graph.edges: + incoming[edge.target].append(edge) + + parent = "trace_first_token_anchor:0" + + # parent -> agent_pre:0: pre-TTFT (recorded start 1.0 < ttft 2.0) => pure + # dispatch anchor, D=1.0s, no first-token refinement. + (pre,) = incoming["agent_pre:0"] + assert pre.source == parent + assert pre.delay_after_predecessor_start_us == pytest.approx(1.0e6) + assert pre.delay_after_predecessor_first_token_us is None + + # parent -> agent_post:0: post-TTFT (recorded start 4.0 >= ttft 2.0) => + # first-token anchor, D=4.0s with D'=D-ttft=2.0s. + (post,) = incoming["agent_post:0"] + assert post.source == parent + assert post.delay_after_predecessor_start_us == pytest.approx(4.0e6) + assert post.delay_after_predecessor_first_token_us == pytest.approx(2.0e6) + + # parent -> tail: END-anchored tail, from the parent ALONE (async children + # excluded from the AND-join), completion + recorded 1.5s gap. + (tail,) = incoming["trace_first_token_anchor:1"] + assert tail.source == parent + assert tail.delay_after_predecessor_us == pytest.approx(1.5e6) + assert tail.delay_after_predecessor_start_us is None + assert tail.delay_after_predecessor_first_token_us is None + + # The parent must be STREAMING so the worker emits a FirstToken (the store + # builder derives the wire ``stream`` flag from the native field). + assert parsed.graph.nodes[parent].streaming is True + + +# --- live proof ------------------------------------------------------------ + + +@pytest.mark.integration +def test_first_token_live_timing(tmp_path: Path) -> None: + """The post-TTFT child tracks the OBSERVED first token; the pre-TTFT holds. + + Two runs identical except ``--ttft`` (1000ms vs 5000ms): the controlled 4.0s + TTFT delta transfers EXACTLY to the post-TTFT child (observed_first_token + + D'), while the pre-TTFT dispatch-anchored control stays at ~1.0s. + """ + raw_low = _run_profile(tmp_path, ttft=1000, itl=60) + d_pre_low, d_post_low, finish_low = _offsets(raw_low) + n_parent_low, n_pre_low, n_post_low, n_tail_low = _assert_mixed_mode_wire( + raw_low, "low" + ) + src_low, rc_low, ttft_avg_low = _assert_streaming_metrics(raw_low, "low", 1.0) + print( + f"[low ] ttft=1000 itl=60: parent_finish={finish_low:.3f}s " + f"d_pre={d_pre_low:.3f}s d_post={d_post_low:.3f}s " + f"resp(parent={n_parent_low} pre={n_pre_low} post={n_post_low} " + f"tail={n_tail_low}) " + f"streamed_request_count={src_low} request_count={rc_low} " + f"ttft_avg={ttft_avg_low:.3f}s" + ) + + raw_high = _run_profile(tmp_path, ttft=5000, itl=60) + d_pre_high, d_post_high, finish_high = _offsets(raw_high) + n_parent_high, n_pre_high, n_post_high, n_tail_high = _assert_mixed_mode_wire( + raw_high, "high" + ) + src_high, rc_high, ttft_avg_high = _assert_streaming_metrics(raw_high, "high", 5.0) + print( + f"[high] ttft=5000 itl=60: parent_finish={finish_high:.3f}s " + f"d_pre={d_pre_high:.3f}s d_post={d_post_high:.3f}s " + f"resp(parent={n_parent_high} pre={n_pre_high} post={n_post_high} " + f"tail={n_tail_high}) " + f"streamed_request_count={src_high} request_count={rc_high} " + f"ttft_avg={ttft_avg_high:.3f}s" + ) + + for label, ttft_s, d_pre, d_post, finish in ( + ("low", 1.0, d_pre_low, d_post_low, finish_low), + ("high", 5.0, d_pre_high, d_post_high, finish_high), + ): + # Post-TTFT child = OBSERVED first token (~ttft) + D' (2.0s). + assert abs(d_post - (ttft_s + _D_PRIME_S)) <= _TOL_S, ( + f"[{label}] post-TTFT child at {d_post:.3f}s, " + f"expected {ttft_s + _D_PRIME_S:.3f}s (first token + D')" + ) + # Pre-TTFT child = dispatch anchor, INVARIANT to server speed. + assert abs(d_pre - _D_PRE_S) <= _TOL_S, ( + f"[{label}] pre-TTFT child at {d_pre:.3f}s, expected {_D_PRE_S:.3f}s " + "(dispatch anchor must NOT move with --ttft)" + ) + # Both children fired while the parent was verifiably IN FLIGHT. + assert d_post < finish and d_pre < finish, ( + f"[{label}] children did not fire mid-parent-flight: " + f"d_pre={d_pre:.3f} d_post={d_post:.3f} parent_finish={finish:.3f}" + ) + + # The controlled +4.0s TTFT delta transfers EXACTLY to the post-TTFT child. + assert abs((d_post_high - d_post_low) - 4.0) <= _TOL_S, ( + f"post-TTFT delta {d_post_high - d_post_low:.3f}s != 4.0s " + f"(low={d_post_low:.3f} high={d_post_high:.3f})" + ) + + +@pytest.mark.integration +def test_first_token_live_timing_no_global_streaming(tmp_path: Path) -> None: + """The recorded ``"s"`` parent streams WITHOUT the global ``--streaming`` flag. + + Same fixture and levers as the low run of :func:`test_first_token_live_timing` + (``ttft=1000`` / ``itl=60``) but with NO global ``--streaming``. The recorded + per-node ``dispatch_overrides.stream=True`` wins for graph credits, so the + parent still streams and emits a mid-flight ``FirstToken``: the post-TTFT + child anchors at observed_first_token + D' (~3.0s) and the pre-TTFT child + holds at its dispatch anchor (~1.0s), IDENTICAL to the global-on low run. The + recorded ``"n"`` children stay non-streaming here too (single JSON body). + """ + raw = _run_profile(tmp_path, ttft=1000, itl=60, streaming=False) + d_pre, d_post, finish = _offsets(raw) + n_parent, n_pre, n_post, n_tail = _assert_mixed_mode_wire(raw, "no-stream") + # SAME expectations as the flagged ttft=1000 low run: per-record semantics, not + # flag semantics -- streamed_request_count == parent dispatches, request_count + # == 4x, TTFT avg ~= the controlled 1.0s (children EXCLUDED). Asserting against + # the identical targets IS the "same as global-on low run" proof. + src, rc, ttft_avg = _assert_streaming_metrics(raw, "no-stream", 1.0) + print( + f"[no-stream] ttft=1000 itl=60 (NO --streaming): " + f"parent_finish={finish:.3f}s d_pre={d_pre:.3f}s d_post={d_post:.3f}s " + f"resp(parent={n_parent} pre={n_pre} post={n_post} tail={n_tail}) " + f"streamed_request_count={src} request_count={rc} " + f"ttft_avg={ttft_avg:.3f}s" + ) + + # Post-TTFT child = OBSERVED first token (~ttft 1.0s) + D' (2.0s) = 3.0s: the + # parent MUST have streamed a mid-flight FirstToken despite no global flag. + assert abs(d_post - (1.0 + _D_PRIME_S)) <= _TOL_S, ( + f"[no-stream] post-TTFT child at {d_post:.3f}s, expected " + f"{1.0 + _D_PRIME_S:.3f}s (first token + D') -- recorded parent must " + "stream without --streaming" + ) + # Pre-TTFT child = pure dispatch anchor (~1.0s), unaffected by streaming mode. + assert abs(d_pre - _D_PRE_S) <= _TOL_S, ( + f"[no-stream] pre-TTFT child at {d_pre:.3f}s, expected {_D_PRE_S:.3f}s" + ) + # Both children fired while the parent was verifiably IN FLIGHT. + assert d_post < finish and d_pre < finish, ( + f"[no-stream] children did not fire mid-parent-flight: " + f"d_pre={d_pre:.3f} d_post={d_post:.3f} parent_finish={finish:.3f}" + ) diff --git a/tests/integration/graph/test_issue_1106_selection.py b/tests/integration/graph/test_issue_1106_selection.py new file mode 100644 index 0000000000..99fe318b56 --- /dev/null +++ b/tests/integration/graph/test_issue_1106_selection.py @@ -0,0 +1,308 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Real-mp E2E: the #1106 graph dataset-selection contract, scenarios 1-3. + +Runs the full multiprocess ``aiperf`` stack (subprocess ``python -m aiperf +profile`` + live mock server + real Worker + HTTP) to lock the ai-dynamo/aiperf +#1106 dataset-selection contract end-to-end: + +* **Scenario 1 (fail loud).** ``--max-context-length`` rejects most traces, + leaving far fewer distinct traces than the requested ``--concurrency 100``; + with cache-bust OFF and wrapping unset the run FAILS LOUDLY -- the wrap-guard's + ``ConfigurationError`` surfaces and no records are produced, instead of the old + silent clone-to-fill. (The guard fires in the AWAITED ``setup_phase``, so it + propagates through the phase-failure path to a non-zero exit; it does not hang.) +* **Scenario 2 (cache-bust rescue).** The SAME over-subscribing shape SUCCEEDS + once ``--cache-bust first_turn_prefix`` turns wrapping on: the finite eligible + pool is intentionally recycled to fill the requested concurrency and the run + completes with real records. +* **Scenario 3 (exact-N selection).** A corpus of >= 100 eligible traces with + ``--num-dataset-entries 100`` dispatches EXACTLY 100 distinct traces -- not the + whole corpus, not fewer. + +The determinism env (t* window pinned to ``[0, 0]`` = full replay) mirrors the +sibling weka/dynamo e2e tests. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import orjson +import pytest + +from tests.harness.dynamo_synth_corpus import write_synthetic_dynamo_capture +from tests.harness.utils import AIPerfCLI, AIPerfMockServer + +_MODEL = "m" + +# --max-context-length threshold: small eligible traces (peak 144) pass, big +# traces (peak 8200) are rejected by the graph filter-then-cap selection. +_MAX_CTX = 1000 +_ELIGIBLE = 3 # small traces surviving the peak-context filter +_REJECTED = 5 # oversized traces the filter drops +_NUM_CONVERSATIONS = 6 # session cap: forces the eligible pool to wrap (> _ELIGIBLE) + + +def _weka_trace(trace_id: str, *, in_tokens: int, out_tokens: int) -> dict[str, Any]: + """One minimal single-turn weka trace (peak context == ``in + out``).""" + return { + "id": trace_id, + "models": [_MODEL], + "block_size": 64, + "hash_id_scope": "local", + "requests": [ + { + "t": 0.0, + "type": "n", + "model": _MODEL, + "in": in_tokens, + "out": out_tokens, + "hash_ids": [1, 2], + "input_types": ["text"], + "output_types": ["text"], + "stop": "end_turn", + "api_time": 0.5, + "think_time": 0.0, + } + ], + } + + +def _write_1106_weka_corpus(directory: Path) -> None: + """Write a #1106-shaped weka dir: a few small traces + many oversized ones.""" + directory.mkdir(parents=True, exist_ok=True) + for i in range(_ELIGIBLE): + (directory / f"small_{i:03d}.json").write_bytes( + orjson.dumps(_weka_trace(f"small-{i}", in_tokens=128, out_tokens=16)) + ) + for i in range(_REJECTED): + (directory / f"big_{i:03d}.json").write_bytes( + orjson.dumps(_weka_trace(f"big-{i}", in_tokens=8000, out_tokens=200)) + ) + + +def _template_ids(records: list[Any]) -> set[str]: + """Distinct template trace ids over records. + + ``conversation_id`` is the nonce-less trajectory TEMPLATE (``{trace}`` for + the root scope, ``{trace}::{scope}`` for children); the base trace id is + its first ``::`` segment. + """ + templates: set[str] = set() + for rec in records: + conv = rec.metadata.conversation_id + assert conv is not None, f"record missing conversation_id: {rec.metadata}" + templates.add(conv.split("::", 1)[0]) + return templates + + +def _instance_ids(records: list[Any]) -> set[str]: + """Distinct per-INSTANCE ids over records. + + Instance identity rides ``x_correlation_id`` (``{conversation}::{nonce}``, + fresh per trajectory instance): every wrap/recycle of one template mints a + new corr, so distinct corrs > distinct templates proves cloning happened. + """ + return { + rec.metadata.x_correlation_id + for rec in records + if rec.metadata.x_correlation_id is not None + } + + +@pytest.mark.integration +@pytest.mark.asyncio +class TestIssue1106Selection: + async def test_oversubscription_without_cache_bust_fails_loud( + self, + cli: AIPerfCLI, + aiperf_mock_server: AIPerfMockServer, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + ) -> None: + """Scenario 1: filter drops most traces, then concurrency 100 FAILS LOUD. + + ``--max-context-length 1000`` rejects the 5 oversized traces (peak 8200), + leaving 3 eligible; ``--num-dataset-entries 100`` does not cap below that; + resolved concurrency 100 exceeds the 3 distinct loaded traces with + cache-bust OFF and wrapping unset. The wrap-guard raises a + ``ConfigurationError`` in the awaited ``setup_phase``, so the run exits + NON-ZERO with the guard's message surfaced -- NOT a silent clone-to-fill, + and NOT a hang (no ``--benchmark-duration`` is set; the old orphaned-task + behavior would have waited on the sending event forever). + """ + + corpus_dir = tmp_path / "corpus" + _write_1106_weka_corpus(corpus_dir) + + result = await cli.run( + f""" + aiperf profile \ + --model {_MODEL} \ + --url {aiperf_mock_server.url} \ + --endpoint-type chat \ + --input-file {corpus_dir} \ + --tokenizer builtin \ + --random-seed 1234 \ + --max-context-length {_MAX_CTX} \ + --num-dataset-entries 100 \ + --concurrency 100 \ + --cache-bust none \ + --workers-max 2 \ + --ui simple + """, + timeout=90.0, + assert_success=False, + ) + # Fails loudly: non-zero exit (the guard's ConfigurationError propagates + # through setup_phase -> the phase-failure path), NOT a clean success. + assert result.exit_code != 0, "over-subscription must fail, not succeed" + + combined = (result.stderr + result.stdout + (result.log or "")).lower() + # The guard's actionable message surfaced -- the operator sees WHY it + # failed, not a generic timeout / no-records error. + assert ( + f"concurrency 100 exceeds {_ELIGIBLE} distinct loaded traces" in combined + ), ( + f"expected the wrap-guard message in the run output; got tail: " + f"{combined[-1500:]}" + ) + assert "dataset wrapping is disabled" in combined + assert "--cache-bust first_turn_prefix" in combined + # NOT a silent clone-to-fill: no profiling records were produced. + assert result.request_count == 0, ( + "the guard must fail BEFORE dispatch, not clone traces to fill lanes" + ) + assert not result.jsonl, "no records should be produced when the guard fires" + + async def test_cache_bust_rescues_oversubscribed_corpus( + self, + cli: AIPerfCLI, + aiperf_mock_server: AIPerfMockServer, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + ) -> None: + """Scenario 2: over-subscription + ``--cache-bust`` wraps and SUCCEEDS. + + ``--max-context-length 1000`` rejects the 5 oversized traces, leaving 3 + eligible; ``--num-dataset-entries 100`` keeps all 3; resolved concurrency + 100 exceeds those 3. With cache-bust ON, wrapping is allowed (the guard + does not fire), the eligible pool recycles to fill the lanes, and the run + completes with real records. ``--num-conversations 6`` bounds the recycle. + """ + + corpus_dir = tmp_path / "corpus" + _write_1106_weka_corpus(corpus_dir) + + result = await cli.run( + f""" + aiperf profile \ + --model {_MODEL} \ + --url {aiperf_mock_server.url} \ + --endpoint-type chat \ + --input-file {corpus_dir} \ + --tokenizer builtin \ + --random-seed 1234 \ + --max-context-length {_MAX_CTX} \ + --num-dataset-entries 100 \ + --concurrency 100 \ + --num-conversations {_NUM_CONVERSATIONS} \ + --cache-bust first_turn_prefix \ + --workers-max 2 \ + --export-level raw \ + --ui simple + """, + timeout=300.0, + assert_success=True, + ) + assert result.exit_code == 0, result.stderr[-2000:] + assert result.request_count > 0, "over-subscribed run produced no records" + assert result.jsonl, "no per-record metrics captured" + + # Selection kept exactly the 3 eligible traces (the 5 oversized ones were + # filtered out by --max-context-length before the build). + templates = _template_ids(result.jsonl) + assert templates == {f"small-{i}" for i in range(_ELIGIBLE)}, ( + f"expected exactly the {_ELIGIBLE} small eligible templates, got {templates}" + ) + # Wrapping happened: more distinct instances than distinct templates means + # the finite eligible pool was cloned/recycled to fill the lanes (the + # duplication cache-bust made safe), which is the whole point of the + # cache-bust rescue vs. scenario 1's fail-loud. + instances = _instance_ids(result.jsonl) + assert len(instances) > len(templates), ( + f"expected wrapping to clone the {len(templates)} eligible templates " + f"into more instances, got {len(instances)} instances" + ) + + async def test_num_dataset_entries_selects_exactly_100( + self, + cli: AIPerfCLI, + aiperf_mock_server: AIPerfMockServer, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + ) -> None: + """Scenario 3: >=100 eligible traces + ``--num-dataset-entries 100`` -> 100 distinct. + + A 130-session dynamo capture (each session one tiny single-turn trace, all + well under any context cap) is capped to 100 by ``--num-dataset-entries``. + + The discrimination is that MORE conversations run than the cap: + ``--num-conversations 130`` with ``--allow-dataset-wrap`` dispatches 130 + sessions over the loaded pool. A WORKING cap loads 100 traces, so the 130 + sessions wrap (sequential draw ``x % 100``) and cover exactly 100 distinct + templates; a BROKEN cap (a no-op that loaded all 130) would draw ``x % 130`` + over 130 sessions and produce 130 distinct templates. Asserting exactly 100 + therefore fails if the cap did not bind -- unlike a run whose conversation + count equals the cap, where the draw yields 100 distinct regardless. + """ + + capture = tmp_path / "dynamo_130.jsonl" + write_synthetic_dynamo_capture( + capture, + sessions=130, + turns_per_session=1, + new_blocks_per_turn=2, + block_size=16, + seed=1234, + ) + + result = await cli.run( + f""" + aiperf profile \ + --model {_MODEL} \ + --url {aiperf_mock_server.url} \ + --endpoint-type chat \ + --input-file {capture} \ + --tokenizer builtin \ + --random-seed 1234 \ + --num-dataset-entries 100 \ + --num-conversations 130 \ + --allow-dataset-wrap \ + --concurrency 10 \ + --workers-max 2 \ + --ui simple + """, + timeout=300.0, + assert_success=True, + ) + assert result.exit_code == 0, result.stderr[-2000:] + assert result.jsonl, "no per-record metrics captured" + + # 130 conversations ran (more than the cap), but the cap bound the loaded + # corpus to 100 of the 130 sessions, so the wrapped draw covers EXACTLY 100 + # distinct templates -- a no-op cap (all 130 loaded) would yield 130. + templates = _template_ids(result.jsonl) + assert len(templates) == 100, ( + f"expected exactly 100 distinct dispatched templates from the " + f"--num-dataset-entries cap over a 130-session corpus, got " + f"{len(templates)} (130 => the cap was a no-op)" + ) + # Sanity: more distinct INSTANCES than templates confirms wrapping actually + # happened (so the 100 is a wrapped-cover, not a short 100-session pass). + assert len(_instance_ids(result.jsonl)) > len(templates), ( + "expected wrapping (distinct instances > distinct templates) so the " + "100-distinct-templates result is a genuine capped-then-wrapped cover" + ) diff --git a/tests/integration/graph/test_start_anchor_live_timing.py b/tests/integration/graph/test_start_anchor_live_timing.py new file mode 100644 index 0000000000..65b69054bf --- /dev/null +++ b/tests/integration/graph/test_start_anchor_live_timing.py @@ -0,0 +1,327 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""LIVE proof that start-anchored edges dispatch at parent DISPATCH, not completion. + +Drives ACTUAL ``aiperf profile --export-level raw`` runs of a purpose-built weka +trace (``weka_start_anchor.json``) against the in-repo mock server at two server +speeds and compares the EXACT firing timestamps of four nodes, keyed by the +recorded ``out`` value that survives to the wire as ``payload.max_tokens``: + +* ``trace_start_anchor:0`` -- the long parent P (``out=200``); ``ttft 500ms + + 200 x itl`` per token. +* ``agent_spawn:0`` -- a mid-flight spawn child C (``out=30``) start-anchored to + P at +2.5s. +* ``agent_chain:0`` -- a mid-flight overlap child Q (``out=45``) start-anchored + to P at +5.0s. +* ``trace_start_anchor:1`` -- the END-anchored tail R (``out=60``), fired + ``1.0s`` after P COMPLETES. + +The proof: the two start-anchored children dispatch at a CONSTANT dispatch-to- +dispatch offset (2.5s / 5.0s from P's dispatch) at BOTH server speeds -- even in +the slow run, where P is still in flight (~12.5s) when both children fire -- while +the end-anchored tail tracks P's actual COMPLETION (compresses to ~2.5s when the +server is fast, expands to ~13.5s when slow). Start-anchored offsets are invariant +to server speed; end-anchored offsets are not. That contrast IS the feature. + +Two operational levers make the four nodes measurable: + +* The t* snapshot window is off by default (``--trajectory-start-min/max-ratio`` + unset => t*=0), so every node dispatches LIVE in profiling. An engaged window + could snapshot the pre-t* nodes (P, C, Q) into cache-priming warmup HISTORY + and leave only the tail dispatched. +* ``--extra-inputs ignore_eos:true`` makes the mock emit EXACTLY ``max_tokens`` + output tokens (the mock otherwise samples a variable, prompt-hash-seeded count), + so P's duration is the intended ``ttft + out x itl`` and the end-anchored tail's + completion offset is deterministic. + +Lives in the INTEGRATION lane (not component): the fidelity report functions +rebuild the trie graph in-process with the real ``gpt2`` tokenizer the subprocess +uses, so the reconstruction must not diverge on a patched FakeTokenizer. +""" + +from __future__ import annotations + +import os +import socket +import subprocess +import time +from pathlib import Path + +import orjson +import pytest +from pytest import param + +from tools.weka_trace_fidelity import ( + causality_timing_vs_real_trace, + content_vs_real_trace, +) + +_REPO = Path(__file__).resolve().parents[3] +_FIX = _REPO / "tests" / "unit" / "graph" / "fixtures" / "weka_start_anchor.json" +_MODEL = "claude-opus-4-5-20251101" + +# Record -> node identity key: each node's recorded ``out`` reaches the wire as +# ``payload.max_tokens`` (``dispatch_overrides.max_tokens == out``). +_PARENT_OUT = 200 # trace_start_anchor:0 -- the long parent P +_SPAWN_OUT = 30 # agent_spawn:0 -- spawn child, start-anchored to P at +2.5s +_CHAIN_OUT = 45 # agent_chain:0 -- overlap child, start-anchored to P at +5.0s +_TAIL_OUT = 60 # trace_start_anchor:1 -- END-anchored tail, 1.0s after P ends + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _wait_for_port(port: int, timeout_s: float = 15.0) -> bool: + deadline = time.time() + timeout_s + while time.time() < deadline: + with socket.socket() as s: + if s.connect_ex(("127.0.0.1", port)) == 0: + return True + time.sleep(0.1) + return False + + +def _run_profile(tmp_path: Path, *, ttft: int, itl: int, ignore_eos: bool) -> Path: + """Drive one live ``aiperf profile`` of the start-anchor fixture; return the raw export. + + ``ttft`` / ``itl`` set the mock latency model (ms). ``ignore_eos`` forces the + mock to emit exactly ``max_tokens`` tokens (needed for a deterministic parent + duration; irrelevant at zero latency). + """ + venv = _REPO / ".venv" / "bin" + mock_bin = venv / "aiperf-mock-server" + aiperf_bin = venv / "aiperf" + if not mock_bin.exists() or not aiperf_bin.exists(): + pytest.skip("aiperf / aiperf-mock-server not installed in .venv") + + port = _free_port() + artifact_dir = tmp_path / "artifacts" + env = { + **os.environ, + "NO_PROXY": "127.0.0.1,localhost", + "HF_HUB_OFFLINE": "1", + "PYTHONUNBUFFERED": "1", + # t*=0 => full native replay: EVERY node dispatches live in profiling. + } + + mock = subprocess.Popen( + [str(mock_bin), "--port", str(port), "--ttft", str(ttft), "--itl", str(itl)], + env=env, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + try: + assert _wait_for_port(port), "mock server did not come up" + cmd = [ + str(aiperf_bin), + "profile", + "--input-file", + str(_FIX), + "--url", + f"http://127.0.0.1:{port}", + "--endpoint-type", + "chat", + "--model", + _MODEL, + "--tokenizer", + "gpt2", + "--num-conversations", + "1", + "--concurrency", + "4", + "--benchmark-duration", + "60", + "--export-level", + "raw", + "--artifact-dir", + str(artifact_dir), + "--random-seed", + "42", + ] + if ignore_eos: + cmd += ["--extra-inputs", "ignore_eos:true"] + proc = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=300) + finally: + mock.terminate() + try: + mock.wait(timeout=10) + except subprocess.TimeoutExpired: + mock.kill() + + assert proc.returncode == 0, ( + f"aiperf profile (start-anchor) exited {proc.returncode}\n" + f"STDERR tail:\n{proc.stderr[-3000:]}" + ) + raw = next(artifact_dir.rglob("profile_export_raw.jsonl"), None) + assert raw is not None, "no profile_export_raw.jsonl produced" + return raw + + +def _profiling_starts(raw: Path) -> dict[int, int]: + """Map ``payload.max_tokens`` -> earliest profiling ``request_start_ns`` (ns). + + ``max_tokens == out`` is the record->node identity key. Full-replay with + ``--num-conversations 1`` dispatches each node once, so each key maps to one + record; ``min`` is defensive against any recycle re-fire. + """ + starts: dict[int, int] = {} + for line in raw.read_text().splitlines(): + if not line.strip(): + continue + obj = orjson.loads(line) + meta = obj.get("metadata", {}) or {} + if meta.get("benchmark_phase") != "profiling": + continue + start = meta.get("request_start_ns") + payload = obj.get("payload", {}) or {} + # The native LlmNode.max_tokens cap is endpoint-mapped to the wire + # token field (max_completion_tokens for chat). + max_tokens = payload.get("max_tokens") or payload.get("max_completion_tokens") + if start is None or max_tokens is None: + continue + prev = starts.get(max_tokens) + if prev is None or start < prev: + starts[max_tokens] = start + return starts + + +def _assert_no_streaming_metrics(raw: Path, label: str) -> tuple[str, str]: + """Regression: an all-``"n"`` graph run surfaces NO fabricated streaming metrics. + + This fixture records only non-streaming (``"n"``) nodes and the run is launched + with the GLOBAL ``--streaming`` flag OFF. The run-level gate is per-record, + so graph workloads do not drop STREAMING_ONLY metrics wholesale (a mixed + graph CAN surface TTFT for its streaming nodes) -- but here ZERO records + stream, so the per-record predicate must exclude every record and no TTFT may + be fabricated. Asserts ``time_to_first_token`` is ABSENT and + ``streamed_request_count`` is absent or exactly 0. Returns + ``(ttft_state, streamed_count_state)`` strings for the report. + """ + summary = raw.parent / "profile_export_aiperf.json" + assert summary.exists(), f"[{label}] no profile_export_aiperf.json beside {raw}" + data = orjson.loads(summary.read_bytes()) + + # Positive control: a valid export always carries request_count, so an empty + # or malformed export cannot let the absence assertions below pass vacuously. + assert "request_count" in data, ( + f"[{label}] request_count missing from export; empty/malformed summary " + f"would make the streaming-absence checks vacuous" + ) + + assert "time_to_first_token" not in data, ( + f"[{label}] time_to_first_token present in an all-'n' (zero-stream) run: " + f"{data.get('time_to_first_token')}" + ) + src = data.get("streamed_request_count") + if src is None: + streamed_state = "absent" + else: + assert src.get("avg") in (0, 0.0), ( + f"[{label}] streamed_request_count={src.get('avg')} in a zero-stream run" + ) + streamed_state = f"present avg={src.get('avg')}" + return "absent", streamed_state + + +@pytest.mark.integration +@pytest.mark.parametrize( + ("ttft", "itl", "label"), + [ + param(500, 5, "fast", id="fast"), + param(500, 60, "slow", id="slow"), + ], +) # fmt: skip +def test_start_anchor_live_timing( + tmp_path: Path, ttft: int, itl: int, label: str +) -> None: + """Start-anchored children fire at CONSTANT parent-dispatch offsets; the tail tracks completion.""" + raw = _run_profile(tmp_path, ttft=ttft, itl=itl, ignore_eos=True) + starts = _profiling_starts(raw) + for out in (_PARENT_OUT, _SPAWN_OUT, _CHAIN_OUT, _TAIL_OUT): + assert out in starts, ( + f"[{label}] no profiling record for out={out}; got {sorted(starts)}" + ) + + # Regression: relaxed gate must NOT fabricate streaming metrics for an + # all-"n" (zero-stream) graph run launched without global --streaming. + ttft_state, streamed_state = _assert_no_streaming_metrics(raw, label) + print( + f"[{label}] no-stream metrics: time_to_first_token={ttft_state} " + f"streamed_request_count={streamed_state}" + ) + + p0 = starts[_PARENT_OUT] + d_spawn = (starts[_SPAWN_OUT] - p0) / 1e9 + d_chain = (starts[_CHAIN_OUT] - p0) / 1e9 + d_tail = (starts[_TAIL_OUT] - p0) / 1e9 + parent_finish = _parent_finish_offset(raw, p0) + print( + f"[{label}] ttft={ttft} itl={itl}: parent_finish={parent_finish:.3f}s " + f"d_spawn={d_spawn:.3f}s d_chain={d_chain:.3f}s d_tail={d_tail:.3f}s" + ) + + # Start-anchored edges gate off the parent's DISPATCH, so their dispatch-to- + # dispatch offset is INVARIANT to server speed (identical fast and slow). + assert abs(d_spawn - 2.5) <= 0.25, f"[{label}] spawn offset {d_spawn:.3f}s != 2.5s" + assert abs(d_chain - 5.0) <= 0.25, f"[{label}] chain offset {d_chain:.3f}s != 5.0s" + + if label == "slow": + # Parent (~12.5s) is verifiably IN FLIGHT when both children dispatch. + assert d_spawn < 12.0 and d_chain < 12.0, ( + f"[{label}] children did not dispatch mid-parent-flight: " + f"d_spawn={d_spawn:.3f} d_chain={d_chain:.3f}" + ) + # END-anchored tail = parent_finish (~12.5s) + recorded 1.0s gap. + assert abs(d_tail - 13.5) <= 1.0, ( + f"[slow] tail at {d_tail:.3f}s, expected ~13.5s" + ) + else: + # END-anchored tail COMPRESSES with a fast server: parent_finish(~1.5)+1.0. + assert abs(d_tail - 2.5) <= 1.0, f"[fast] tail at {d_tail:.3f}s, expected ~2.5s" + # Content fidelity is latency-independent -- assert it on the fast run. + # The subprocess ran with ``--tokenizer gpt2``; pass the run's knob. + content = content_vs_real_trace(raw, _FIX, tokenizer_name="gpt2") + assert content.passed, content.render() + + +def _parent_finish_offset(raw: Path, p0: int) -> float: + """Parent P's completion offset (s) from its own dispatch -- for print/context.""" + for line in raw.read_text().splitlines(): + if not line.strip(): + continue + obj = orjson.loads(line) + meta = obj.get("metadata", {}) or {} + if meta.get("benchmark_phase") != "profiling": + continue + payload = obj.get("payload", {}) or {} + cap = payload.get("max_tokens") or payload.get("max_completion_tokens") + if cap == _PARENT_OUT: + end = meta.get("request_end_ns") + if end is not None: + return (end - p0) / 1e9 + return float("nan") + + +@pytest.mark.integration +def test_start_anchor_content_causality_fidelity(tmp_path: Path) -> None: + """Reconstruction content + causality both PASS at zero server latency. + + :func:`causality_timing_vs_real_trace` reconstructs each END-anchored edge + from its predecessor's DISPATCH + the warped edge delay -- valid only when the + predecessor returns ~instantly (its docstring drives ``--ttft 0 --itl 0``). On + a latency-bearing run the tail's observed dispatch is parent_FINISH + delay, + off by the parent's processing time -- which is exactly what + :func:`test_start_anchor_live_timing` measures DIRECTLY. So the causality + reconstruction is asserted at zero latency (its design point), over the + IDENTICAL graph geometry. Content fidelity is latency-independent and is also + asserted on the fast run above. + """ + raw = _run_profile(tmp_path, ttft=0, itl=0, ignore_eos=False) + # The subprocess ran with ``--tokenizer gpt2``; pass the run's knob. + content = content_vs_real_trace(raw, _FIX, tokenizer_name="gpt2") + causality = causality_timing_vs_real_trace(raw, _FIX, tokenizer_name="gpt2") + assert content.passed, content.render() + assert causality.passed, causality.render() diff --git a/tests/integration/graph/test_weka_cachebust_raw_export.py b/tests/integration/graph/test_weka_cachebust_raw_export.py new file mode 100644 index 0000000000..a06eca3e93 --- /dev/null +++ b/tests/integration/graph/test_weka_cachebust_raw_export.py @@ -0,0 +1,123 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Real-mp assertion that --cache-bust first_turn_prefix stamps the marker. + +Runs the full multiprocess ``aiperf`` stack over a 2-trace weka directory with +``--cache-bust first_turn_prefix --export-level raw`` and verifies on the EXPORTED +raw artifact that the AgentX-format marker ``[rid:<12hex>]\\n\\n`` is: + +* present on the FIRST user message of every dispatched request, +* SHARED across all of one trace instance's dispatches (per-trace value), and +* DISTINCT across the two trace instances, + +plus a ``--cache-bust none`` control proving the stamp is a no-op when off. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import orjson +import pytest + +from tests.harness.utils import AIPerfCLI, AIPerfMockServer + +MULTIGRAPH_DIR = Path(__file__).parent / "fixtures" / "weka_multigraph_dir" +# AgentX marker: ``[rid:<12 hex>]`` followed by a blank line, at content start. +_RID_PREFIX = re.compile(r"^\[rid:([0-9a-f]{12})\]\n\n") + + +def _first_user_contents(raw_lines: list[str]) -> list[tuple[str, str]]: + """Return ``(conversation_id, first_user_message_content)`` per raw record.""" + out: list[tuple[str, str]] = [] + for line in raw_lines: + if not line.strip(): + continue + rec = orjson.loads(line) + conv = rec["metadata"]["conversation_id"] + messages = rec["payload"]["messages"] + user = next((m for m in messages if m.get("role") == "user"), None) + if user is not None and isinstance(user.get("content"), str): + out.append((conv, user["content"])) + return out + + +async def _run(cli, mock, monkeypatch, cache_bust: str): + result = await cli.run( + f""" + aiperf profile \ + --model claude-opus-4-5-20251101 \ + --url {mock.url} \ + --endpoint-type chat \ + --input-file {MULTIGRAPH_DIR} \ + --tokenizer gpt2 \ + --cache-bust {cache_bust} \ + --num-conversations 2 \ + --concurrency 2 \ + --workers-max 2 \ + --export-level raw \ + --ui simple + """, + timeout=300.0, + ) + assert result.exit_code == 0, result.stderr[-2000:] + raw = next(result.artifacts_dir.glob("**/*profile_export_raw.jsonl"), None) + assert raw is not None, ( + "profile_export_raw.jsonl must exist with --export-level raw" + ) + return _first_user_contents(raw.read_text(encoding="utf-8").splitlines()) + + +@pytest.mark.integration +@pytest.mark.asyncio +class TestCacheBustRawExport: + async def test_first_turn_prefix_stamped_per_trace( + self, + cli: AIPerfCLI, + aiperf_mock_server: AIPerfMockServer, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + contents = await _run(cli, aiperf_mock_server, monkeypatch, "first_turn_prefix") + assert contents, "no dispatched user messages in the raw export" + + # Every dispatched first user message carries the marker at content start. + # The cache-bust marker is stamped per TRACE INSTANCE; a subagent now + # carries its own conversation_id (``{instance}::{subagent}``) but SHARES + # its parent trace instance's marker -- so group rids by the trace + # instance (the conversation_id minus any ``::subagent`` suffix). + rid_by_instance: dict[str, set[str]] = {} + for conv, content in contents: + m = _RID_PREFIX.match(content) + assert m is not None, ( + f"cache-bust marker missing/misplaced on {conv}: {content[:80]!r}" + ) + instance = conv.split("::", 1)[0] + rid_by_instance.setdefault(instance, set()).add(m.group(1)) + + # Per trace instance: exactly one rid, shared across the root + every + # subagent dispatch of that instance. + for instance, rids in rid_by_instance.items(): + assert len(rids) == 1, ( + f"{instance} has multiple rids (not per-trace-instance): {rids}" + ) + + # Distinct across the two trace instances. + all_rids = {next(iter(r)) for r in rid_by_instance.values()} + assert len(rid_by_instance) >= 2 and len(all_rids) == len(rid_by_instance), ( + f"rids must be distinct per trace instance: {rid_by_instance}" + ) + + async def test_cache_bust_none_is_noop( + self, + cli: AIPerfCLI, + aiperf_mock_server: AIPerfMockServer, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + contents = await _run(cli, aiperf_mock_server, monkeypatch, "none") + assert contents, "no dispatched user messages in the raw export" + for conv, content in contents: + assert _RID_PREFIX.match(content) is None, ( + f"cache-bust=none must NOT stamp a marker, but {conv} has one: " + f"{content[:60]!r}" + ) diff --git a/tests/integration/graph/test_weka_isl_endpoint_osl_e2e.py b/tests/integration/graph/test_weka_isl_endpoint_osl_e2e.py new file mode 100644 index 0000000000..4d2f7a0615 --- /dev/null +++ b/tests/integration/graph/test_weka_isl_endpoint_osl_e2e.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Real-mp E2E locks for two general weka-export fidelity behaviors on the trie. + +Runs the full multiprocess ``aiperf`` stack over the 2-trace weka multigraph +directory against the in-repo mock server (segment-trie IR default) and asserts, +on the EXPORTED artifacts, two endpoint-agnostic behaviors that the trie path +carries verbatim: + +* **ISL present:** ``input_sequence_length`` appears in the JSON export and + ``token_counts.input`` is populated on every profiling record, tracking the + block-aligned covered-token count of the wire prompt (builtin/o200k tokenization, + modulo small client-side boundary drift). + +* **Non-chat rejected up front:** ``--endpoint-type completions`` over a weka + workload exits non-zero with a clear ``GraphEndpointUnsupportedError`` message at + configure time, NOT a per-request 422. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from tests.harness.utils import AIPerfCLI, AIPerfMockServer + +MULTIGRAPH_DIR = Path(__file__).parent / "fixtures" / "weka_multigraph_dir" + + +async def _run( + cli: AIPerfCLI, + mock: AIPerfMockServer, + monkeypatch, + *extra: str, + assert_success: bool = True, +): + # Pin the seed so the synthesized weka content -- and thus the block-aligned + # ISL multiset asserted below -- is deterministic across runs and store paths. + # Without a pinned seed, unseeded runs draw different random content and the + # exact ISL floor drifts +/-1 token at block boundaries (see the A/B parity + # test, which pins the same seed to prove legacy == unified byte-for-byte). + return await cli.run( + f""" + aiperf profile \ + --model claude-opus-4-5-20251101 \ + --url {mock.url} \ + --input-file {MULTIGRAPH_DIR} \ + --tokenizer builtin \ + --random-seed 1234 \ + --num-conversations 2 \ + --concurrency 2 \ + --workers-max 2 \ + --export-level raw \ + --ui simple \ + {" ".join(extra)} + """, + timeout=300.0, + assert_success=assert_success, + ) + + +@pytest.mark.integration +@pytest.mark.asyncio +class TestWekaIslEndpointOslE2E: + async def test_isl_present_in_exports( + self, + cli: AIPerfCLI, + aiperf_mock_server: AIPerfMockServer, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """ISL is present (block-aligned covered-count) on every record.""" + result = await _run( + cli, aiperf_mock_server, monkeypatch, "--endpoint-type", "chat" + ) + assert result.exit_code == 0, result.stderr[-2000:] + + # Aggregate ISL is present (was entirely absent before the fix). + assert result.json is not None and result.json.input_sequence_length is not None + assert result.json.input_sequence_length.avg is not None + + # Every profiling metric record carries an input token count. + assert result.jsonl is not None and result.jsonl + for rec in result.jsonl: + isl = rec.metrics.get("input_sequence_length") + assert isl is not None, f"record missing ISL: {rec.metadata}" + assert isl.value > 0 + + # The trie IR materializes BLOCK-ALIGNED prompts: a node's prompt is the + # message-unit concatenation of its covered whole blocks + # (``min(len(hash_ids), in // block_size)`` of them, block_size 64), with + # the recorded ``in % block_size`` partial tail deliberately excluded (see + # ``_assemble_messages`` / ``compute_turn_block_geometry`` and the + # ``_assert_isl`` covered-count gate). So the recovered ISLs are the + # block-aligned covered-token counts, NOT the raw recorded ``in``: + # turn-0 (in=180/200, 2 covered blocks) -> 128; the deeper turns cover + # 3 or 4 blocks (~192/~256) modulo small builtin (o200k) decode/re-encode + # boundary drift. The minimum is the 2-block floor (128), never the raw 180. + isls = sorted( + rec.metrics["input_sequence_length"].value for rec in result.jsonl + ) + assert min(isls) == 128 # 2 covered blocks * block_size (64), no partial tail + assert ( + max(isls) <= 265 + ) # largest is the 4-block prompt (~256) + tokenizer drift + + async def test_non_chat_endpoint_rejected_up_front( + self, + cli: AIPerfCLI, + aiperf_mock_server: AIPerfMockServer, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """A completions endpoint is rejected at configure time, not per-request.""" + result = await _run( + cli, + aiperf_mock_server, + monkeypatch, + "--endpoint-type", + "completions", + assert_success=False, + ) + assert result.exit_code != 0 + combined = (result.stderr + result.stdout + (result.log or "")).lower() + assert ( + "wekaendpointunsupported" in combined or "only supports a chat" in combined + ), f"expected up-front endpoint rejection, got: {combined[-1500:]}" + # It must fail UP FRONT at configure, not via a per-request 422 (the mock's + # "Field required: prompt" rejection of a chat body at /v1/completions). + assert "field required" not in combined, ( + "non-chat endpoint reached per-request dispatch (422) instead of the " + "configure-time guard" + ) diff --git a/tests/integration/graph/test_weka_sticky_placement.py b/tests/integration/graph/test_weka_sticky_placement.py new file mode 100644 index 0000000000..3e1d0fb578 --- /dev/null +++ b/tests/integration/graph/test_weka_sticky_placement.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-trajectory sticky placement over the full multiprocess stack. + +The live A/B for sticky routing: a real `aiperf profile` run +over the 2-trace weka multigraph directory against the in-repo mock server, with +two workers, must place ALL of a trajectory's credits on ONE worker. + +Graph credits mint `turn_index` per node and key their router session on the +instance `trace_id`; without the session the router would scatter a trace's +nodes across workers (least-loaded per credit). The instance session pins +them; the decisive proof is worker consolidation grouped by the template +`conversation_id` (one instance per template in this single-pass run), which +accounting-only assertions cannot see. + +This is also the non-regression gate for sticky routing: the weka corpus +still runs clean and still produces ISL-bearing records under sticky routing. +""" + +from __future__ import annotations + +from collections import defaultdict +from pathlib import Path + +import pytest + +from tests.harness.utils import AIPerfCLI, AIPerfMockServer + +MULTIGRAPH_DIR = Path(__file__).parent / "fixtures" / "weka_multigraph_dir" + + +@pytest.mark.integration +@pytest.mark.asyncio +class TestWekaStickyPlacement: + async def test_each_trace_pins_to_one_worker( + self, + cli: AIPerfCLI, + aiperf_mock_server: AIPerfMockServer, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + result = await cli.run( + f""" + aiperf profile \ + --model claude-opus-4-5-20251101 \ + --url {aiperf_mock_server.url} \ + --input-file {MULTIGRAPH_DIR} \ + --endpoint-type chat \ + --tokenizer builtin \ + --random-seed 1234 \ + --num-conversations 2 \ + --concurrency 2 \ + --workers-max 2 \ + --export-level raw \ + --ui simple + """, + timeout=300.0, + assert_success=True, + ) + assert result.exit_code == 0, result.stderr[-2000:] + + raw = result.raw_records + assert raw is not None and raw, "no raw records exported" + + # Group every request by its trajectory template id (graph credits + # stamp conversation_id = the nonce-less trajectory template; this + # single-pass run mints exactly one instance per template). + by_trace: dict[str, set[str]] = defaultdict(set) + multi_node_traces = defaultdict(int) + for rec in raw: + conv = rec.metadata.conversation_id + worker = rec.metadata.worker_id + assert conv is not None, f"record missing conversation_id: {rec.metadata}" + assert worker is not None, f"record missing worker_id: {rec.metadata}" + by_trace[conv].add(worker) + multi_node_traces[conv] += 1 + + # THE sticky invariant: no trace instance spans more than one worker. + scattered = {c: w for c, w in by_trace.items() if len(w) > 1} + assert not scattered, f"traces scattered across workers: {scattered}" + + # The check is non-vacuous: at least one trace issued multiple credits + # (so single-worker placement is a real consolidation, not one request). + assert any(n > 1 for n in multi_node_traces.values()), ( + f"expected a multi-node trace to make consolidation meaningful, " + f"got per-trace request counts {dict(multi_node_traces)}" + ) + + async def test_weka_still_runs_clean_under_sticky( + self, + cli: AIPerfCLI, + aiperf_mock_server: AIPerfMockServer, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Non-regression: sticky routing does not break weka ISL/records.""" + + result = await cli.run( + f""" + aiperf profile \ + --model claude-opus-4-5-20251101 \ + --url {aiperf_mock_server.url} \ + --input-file {MULTIGRAPH_DIR} \ + --endpoint-type chat \ + --tokenizer builtin \ + --random-seed 1234 \ + --num-conversations 2 \ + --concurrency 2 \ + --workers-max 2 \ + --ui simple + """, + timeout=300.0, + assert_success=True, + ) + assert result.exit_code == 0, result.stderr[-2000:] + assert result.json is not None + assert result.json.input_sequence_length is not None + assert result.json.input_sequence_length.avg is not None + assert result.jsonl is not None and result.jsonl + for rec in result.jsonl: + isl = rec.metrics.get("input_sequence_length") + assert isl is not None and isl.value > 0 diff --git a/tests/integration/graph/test_weka_trace_fidelity_real.py b/tests/integration/graph/test_weka_trace_fidelity_real.py new file mode 100644 index 0000000000..82dccdaee2 --- /dev/null +++ b/tests/integration/graph/test_weka_trace_fidelity_real.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""REAL empirical proof: a live trie ``--export-level raw`` run matches the recorded trace. + +This is the Task-7 acceptance gate's real-run half. It drives an ACTUAL +``aiperf profile --export-level raw`` of the weka segment-trie IR path (a +subprocess run against the in-repo mock server) over the +subagent fixture, then runs the offline fidelity reports +(:mod:`tools.weka_trace_fidelity`) on the produced ``profile_export_raw.jsonl`` and +asserts BOTH: + +* :func:`content_vs_real_trace` (criterion 2) -- each dispatched record's + reconstructed prompt equals what the recorded trace prescribes, and +* :func:`causality_timing_vs_real_trace` (criterion 3) -- the dispatch causal order + and relative inter-request timing match the recorded trace. + +It lives in the INTEGRATION lane (not component) on purpose: the report functions +rebuild the trie graph in-process via ``build_trie_graph``, so the in-process +tokenizer must be the SAME real ``gpt2`` the subprocess uses. The +component-integration package patches ``Tokenizer.from_pretrained`` to a +FakeTokenizer (which the subprocess would NOT see), so the reconstruction would +diverge there; the integration package uses the real tokenizer end-to-end. +""" + +from __future__ import annotations + +import os +import socket +import subprocess +import time +from pathlib import Path + +import pytest + +from tools.weka_trace_fidelity import ( + causality_timing_vs_real_trace, + content_vs_real_trace, + load_raw_export, +) + +_REPO = Path(__file__).resolve().parents[3] +_FIX = _REPO / "tests" / "unit" / "graph" / "fixtures" / "weka_subagent.json" +_MODEL = "claude-opus-4-5-20251101" + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _wait_for_port(port: int, timeout_s: float = 15.0) -> bool: + deadline = time.time() + timeout_s + while time.time() < deadline: + with socket.socket() as s: + if s.connect_ex(("127.0.0.1", port)) == 0: + return True + time.sleep(0.1) + return False + + +@pytest.mark.integration +def test_real_trie_raw_export_matches_real_trace(tmp_path: Path) -> None: + """A live trie raw-export run PASSES content + causality/timing vs the real trace.""" + venv = _REPO / ".venv" / "bin" + mock_bin = venv / "aiperf-mock-server" + aiperf_bin = venv / "aiperf" + if not mock_bin.exists() or not aiperf_bin.exists(): + pytest.skip("aiperf / aiperf-mock-server not installed in .venv") + + port = _free_port() + artifact_dir = tmp_path / "artifacts" + env = { + **os.environ, + "NO_PROXY": "127.0.0.1,localhost", + "HF_HUB_OFFLINE": "1", + "PYTHONUNBUFFERED": "1", + } + + mock = subprocess.Popen( + [str(mock_bin), "--port", str(port), "--ttft", "0", "--itl", "0"], + env=env, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + try: + assert _wait_for_port(port), "mock server did not come up" + cmd = [ + str(aiperf_bin), + "profile", + "--input-file", + str(_FIX), + "--url", + f"http://127.0.0.1:{port}", + "--endpoint-type", + "chat", + "--model", + _MODEL, + "--tokenizer", + "gpt2", + "--num-conversations", + "1", + "--concurrency", + "4", + "--benchmark-duration", + "30", + "--export-level", + "raw", + "--artifact-dir", + str(artifact_dir), + "--random-seed", + "42", + ] + proc = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=300) + finally: + mock.terminate() + try: + mock.wait(timeout=10) + except subprocess.TimeoutExpired: + mock.kill() + + assert proc.returncode == 0, ( + f"aiperf profile (trie path) exited {proc.returncode}\n" + f"STDERR tail:\n{proc.stderr[-3000:]}" + ) + + raw = next(artifact_dir.rglob("profile_export_raw.jsonl"), None) + assert raw is not None, "no profile_export_raw.jsonl produced" + records = load_raw_export(raw) + profiling = [r for r in records if r.phase == "profiling"] + assert profiling, "no profiling records dispatched on the trie path" + + # The subprocess ran with ``--tokenizer gpt2``; the tool defaults to the + # bare live-run builtin tokenizer, so the run's knob is passed explicitly. + content = content_vs_real_trace(raw, _FIX, tokenizer_name="gpt2") + causality = causality_timing_vs_real_trace(raw, _FIX, tokenizer_name="gpt2") + assert content.passed, content.render() + assert causality.passed, causality.render() diff --git a/tests/integration/graph/test_weka_unified_store_ab.py b/tests/integration/graph/test_weka_unified_store_ab.py new file mode 100644 index 0000000000..dca0a91c13 --- /dev/null +++ b/tests/integration/graph/test_weka_unified_store_ab.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Real-mp E2E: the unified segment store drives a full profile run. + +Runs the full multiprocess ``aiperf`` stack (subprocess ``python -m aiperf +profile`` + live mock server + real Worker + HTTP) over the 2-trace weka +multigraph directory with a pinned seed and asserts the unified-store Worker +path actually dispatches and produces per-record ISL metrics. + +The unified store is the SOLE trie store shape (the legacy split +segment+delta stores are retired); the streaming-vs-eager unified store A/B +lives in ``tests/unit/graph/test_hf_streaming_trie_stores.py:: +test_streaming_unified_store_byte_matches_eager_interned``. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from tests.harness.utils import AIPerfCLI, AIPerfMockServer + +MULTIGRAPH_DIR = Path(__file__).parent / "fixtures" / "weka_multigraph_dir" + +# Pinned seed so the synthesized weka content is deterministic +# (the weka content seed derives from run.random_seed). +_SEED = "1234" + + +async def _run( + cli: AIPerfCLI, + mock: AIPerfMockServer, + monkeypatch: pytest.MonkeyPatch, +): + """Run one full profile pass over the weka multigraph directory.""" + return await cli.run( + f""" + aiperf profile \ + --model claude-opus-4-5-20251101 \ + --url {mock.url} \ + --endpoint-type chat \ + --input-file {MULTIGRAPH_DIR} \ + --tokenizer builtin \ + --random-seed {_SEED} \ + --num-conversations 2 \ + --concurrency 2 \ + --workers-max 2 \ + --export-level raw \ + --ui simple + """, + timeout=300.0, + assert_success=True, + ) + + +@pytest.mark.integration +@pytest.mark.asyncio +class TestWekaUnifiedStoreAB: + async def test_unified_store_runs_end_to_end( + self, + cli: AIPerfCLI, + aiperf_mock_server: AIPerfMockServer, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """The unified Worker path actually dispatches and produces records.""" + result = await _run(cli, aiperf_mock_server, monkeypatch) + assert result.exit_code == 0, result.stderr[-2000:] + + assert result.request_count > 0 + assert result.jsonl is not None and result.jsonl, "no profiling records" + for rec in result.jsonl: + isl = rec.metrics.get("input_sequence_length") + assert isl is not None, f"record missing ISL: {rec.metadata}" + assert isl.value > 0 diff --git a/tests/integration/test_session_routing_raw_export.py b/tests/integration/test_session_routing_raw_export.py new file mode 100644 index 0000000000..a5d16e7fe5 --- /dev/null +++ b/tests/integration/test_session_routing_raw_export.py @@ -0,0 +1,188 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end raw-export proof that every ``--session-routing`` mode reaches the wire. + +Runs a real ``aiperf profile`` subprocess against the in-repo mock server with +``--export-level raw`` for each session-routing mode, then reads the exported +wire payloads / request headers and asserts the per-mode contract: + +- ``dynamo_headers``: ``X-Dynamo-Session-ID`` on every request equals the + session's ``x_correlation_id``; no parent header on root sessions; body is + untouched (``"nvext" not in payload``). +- ``dynamo_nvext``: ``nvext.session_control`` carries ``bind`` (+ timeout) on + every non-final turn and ``close`` (no timeout) on the final turn, with one + stable ``session_id == x_correlation_id`` across the session. +- ``smg_routing_key``: ``X-SMG-Routing-Key`` equals ``x_correlation_id``. +- ``session_id_header`` (``--session-routing-opt header_name=X-Affinity``): the + configured header equals ``x_correlation_id``. +""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Callable, Sequence + +import pytest +from pytest import param + +from aiperf.common.models import RawRecordInfo +from tests.harness.utils import AIPerfCLI, AIPerfMockServer + +_TOKENIZER = "openai/gpt-oss-120b" # pre-cached + offline in integration conftest + +_NUM_SESSIONS = 2 +_TURNS_PER_SESSION = 3 +# Non-default (plugin default is 300) so the assertions prove the CLI -> +# config -> worker numeric plumb actually carries this value to the wire, +# rather than coinciding with the default. +_TIMEOUT_SECONDS = 123 + + +def _build_cmd(url: str, *, mode: str, opts: Sequence[str]) -> str: + opt_flags = " ".join(f"--session-routing-opt {kv}" for kv in opts) + return f""" + aiperf profile \ + --model {_TOKENIZER} \ + --url {url} \ + --endpoint-type chat \ + --num-sessions {_NUM_SESSIONS} \ + --session-turns-mean {_TURNS_PER_SESSION} \ + --session-turns-stddev 0 \ + --random-seed 42 \ + --workers-max 1 \ + --session-routing {mode} \ + {opt_flags} \ + --export-level raw \ + --ui simple + """ + + +async def _records_by_session( + cli: AIPerfCLI, url: str, *, mode: str, opts: Sequence[str] +) -> dict[str, list[RawRecordInfo]]: + """Run a benchmark and return each session's raw records ordered by + turn_index, keyed by X-Correlation-ID.""" + result = await cli.run(_build_cmd(url, mode=mode, opts=opts), timeout=300.0) + + records = list(result.raw_records or []) + assert records, f"no raw records\n{(result.log or '')[-1500:]}" + + grouped: dict[str, list[RawRecordInfo]] = defaultdict(list) + for rec in records: + grouped[rec.metadata.x_correlation_id].append(rec) + + assert len(grouped) == _NUM_SESSIONS, ( + f"expected {_NUM_SESSIONS} sessions, got {len(grouped)}" + ) + out: dict[str, list[RawRecordInfo]] = {} + for xcorr, recs in grouped.items(): + assert len(recs) == _TURNS_PER_SESSION, ( + f"session {xcorr}: expected {_TURNS_PER_SESSION} turns, got {len(recs)}" + ) + recs.sort(key=lambda r: r.metadata.turn_index) + out[xcorr] = recs + return out + + +def _verify_dynamo_headers(by_session: dict[str, list[RawRecordInfo]]) -> None: + for xcorr, recs in by_session.items(): + for rec in recs: + headers = rec.request_headers or {} + assert headers.get("X-Dynamo-Session-ID") == xcorr, ( + f"session {xcorr} turn {rec.metadata.turn_index}: headers={headers}" + ) + # Root sessions have no parent, so the parent header must be absent. + assert "X-Dynamo-Parent-Session-ID" not in headers, ( + f"session {xcorr}: root session emitted a parent header; {headers}" + ) + assert "nvext" not in rec.payload, ( + f"session {xcorr}: header-mode must not mutate the body" + ) + + +def _verify_smg_routing_key(by_session: dict[str, list[RawRecordInfo]]) -> None: + for xcorr, recs in by_session.items(): + for rec in recs: + headers = rec.request_headers or {} + assert headers.get("X-SMG-Routing-Key") == xcorr, ( + f"session {xcorr} turn {rec.metadata.turn_index}: headers={headers}" + ) + assert "nvext" not in rec.payload, ( + f"session {xcorr}: header-mode must not mutate the body" + ) + + +def _verify_session_id_header(by_session: dict[str, list[RawRecordInfo]]) -> None: + for xcorr, recs in by_session.items(): + for rec in recs: + headers = rec.request_headers or {} + assert headers.get("X-Affinity") == xcorr, ( + f"session {xcorr} turn {rec.metadata.turn_index}: headers={headers}" + ) + assert "nvext" not in rec.payload, ( + f"session {xcorr}: header-mode must not mutate the body" + ) + + +def _verify_dynamo_nvext(by_session: dict[str, list[RawRecordInfo]]) -> None: + for xcorr, recs in by_session.items(): + scs = [] + for rec in recs: + sc = rec.payload.get("nvext", {}).get("session_control") + assert sc is not None, ( + f"session {xcorr} turn {rec.metadata.turn_index}: " + "every request must carry nvext.session_control" + ) + scs.append(sc) + + actions = [sc.get("action") for sc in scs] + assert all(a == "bind" for a in actions[:-1]), ( + f"session {xcorr}: non-final turns must bind; {actions}" + ) + assert actions[-1] == "close", f"session {xcorr}: actions={actions}" + assert "open" not in actions, f"session {xcorr}: emitted open; {actions}" + + # Every non-final 'bind' carries the timeout; 'close' does not. + for sc in scs[:-1]: + assert sc["timeout"] == _TIMEOUT_SECONDS, f"session {xcorr}: {sc}" + assert "timeout" not in scs[-1], ( + f"session {xcorr}: close carried timeout; {scs[-1]}" + ) + + # One stable session_id == the X-Correlation-ID, on every turn. + assert {sc["session_id"] for sc in scs} == {xcorr}, ( + f"session {xcorr}: session_id drift; {[sc.get('session_id') for sc in scs]}" + ) + + +_VERIFIERS: dict[str, Callable[[dict[str, list[RawRecordInfo]]], None]] = { + "dynamo_headers": _verify_dynamo_headers, + "dynamo_nvext": _verify_dynamo_nvext, + "smg_routing_key": _verify_smg_routing_key, + "session_id_header": _verify_session_id_header, +} + + +@pytest.mark.integration +@pytest.mark.asyncio +@pytest.mark.parametrize( + "mode, opts", + [ + param("dynamo_headers", (), id="dynamo_headers"), + param("dynamo_nvext", (f"timeout_seconds={_TIMEOUT_SECONDS}",), id="dynamo_nvext"), + param("smg_routing_key", (), id="smg_routing_key"), + param("session_id_header", ("header_name=X-Affinity",), id="session_id_header"), + ], +) # fmt: skip +async def test_session_routing_mode_reaches_wire( + cli: AIPerfCLI, + aiperf_mock_server: AIPerfMockServer, + mode: str, + opts: tuple[str, ...], +): + """Each session-routing mode stamps its per-session identity on the wire and + it survives to the raw export exactly as the plugin specifies.""" + by_session = await _records_by_session( + cli, aiperf_mock_server.url, mode=mode, opts=opts + ) + _VERIFIERS[mode](by_session) diff --git a/tests/unit/cli_commands/test_dynamo_trace_report.py b/tests/unit/cli_commands/test_dynamo_trace_report.py new file mode 100644 index 0000000000..b4fcd9d6fc --- /dev/null +++ b/tests/unit/cli_commands/test_dynamo_trace_report.py @@ -0,0 +1,260 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""`aiperf dynamo trace-report` against the REAL `dynamo.request.trace.v1` schema. + +The original report was written against a nonexistent workflow/program schema +(WK1): it passed a `workflow_id=` kwarg `iter_trace_records` does not accept +and read `agent_context.workflow_id`/`.program_id` fields `AgentContext` does +not have -- dead on arrival for every trace the in-repo reader parses. These +tests drive `aggregate_by_session` end-to-end over real minimal trace files: +session grouping, parent linkage, replay-only skip accounting, the reader-level +session filter, and nearest-rank percentile math. +""" + +from __future__ import annotations + +import gzip +from pathlib import Path +from typing import Any + +import orjson +import pytest +from pytest import param + +from aiperf.cli_commands.dynamo_trace_report import ( + _format_csv, + _format_json, + _percentiles, + aggregate_by_session, +) + + +def _request_end( + *, + session_id: str | None, + parent_session_id: str | None = None, + model: str | None = "test-model", + ts_ms: int = 1_000, + ttft_ms: float | None = None, + input_tokens: int | None = None, + replay_hashes: list[int] | None = None, +) -> dict[str, Any]: + """A minimal bare `request_end` record the reader accepts. + + `session_id=None` omits `agent_context` entirely (a replay-only record). + """ + request: dict[str, Any] = {"request_id": f"r-{ts_ms}"} + if model is not None: + request["model"] = model + if ttft_ms is not None: + request["ttft_ms"] = ttft_ms + if input_tokens is not None: + request["input_tokens"] = input_tokens + if replay_hashes is not None: + request["replay"] = { + "trace_block_size": 16, + "input_length": 32, + "input_sequence_hashes": replay_hashes, + } + record: dict[str, Any] = { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts_ms, + "request": request, + } + if session_id is not None: + ctx: dict[str, Any] = {"session_id": session_id} + if parent_session_id is not None: + ctx["parent_session_id"] = parent_session_id + record["agent_context"] = ctx + return record + + +def _write_trace(tmp_path: Path, records: list[dict[str, Any]]) -> Path: + trace = tmp_path / "trace.jsonl" + trace.write_bytes(b"\n".join(orjson.dumps(r) for r in records) + b"\n") + return trace + + +def test_aggregate_by_session_groups_and_links_parent(tmp_path: Path) -> None: + """Minimal valid request_end records -> per-session rows with parent links.""" + trace = _write_trace( + tmp_path, + [ + _request_end(session_id="s-root", ts_ms=1_000, ttft_ms=5.0), + _request_end(session_id="s-root", ts_ms=3_000, ttft_ms=7.0), + _request_end( + session_id="s-child", + parent_session_id="s-root", + ts_ms=2_000, + ttft_ms=9.0, + replay_hashes=[11, 22], + ), + ], + ) + report = aggregate_by_session(trace) + + assert report.skipped_no_agent_context == 0 + assert [row["session_id"] for row in report.rows] == ["s-child", "s-root"] + + child, root = report.rows + assert root["parent_session_id"] is None + assert root["request_count"] == 2 + assert root["child_session_count"] == 1 + assert root["time_range_ms"] == [1_000, 3_000] + assert root["models"] == ["test-model"] + assert root["metrics"]["ttft_ms"]["count"] == 2.0 + assert root["metrics"]["ttft_ms"]["mean"] == pytest.approx(6.0) + + assert child["parent_session_id"] == "s-root" + assert child["parent_session_id_conflict"] is False + assert child["child_session_count"] == 0 + assert child["replay_records"] == 1 + assert child["unique_replay_hashes"] == 2 + + +def test_duplicate_record_across_dual_sink_files_folded_once(tmp_path: Path) -> None: + """Dynamo's dual file sinks can write the SAME record into two files of one + capture dir (`discover_segments` reads ALL *.jsonl + *.jsonl.gz); the + aggregate must fold it once and count the duplicate, matching the chain + parser's ("request_end", session_id, request_id) dedup identity.""" + rec = _request_end(session_id="s-a", ts_ms=1_000, ttft_ms=5.0) + envelope = orjson.dumps({"timestamp": 5, "event": rec}) + b"\n" + (tmp_path / "trace.jsonl").write_bytes(envelope) + with gzip.open(tmp_path / "trace.000000.jsonl.gz", "wb") as f: + f.write(envelope) + + report = aggregate_by_session(tmp_path) + assert [row["session_id"] for row in report.rows] == ["s-a"] + assert report.rows[0]["request_count"] == 1 + assert report.rows[0]["metrics"]["ttft_ms"]["count"] == 1.0 + assert report.duplicate_records == 1 + assert orjson.loads(_format_json(report))["duplicate_records"] == 1 + + +def test_distinct_request_ids_are_not_deduped(tmp_path: Path) -> None: + """Two records of one session with different request_ids both fold in.""" + trace = _write_trace( + tmp_path, + [ + _request_end(session_id="s-a", ts_ms=1_000), + _request_end(session_id="s-a", ts_ms=2_000), + ], + ) + report = aggregate_by_session(trace) + assert report.duplicate_records == 0 + assert report.rows[0]["request_count"] == 2 + + +def test_replay_only_records_skipped_with_counter(tmp_path: Path) -> None: + """Records without agent_context produce no rows but are counted.""" + trace = _write_trace( + tmp_path, + [ + _request_end(session_id=None, ts_ms=1_000), + _request_end(session_id=None, ts_ms=2_000), + _request_end(session_id="s-live", ts_ms=3_000), + ], + ) + report = aggregate_by_session(trace) + assert report.skipped_no_agent_context == 2 + assert [row["session_id"] for row in report.rows] == ["s-live"] + + +def test_session_filter_uses_reader_session_id_param(tmp_path: Path) -> None: + """The --session-id filter is pushed down to the reader's parse-time filter.""" + trace = _write_trace( + tmp_path, + [ + _request_end(session_id="s-a", ts_ms=1_000), + _request_end(session_id="s-b", ts_ms=2_000), + ], + ) + report = aggregate_by_session(trace, session_id="s-a") + assert [row["session_id"] for row in report.rows] == ["s-a"] + assert report.rows[0]["request_count"] == 1 + + +def test_model_none_excluded_from_model_set(tmp_path: Path) -> None: + """`request.model` is optional in the schema; None must not enter the model set.""" + trace = _write_trace( + tmp_path, + [ + _request_end(session_id="s-a", model=None, ts_ms=1_000), + _request_end(session_id="s-a", model="test-model", ts_ms=2_000), + ], + ) + report = aggregate_by_session(trace) + assert report.rows[0]["models"] == ["test-model"] + + +def test_limit_stops_new_sessions_but_keeps_existing(tmp_path: Path) -> None: + trace = _write_trace( + tmp_path, + [ + _request_end(session_id="s-a", ts_ms=1_000), + _request_end(session_id="s-b", ts_ms=2_000), + _request_end(session_id="s-a", ts_ms=3_000), + ], + ) + report = aggregate_by_session(trace, limit=1) + assert [row["session_id"] for row in report.rows] == ["s-a"] + assert report.rows[0]["request_count"] == 2 + + +@pytest.mark.parametrize( + "values, stat, expected", + [ + param([1.0, 2.0, 3.0, 4.0], "p50", 2.0, id="p50_even_n_nearest_rank"), + param([1.0, 2.0, 3.0, 4.0], "p90", 4.0, id="p90_even_n"), + param([1.0, 2.0, 3.0], "p50", 2.0, id="p50_odd_n"), + param([1.0, 2.0, 3.0, 4.0, 5.0], "p99", 5.0, id="p99_top_rank"), + param([7.0], "p50", 7.0, id="single_value_all_stats"), + param([1.0, 2.0, 3.0, 4.0], "min", 1.0, id="min"), + param([1.0, 2.0, 3.0, 4.0], "max", 4.0, id="max"), + param([1.0, 2.0, 3.0, 4.0], "mean", 2.5, id="mean"), + param([1.0, 2.0, 3.0, 4.0], "count", 4.0, id="count"), + ], +) # fmt: skip +def test_percentiles_nearest_rank( + values: list[float], stat: str, expected: float +) -> None: + """Nearest-rank percentiles: rank = ceil(p/100 * n), 1-based. + + The original implementation used int(p/100 * n) as a 0-based index, which + is off by one (p50 of [1,2,3,4] returned 3 instead of 2). + """ + assert _percentiles(values)[stat] == pytest.approx(expected) + + +def test_percentiles_empty_returns_empty_dict() -> None: + assert _percentiles([]) == {} + + +def test_format_json_envelope_carries_skip_counter(tmp_path: Path) -> None: + trace = _write_trace( + tmp_path, + [ + _request_end(session_id="s-a", ts_ms=1_000), + _request_end(session_id=None, ts_ms=2_000), + ], + ) + report = aggregate_by_session(trace) + payload = orjson.loads(_format_json(report)) + assert payload["skipped_no_agent_context"] == 1 + assert [s["session_id"] for s in payload["sessions"]] == ["s-a"] + + +def test_format_csv_headers_and_rows(tmp_path: Path) -> None: + trace = _write_trace( + tmp_path, + [_request_end(session_id="s-a", ts_ms=1_000, ttft_ms=5.0, input_tokens=10)], + ) + report = aggregate_by_session(trace) + lines = _format_csv(report.rows).strip().splitlines() + assert len(lines) == 2 + header = lines[0].split(",") + assert header[0] == "session_id" + assert "parent_session_id" in header + assert "ttft_ms_p50" in header + assert lines[1].startswith("s-a,") diff --git a/tests/unit/cli_runner/test_aggregate_submission_writer.py b/tests/unit/cli_runner/test_aggregate_submission_writer.py new file mode 100644 index 0000000000..c524ced302 --- /dev/null +++ b/tests/unit/cli_runner/test_aggregate_submission_writer.py @@ -0,0 +1,398 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Writer side of the multi-run AGGREGATE scenario-submission verdict. + +The reader (``AggregateConfidenceJsonExporter._build_submission_metadata``) was +ported but the WRITER was not, so multi-run aggregate exports emitted an EMPTY +submission verdict (no ``scenario`` / ``submission_valid`` keys). These tests +cover ``cli_runner._aggregate._stamp_scenario_submission_metadata``: it stamps +the carrier keys the reader pops, and the writer + reader round-trip produces +the right aggregate verdict (including cross-run overflow rate, cancellation, +and the ``--unsafe-override`` lock fold). + +Uses real ``BenchmarkPlan`` / ``BenchmarkConfig`` / ``RunResult`` / +``AggregateResult`` objects (no MagicMock) so the full plumbing is exercised. +""" + +from __future__ import annotations + +import orjson + +from aiperf.cli_runner._aggregate import ( + SCENARIO_SUBMISSION_CARRIER_KEYS, + _stamp_scenario_submission_metadata, + _sum_runtime_response_counts, + strip_scenario_submission_carrier_keys, +) +from aiperf.common.models.export_models import JsonMetricResult +from aiperf.config import BenchmarkConfig, BenchmarkPlan +from aiperf.exporters.aggregate import ( + AggregateConfidenceCsvExporter, + AggregateConfidenceJsonExporter, + AggregateExporterConfig, +) +from aiperf.orchestrator.aggregation.base import AggregateResult +from aiperf.orchestrator.models import RunResult + +_MINIMAL_CONFIG = { + "models": ["test-model"], + "endpoint": { + "urls": ["http://localhost:8000/v1/chat/completions"], + "wait_for_model_timeout": 0, + }, + "datasets": [ + { + "name": "default", + "type": "synthetic", + "entries": 100, + "prompts": {"isl": 128, "osl": 64}, + } + ], + "phases": [ + { + "name": "profiling", + "type": "concurrency", + "requests": 100, + "concurrency": 1, + } + ], +} + + +def _make_plan(*, scenario: str | None, unsafe_override: bool = False) -> BenchmarkPlan: + """A non-sweep two-trial plan, optionally carrying a scenario lock.""" + cfg_dict = dict(_MINIMAL_CONFIG) + if scenario is not None: + cfg_dict = { + **cfg_dict, + "scenario": scenario, + "unsafe_override": unsafe_override, + } + cfg = BenchmarkConfig.model_validate(cfg_dict) + return BenchmarkPlan(configs=[cfg], trials=2, confidence_level=0.95) + + +def _count(value: float) -> JsonMetricResult: + """A count-style ``JsonMetricResult`` whose per-run total is ``avg``.""" + return JsonMetricResult(unit="requests", avg=value, min=value, max=value) + + +def _run( + *, + label: str, + artifacts_path, + request_count: float, + overflow: float = 0.0, + errors: float = 0.0, + was_cancelled: bool = False, + success: bool = True, + submission_valid: bool | None = True, + submission_invalid_reasons: list[str] | None = None, +) -> RunResult: + """A ``RunResult`` carrying the real lock-only verdict + a per-run JSON. + + ``submission_valid`` / ``submission_invalid_reasons`` mirror the + ``ScenarioOutcome`` the orchestrator carries onto each run from + ``run.resolved.scenario_outcome`` (defaults to a clean ``(True, [])`` lock). + The per-run JSON records ``was_cancelled`` so the writer's cross-run + cancellation fold (which reads it back) is exercised. + """ + artifacts_path.mkdir(parents=True, exist_ok=True) + (artifacts_path / "profile_export_aiperf.json").write_bytes( + orjson.dumps({"was_cancelled": was_cancelled}) + ) + return RunResult( + label=label, + success=success, + artifacts_path=artifacts_path, + summary_metrics={ + "request_count": _count(request_count), + "error_request_count": _count(errors), + "context_overflow_count": _count(overflow), + }, + submission_valid=submission_valid, + submission_invalid_reasons=list(submission_invalid_reasons or []), + ) + + +def _aggregate(metadata: dict) -> AggregateResult: + return AggregateResult( + aggregation_type="confidence", + num_runs=2, + num_successful_runs=2, + failed_runs=[], + metrics={}, + metadata=dict(metadata), + ) + + +def _export_metadata(tmp_path, result: AggregateResult) -> dict: + config = AggregateExporterConfig(result=result, output_dir=tmp_path) + exporter = AggregateConfidenceJsonExporter(config=config) + return orjson.loads(exporter._generate_content())["metadata"] + + +def test_sum_runtime_response_counts_sums_across_runs(tmp_path): + """Cross-run sum = request + error + overflow, with overflow double-counted.""" + runs = [ + _run( + label="r0", + artifacts_path=tmp_path / "r0", + request_count=90, + overflow=2, + errors=3, + ), + _run( + label="r1", + artifacts_path=tmp_path / "r1", + request_count=80, + overflow=4, + errors=1, + ), + ] + total, overflow = _sum_runtime_response_counts(runs) + assert overflow == 6 + assert total == (90 + 3 + 2) + (80 + 1 + 4) + + +def test_writer_stamps_all_carrier_keys_clean_run(tmp_path): + """A clean scenario run stamps a True validator verdict + summed counts.""" + plan = _make_plan(scenario="inferencex-agentx-mvp") + aggregate = _aggregate({"confidence_level": 0.95}) + runs = [ + _run(label="r0", artifacts_path=tmp_path / "r0", request_count=100, overflow=1), + _run(label="r1", artifacts_path=tmp_path / "r1", request_count=100, overflow=0), + ] + _stamp_scenario_submission_metadata(aggregate, runs, plan) + md = aggregate.metadata + assert md["_scenario_name"] == "inferencex-agentx-mvp" + assert md["_validator_submission_valid"] is True + assert md["_validator_submission_invalid_reasons"] == [] + assert md["_total_responses"] == 201 + assert md["_context_overflow_count"] == 1 + assert md["_was_cancelled"] is False + + +def test_writer_reader_roundtrip_clean_run_valid(tmp_path): + """Writer + reader round-trip: a clean 2-run scenario aggregate is valid.""" + plan = _make_plan(scenario="inferencex-agentx-mvp") + aggregate = _aggregate({"confidence_level": 0.95}) + runs = [ + _run( + label="r0", artifacts_path=tmp_path / "r0", request_count=1000, overflow=1 + ), + _run( + label="r1", artifacts_path=tmp_path / "r1", request_count=1000, overflow=0 + ), + ] + _stamp_scenario_submission_metadata(aggregate, runs, plan) + md = _export_metadata(tmp_path / "agg", aggregate) + assert md["scenario"] == "inferencex-agentx-mvp" + assert md["submission_valid"] is True + assert "submission_invalid_reasons" not in md + + +def test_writer_reader_roundtrip_unsafe_override_violation_invalid(tmp_path): + """``--unsafe-override`` WITH a real violation folds invalid across runs. + + The invalid verdict comes from the lock-only ``ScenarioOutcome`` carried on + each run (``submission_valid=False``, reasons ``['unsafe_override']``) -- the + real outcome ``apply_scenario`` produces when violations are downgraded -- NOT + from re-deriving the verdict from the ``unsafe_override`` flag. + """ + plan = _make_plan(scenario="inferencex-agentx-mvp", unsafe_override=True) + aggregate = _aggregate({"confidence_level": 0.95}) + runs = [ + _run( + label="r0", + artifacts_path=tmp_path / "r0", + request_count=100, + submission_valid=False, + submission_invalid_reasons=["unsafe_override"], + ), + _run( + label="r1", + artifacts_path=tmp_path / "r1", + request_count=100, + submission_valid=False, + submission_invalid_reasons=["unsafe_override"], + ), + ] + _stamp_scenario_submission_metadata(aggregate, runs, plan) + md = _export_metadata(tmp_path / "agg", aggregate) + assert md["submission_valid"] is False + assert md["submission_invalid_reasons"] == ["unsafe_override"] + + +def test_writer_reader_roundtrip_clean_override_no_violation_valid(tmp_path): + """C-1 regression: a CLEAN lock carrying ``--unsafe-override`` is VALID. + + ``apply_scenario`` returns ``(True, [])`` for a conforming config even with + ``--unsafe-override`` set (the flag only matters when there are violations to + downgrade). The old writer re-derived the verdict from the flag alone and + falsely stamped ``submission_valid=False`` / ``['unsafe_override']``. With the + real carried outcome on each run, the aggregate matches the single-run path. + """ + plan = _make_plan(scenario="inferencex-agentx-mvp", unsafe_override=True) + aggregate = _aggregate({"confidence_level": 0.95}) + runs = [ + _run( + label="r0", + artifacts_path=tmp_path / "r0", + request_count=100, + submission_valid=True, + submission_invalid_reasons=[], + ), + _run( + label="r1", + artifacts_path=tmp_path / "r1", + request_count=100, + submission_valid=True, + submission_invalid_reasons=[], + ), + ] + _stamp_scenario_submission_metadata(aggregate, runs, plan) + assert aggregate.metadata["_validator_submission_valid"] is True + assert aggregate.metadata["_validator_submission_invalid_reasons"] == [] + md = _export_metadata(tmp_path / "agg", aggregate) + assert md["submission_valid"] is True + assert "submission_invalid_reasons" not in md + + +def test_writer_folds_cancellation_over_failed_run(tmp_path): + """I-1 regression: cancellation folds over ALL runs, not just successful ones. + + A graceful Ctrl+C that completed ZERO requests is classified + ``success=False`` but still wrote its per-run JSON with ``was_cancelled=true``. + Folding only over successful runs would drop it; folding over all runs (as + AgentX does) keeps the aggregate invalid. + """ + plan = _make_plan(scenario="inferencex-agentx-mvp") + aggregate = _aggregate({"confidence_level": 0.95}) + runs = [ + _run(label="r0", artifacts_path=tmp_path / "r0", request_count=100), + _run( + label="r1", + artifacts_path=tmp_path / "r1", + request_count=0, + success=False, + was_cancelled=True, + ), + ] + _stamp_scenario_submission_metadata(aggregate, runs, plan) + assert aggregate.metadata["_was_cancelled"] is True + md = _export_metadata(tmp_path / "agg", aggregate) + assert md["submission_valid"] is False + assert "run_cancelled" in md["submission_invalid_reasons"] + + +def test_writer_reader_roundtrip_cross_run_overflow_rate(tmp_path): + """Cross-run overflow rate > 1% flips the aggregate invalid via summed counts. + + Neither run alone trips it cleanly, but the summed rate + (6 / 200 = 3% > 1%) does, proving the cross-run fold. + """ + plan = _make_plan(scenario="inferencex-agentx-mvp") + aggregate = _aggregate({"confidence_level": 0.95}) + runs = [ + _run(label="r0", artifacts_path=tmp_path / "r0", request_count=97, overflow=3), + _run(label="r1", artifacts_path=tmp_path / "r1", request_count=97, overflow=3), + ] + _stamp_scenario_submission_metadata(aggregate, runs, plan) + assert aggregate.metadata["_total_responses"] == 200 + assert aggregate.metadata["_context_overflow_count"] == 6 + md = _export_metadata(tmp_path / "agg", aggregate) + assert md["submission_valid"] is False + assert "context_overflow_rate_exceeded" in md["submission_invalid_reasons"] + + +def test_writer_reader_roundtrip_cancellation_invalid(tmp_path): + """A cancelled run (read from per-run JSON) flips the aggregate invalid.""" + plan = _make_plan(scenario="inferencex-agentx-mvp") + aggregate = _aggregate({"confidence_level": 0.95}) + runs = [ + _run(label="r0", artifacts_path=tmp_path / "r0", request_count=50), + _run( + label="r1", + artifacts_path=tmp_path / "r1", + request_count=50, + was_cancelled=True, + ), + ] + _stamp_scenario_submission_metadata(aggregate, runs, plan) + assert aggregate.metadata["_was_cancelled"] is True + md = _export_metadata(tmp_path / "agg", aggregate) + assert md["submission_valid"] is False + assert md["submission_invalid_reasons"] == ["run_cancelled"] + + +def test_writer_no_scenario_is_noop(tmp_path): + """No scenario -> no carrier keys stamped -> reader omits submission fields.""" + plan = _make_plan(scenario=None) + aggregate = _aggregate({"confidence_level": 0.95}) + runs = [ + _run(label="r0", artifacts_path=tmp_path / "r0", request_count=10, overflow=9), + _run(label="r1", artifacts_path=tmp_path / "r1", request_count=10, overflow=9), + ] + _stamp_scenario_submission_metadata(aggregate, runs, plan) + for key in ( + "_scenario_name", + "_validator_submission_valid", + "_total_responses", + "_was_cancelled", + ): + assert key not in aggregate.metadata + md = _export_metadata(tmp_path / "agg", aggregate) + assert "scenario" not in md + assert "submission_valid" not in md + + +def test_writer_stamps_exactly_the_carrier_key_constant(tmp_path): + """Drift guard: the stamped key set IS ``SCENARIO_SUBMISSION_CARRIER_KEYS``. + + The CSV exporter strips exactly this constant; a new carrier key stamped + without extending it would leak into user-facing output. + """ + plan = _make_plan(scenario="inferencex-agentx-mvp") + aggregate = _aggregate({}) + runs = [_run(label="r0", artifacts_path=tmp_path / "r0", request_count=10)] + _stamp_scenario_submission_metadata(aggregate, runs, plan) + assert set(aggregate.metadata) == set(SCENARIO_SUBMISSION_CARRIER_KEYS) + + +def test_strip_helper_removes_only_carrier_keys(): + """The shared helper strips every carrier key and nothing else.""" + metadata = {key: "x" for key in SCENARIO_SUBMISSION_CARRIER_KEYS} + metadata["cooldown_seconds"] = 5 + metadata["scenario"] = "kept" + stripped = strip_scenario_submission_carrier_keys(metadata) + assert stripped == {"cooldown_seconds": 5, "scenario": "kept"} + # Input not mutated. + assert set(SCENARIO_SUBMISSION_CARRIER_KEYS) <= set(metadata) + + +def test_csv_exporter_excludes_carrier_keys(tmp_path): + """WK3: the aggregate confidence CSV must not leak the carrier keys. + + The CSV title-cases metadata keys (``_scenario_name`` -> `` Scenario Name``), + so assert on the rendered names as well as the raw keys. + """ + plan = _make_plan(scenario="inferencex-agentx-mvp") + aggregate = _aggregate({"confidence_level": 0.95, "cooldown_seconds": 3}) + runs = [ + _run(label="r0", artifacts_path=tmp_path / "r0", request_count=100, overflow=1), + _run(label="r1", artifacts_path=tmp_path / "r1", request_count=100, overflow=0), + ] + _stamp_scenario_submission_metadata(aggregate, runs, plan) + + exporter = AggregateConfidenceCsvExporter( + config=AggregateExporterConfig(result=aggregate, output_dir=tmp_path / "agg") + ) + content = exporter._generate_content() + + for key in SCENARIO_SUBMISSION_CARRIER_KEYS: + assert key not in content + assert key.replace("_", " ").title() not in content + # Non-carrier metadata still exported. + assert "Cooldown Seconds" in content + assert "Confidence Level" in content diff --git a/tests/unit/cli_runner/test_aggregation_round3.py b/tests/unit/cli_runner/test_aggregation_round3.py index 2bce94e9aa..097ae9ad50 100644 --- a/tests/unit/cli_runner/test_aggregation_round3.py +++ b/tests/unit/cli_runner/test_aggregation_round3.py @@ -1,19 +1,18 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Round-3 multi-run aggregation regression tests. +"""Multi-run aggregation regression tests. -Covers fixes for the round-2 adversarial findings on multi-run/sweep -aggregation correctness: +Covers multi-run/sweep aggregation correctness: - * R2-H4: QMC-collision cells no longer pool into one aggregate dir. - * R2-M5: NaN samples don't poison sibling metric aggregations. - * R2-M7: ``BenchmarkPlan.trials`` and ``MultiRunConfig.num_runs`` caps + * QMC-collision cells never pool into one aggregate dir. + * NaN samples don't poison sibling metric aggregations. + * ``BenchmarkPlan.trials`` and ``MultiRunConfig.num_runs`` caps stay aligned, surfaced via cross-validator. - * R2-M8: ``AIPERF_RAISE_ON_CALLBACK_ERROR`` is read through the + * ``AIPERF_RAISE_ON_CALLBACK_ERROR`` is read through the Pydantic Settings registry (Field on ``_CLIRunnerSettings``) rather than via raw ``os.environ.get``, picking up Pydantic's bool coercion. - * R2-L7: ``zip(plan.configs, plan.variations)`` raises on length + * ``zip(plan.configs, plan.variations)`` raises on length mismatch instead of silently truncating. """ @@ -55,17 +54,18 @@ def _result( # ============================================================================= -# R2-H4: QMC-collision cells get distinct aggregate keys +# QMC-collision cells get distinct aggregate keys # ============================================================================= def test_qmc_collision_cells_get_distinct_aggregate_dirs() -> None: """Sobol over int dims that collides on ``values`` keeps distinct cells. - Repro from the round-2 adversarial report: Sobol over ``lo=1, hi=4, - kind=int, samples=8`` produced 4 unique values from 8 distinct cells. - Pre-fix: one aggregate dir for all 3 cells that mapped to concurrency=3, - inflated ``num_profile_runs=3``. Post-fix: each cell keys distinctly. + Sobol over ``lo=1, hi=4, + kind=int, samples=8`` produces 4 unique values from 8 distinct cells. + Keying aggregate dirs by value would pool all 3 cells that map to + concurrency=3 into one dir and inflate ``num_profile_runs=3``; + each cell must key distinctly. """ # Reproduce the collision: 3 distinct sobol cells map to concurrency=3. results = [ @@ -117,7 +117,7 @@ def test_grid_sweep_distinct_values_still_group_per_label() -> None: # ============================================================================= -# R2-M5: NaN-safe stats +# NaN-safe stats # ============================================================================= @@ -176,7 +176,7 @@ def test_aggregator_nan_in_one_trial_does_not_poison_metric() -> None: # ============================================================================= -# R2-M7: trials cap == num_runs cap (cross-validator) +# trials cap == num_runs cap (cross-validator) # ============================================================================= @@ -221,7 +221,7 @@ def test_multi_run_cooldown_seconds_capped_at_24h() -> None: # ============================================================================= -# R2-M8: RAISE_ON_CALLBACK_ERROR via Pydantic Settings +# RAISE_ON_CALLBACK_ERROR via Pydantic Settings # ============================================================================= @@ -301,7 +301,7 @@ def boom(_completed: CompletedRun) -> None: # ============================================================================= -# R2-L7: strict=True zip surfaces orchestrator config/variation drift +# strict=True zip surfaces orchestrator config/variation drift # ============================================================================= @@ -326,7 +326,7 @@ def test_strict_zip_in_aggregator_raises_on_length_mismatch() -> None: ) assert "zip(plan.configs, plan.variations, strict=False)" not in source, ( "strict=False zip(plan.configs, plan.variations) reintroduced; " - "see round-2 R2-L7." + "a length mismatch would silently truncate." ) diff --git a/tests/unit/cli_runner/test_request_count_override.py b/tests/unit/cli_runner/test_request_count_override.py index 92cf37961b..300835e91e 100644 --- a/tests/unit/cli_runner/test_request_count_override.py +++ b/tests/unit/cli_runner/test_request_count_override.py @@ -2,12 +2,12 @@ # SPDX-License-Identifier: Apache-2.0 """Regression: ``--request-count`` overrides YAML ``phases[*].requests``. -Round-2 R2-H3 reproduced this end-to-end: with ``--config minimal.yaml`` -and ``--request-count 10``, the CLI flag silently no-opped because the +With ``--config minimal.yaml`` +and ``--request-count 10``, the CLI flag would silently no-op if the YAML+CLI resolver only built section-level overrides and never overlaid loadgen-derived values onto the YAML's profiling phase. -The fix in ``aiperf.config.flags.resolver._apply_phase_loadgen_overrides`` +``aiperf.config.flags.resolver._apply_phase_loadgen_overrides`` walks the merged envelope and writes loadgen fields onto the phase named ``profiling`` (or the sole non-warmup entry). This test locks that in. """ diff --git a/tests/unit/common/test_aiperf_logger.py b/tests/unit/common/test_aiperf_logger.py index 6151f7a10f..e915729f7e 100644 --- a/tests/unit/common/test_aiperf_logger.py +++ b/tests/unit/common/test_aiperf_logger.py @@ -302,7 +302,9 @@ def aiperf_plain_string(): def aiperf_plain_string_lazy(): aiperf_logger.debug( - lambda: "Hello, world! This is a test of an example message that will NOT be printed." + lambda: ( + "Hello, world! This is a test of an example message that will NOT be printed." + ) ) def standard_plain_string(): @@ -337,7 +339,9 @@ def aiperf_plain_string(): def aiperf_plain_string_lazy(): aiperf_logger.info( - lambda: "Hello, world! This is a test of an example message that will be printed." + lambda: ( + "Hello, world! This is a test of an example message that will be printed." + ) ) def standard_plain_string(): @@ -390,9 +394,11 @@ def test_lazy_evaluation_and_formatting_debug(self, aiperf_logger, standard_logg def aiperf_formatting_and_lazy_evaluation(): aiperf_logger.debug( - lambda: "Hello, world! This will NOT be printed %s " - * 100 - % tuple([*["test"] * 100]) + lambda: ( + "Hello, world! This will NOT be printed %s " + * 100 + % tuple([*["test"] * 100]) + ) ) def standard_formatting_no_print(): @@ -417,7 +423,9 @@ def test_lazy_evaluation_and_formatting_and_multiple_args_debug( def aiperf_multiple_args(): aiperf_logger.debug( - lambda: f"Hello Mr {time.time_ns() ** 2} {time.time_ns() ** 2} This will NOT be printed" + lambda: ( + f"Hello Mr {time.time_ns() ** 2} {time.time_ns() ** 2} This will NOT be printed" + ) ) def standard_multiple_args(): @@ -496,7 +504,9 @@ def test_large_messages_debug_math( def aiperf_f_string_math(): aiperf_logger.debug( - lambda: f"Request time: {(large_message.end_perf_ns - large_message.start_perf_ns) / NANOS_PER_SECOND:.2f}" + lambda: ( + f"Request time: {(large_message.end_perf_ns - large_message.start_perf_ns) / NANOS_PER_SECOND:.2f}" + ) ) def standard_f_string_math(): diff --git a/tests/unit/common/test_clock.py b/tests/unit/common/test_clock.py new file mode 100644 index 0000000000..f2a9dc767a --- /dev/null +++ b/tests/unit/common/test_clock.py @@ -0,0 +1,228 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the AIPerf clock abstraction. + +Covers the unit-conversion mixin, the stateless ``WallClock``, and the +``VirtualClock`` discrete-event semantics: monotonic ``advance_to``, the +fast-path for already-crossed deadlines, deterministic wake order by +``(deadline, insertion_id)``, and the ``peek_min_waiter_ns`` / +``set_on_waiter_parked`` driver hooks. +""" + +import asyncio + +import pytest + +from aiperf.common.clock import AIPerfClock, VirtualClock, WallClock + + +def test_wallclock_satisfies_protocol(): + assert isinstance(WallClock(), AIPerfClock) + + +def test_virtualclock_satisfies_protocol(): + assert isinstance(VirtualClock(), AIPerfClock) + + +def test_unit_conversions_share_the_ns_layer(): + clock = VirtualClock() + clock._now_ns = 2_500_000_000 # 2.5s + assert clock.now_ns() == 2_500_000_000 + assert clock.now() == pytest.approx(2.5) + assert clock.now_ms() == pytest.approx(2_500.0) + + +@pytest.mark.asyncio +async def test_wallclock_nonpositive_sleep_is_immediate(): + clock = WallClock() + # Must not raise or hang; negative/zero durations return immediately. + await clock.sleep_ns(0) + await clock.sleep_ns(-5) + await clock.sleep(-1.0) + + +@pytest.mark.asyncio +async def test_virtualclock_advance_wakes_crossed_sleeper(): + clock = VirtualClock() + woke_at: list[int] = [] + + async def sleeper() -> None: + await clock.sleep_ns(1_000) + woke_at.append(clock.now_ns()) + + task = asyncio.ensure_future(sleeper()) + await asyncio.sleep(0) # let the sleeper park + assert clock.has_waiters() + assert clock.peek_min_waiter_ns() == 1_000 + + await clock.advance_to(1_000) + await task + assert woke_at == [1_000] + assert not clock.has_waiters() + + +@pytest.mark.asyncio +async def test_virtualclock_advance_is_monotonic(): + clock = VirtualClock() + await clock.advance_to(5_000) + assert clock.now_ns() == 5_000 + # Backwards / equal advances are silently ignored. + await clock.advance_to(4_000) + assert clock.now_ns() == 5_000 + await clock.advance_to(5_000) + assert clock.now_ns() == 5_000 + + +@pytest.mark.asyncio +async def test_virtualclock_fast_path_for_already_crossed_deadline(): + clock = VirtualClock() + await clock.advance_to(10_000) + # Deadline already in the past -> returns immediately, never parks. + await clock.sleep_until_ns(5_000) + assert not clock.has_waiters() + + +@pytest.mark.asyncio +async def test_virtualclock_wakes_in_deadline_then_insertion_order(): + clock = VirtualClock() + order: list[str] = [] + + async def sleeper(name: str, deadline_ns: int) -> None: + await clock.sleep_until_ns(deadline_ns) + order.append(name) + + # Park three sleepers: two share a deadline (must wake in registration + # order), one is later. + t_a = asyncio.ensure_future(sleeper("a", 1_000)) + await asyncio.sleep(0) + t_b = asyncio.ensure_future(sleeper("b", 1_000)) + await asyncio.sleep(0) + t_c = asyncio.ensure_future(sleeper("c", 2_000)) + await asyncio.sleep(0) + + assert clock.peek_min_waiter_ns() == 1_000 + + # One advance crossing both 1_000 deadlines wakes a then b (insertion order). + await clock.advance_to(1_500) + await asyncio.gather(t_a, t_b) + assert order == ["a", "b"] + + await clock.advance_to(2_000) + await t_c + assert order == ["a", "b", "c"] + + +@pytest.mark.asyncio +async def test_virtualclock_on_waiter_parked_callback_fires(): + clock = VirtualClock() + parked: list[int] = [] + clock.set_on_waiter_parked(lambda: parked.append(1)) + + async def sleeper() -> None: + await clock.sleep_ns(500) + + task = asyncio.ensure_future(sleeper()) + await asyncio.sleep(0) + assert parked == [1] # callback fired exactly once when the waiter parked + + await clock.advance_to(500) + await task + + +@pytest.mark.asyncio +async def test_virtualclock_peek_returns_none_when_idle(): + clock = VirtualClock() + assert clock.peek_min_waiter_ns() is None + assert not clock.has_waiters() + + +@pytest.mark.asyncio +async def test_wallclock_nonpositive_sleep_yields_to_event_loop(): + """C4: a tight loop of sleep_ns(<=0) must not starve a concurrent task. + + The protocol docstring promises asyncio.sleep semantics, which yield to + the event loop exactly once even for non-positive durations. + """ + clock = WallClock() + progress: list[int] = [] + + async def side_task() -> None: + for i in range(5): + progress.append(i) + await asyncio.sleep(0) + + task = asyncio.ensure_future(side_task()) + for _ in range(10): + await clock.sleep_ns(0) + await clock.sleep_ns(-5) + await clock.sleep_until_ns(clock.now_ns() - 1_000) + await task + assert progress == [0, 1, 2, 3, 4] + + +@pytest.mark.asyncio +async def test_virtualclock_nonpositive_and_crossed_sleeps_yield_to_event_loop(): + """C4: VirtualClock fast paths (zero duration, already-crossed deadline) + must also yield once instead of returning synchronously.""" + clock = VirtualClock() + await clock.advance_to(10_000) + progress: list[int] = [] + + async def side_task() -> None: + for i in range(3): + progress.append(i) + await asyncio.sleep(0) + + task = asyncio.ensure_future(side_task()) + for _ in range(5): + await clock.sleep_ns(0) + await clock.sleep_until_ns(5_000) # already crossed + await task + assert progress == [0, 1, 2] + assert not clock.has_waiters() + + +@pytest.mark.asyncio +async def test_virtualclock_cancelled_sleeper_is_reaped(): + """C5: cancelling a parked sleeper must remove its phantom deadline so a + driver pump cannot fast-forward sim time to a deadline nobody waits on.""" + clock = VirtualClock() + task = asyncio.ensure_future(clock.sleep_until_ns(1_000)) + await asyncio.sleep(0) # let the sleeper park + assert clock.has_waiters() + assert clock.peek_min_waiter_ns() == 1_000 + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert not clock.has_waiters() + assert clock.peek_min_waiter_ns() is None + + +@pytest.mark.asyncio +async def test_virtualclock_cancelled_sleeper_does_not_mask_live_waiter(): + """C5: a cancelled earlier deadline must neither be reported by peek nor + disturb the wake of a live later waiter.""" + clock = VirtualClock() + dead = asyncio.ensure_future(clock.sleep_until_ns(1_000)) + await asyncio.sleep(0) + woke_at: list[int] = [] + + async def live_sleeper() -> None: + await clock.sleep_until_ns(2_000) + woke_at.append(clock.now_ns()) + + live = asyncio.ensure_future(live_sleeper()) + await asyncio.sleep(0) + + dead.cancel() + with pytest.raises(asyncio.CancelledError): + await dead + + assert clock.peek_min_waiter_ns() == 2_000 + # Advancing across the dead entry's deadline must not error or wake it. + await clock.advance_to(2_000) + await live + assert woke_at == [2_000] + assert not clock.has_waiters() diff --git a/tests/unit/common/test_dynamo_settings.py b/tests/unit/common/test_dynamo_settings.py new file mode 100644 index 0000000000..73391af16d --- /dev/null +++ b/tests/unit/common/test_dynamo_settings.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from aiperf.common.environment import _DynamoSettings + + +def test_dynamo_settings_defaults() -> None: + s = _DynamoSettings() + assert s.MAX_SUBAGENT_DEPTH == 16 + + +def test_dynamo_settings_env_override(monkeypatch) -> None: + monkeypatch.setenv("AIPERF_DYNAMO_MAX_SUBAGENT_DEPTH", "4") + s = _DynamoSettings() + assert s.MAX_SUBAGENT_DEPTH == 4 diff --git a/tests/unit/common/test_finite.py b/tests/unit/common/test_finite.py index 468ef15a76..4c349bbe13 100644 --- a/tests/unit/common/test_finite.py +++ b/tests/unit/common/test_finite.py @@ -248,7 +248,7 @@ def _make_aggregate_with_combos(combos): @pytest.mark.asyncio async def test_aggregate_sweep_csv_header_is_union_across_combos(tmp_path) -> None: - """R2-H6 regression: CSV header must include metrics that appear in any combo, not just combo[0].""" + """CSV header must include metrics that appear in any combo, not just combo[0].""" from aiperf.exporters.aggregate import ( AggregateExporterConfig, AggregateSweepCsvExporter, @@ -281,7 +281,7 @@ async def test_aggregate_sweep_csv_header_is_union_across_combos(tmp_path) -> No @pytest.mark.asyncio async def test_aggregate_sweep_csv_header_when_first_combo_empty(tmp_path) -> None: - """R2-H6 regression: if combo[0] has empty metrics, columns must still appear from later combos.""" + """If combo[0] has empty metrics, columns must still appear from later combos.""" from aiperf.exporters.aggregate import ( AggregateExporterConfig, AggregateSweepCsvExporter, @@ -305,7 +305,7 @@ async def test_aggregate_sweep_csv_header_when_first_combo_empty(tmp_path) -> No def test_aggregate_sweep_csv_format_number_handles_nan_and_inf() -> None: - """R2-M10 regression: NaN, +inf, -inf must all render as empty string (matching None).""" + """NaN, +inf, -inf must all render as empty string (matching None).""" from aiperf.exporters.aggregate.aggregate_sweep_csv_exporter import ( AggregateSweepCsvExporter, ) diff --git a/tests/unit/common/test_hash_id_random_generator.py b/tests/unit/common/test_hash_id_random_generator.py new file mode 100644 index 0000000000..6016f36284 --- /dev/null +++ b/tests/unit/common/test_hash_id_random_generator.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""HashIdRandomGenerator seed derivation (C6). + +``from_base_rng`` must read ``base_rng.seed`` WITHOUT consuming RNG state -- +the weka content synthesizer relies on that stability contract. The old falsy +``or`` fallback conflated the legal seed 0 with "unset", silently drawing a +seed and mutating the base RNG. +""" + +from __future__ import annotations + +from aiperf.common.hash_id_random_generator import HashIdRandomGenerator +from aiperf.common.random_generator import RandomGenerator + + +def _base_rng(seed: int | None) -> RandomGenerator: + return RandomGenerator(seed, _internal=True) + + +def test_from_base_rng_seed_zero_preserved() -> None: + gen = HashIdRandomGenerator.from_base_rng(_base_rng(0)) + assert gen.seed == 0 + + +def test_from_base_rng_seed_zero_does_not_consume_base_state() -> None: + base = _base_rng(0) + HashIdRandomGenerator.from_base_rng(base) + fresh = _base_rng(0) + assert base.randrange(0, 2**64) == fresh.randrange(0, 2**64) + + +def test_from_base_rng_seed_zero_produces_deterministic_output() -> None: + gen_a = HashIdRandomGenerator.from_base_rng(_base_rng(0)) + gen_b = HashIdRandomGenerator.from_base_rng(_base_rng(0)) + gen_a.reseed_for_hash_id(7, trace_id="trace") + gen_b.reseed_for_hash_id(7, trace_id="trace") + assert [gen_a.randrange(0, 10**6) for _ in range(5)] == [ + gen_b.randrange(0, 10**6) for _ in range(5) + ] + + +def test_from_base_rng_seedless_base_draws_fallback_seed() -> None: + gen = HashIdRandomGenerator.from_base_rng(_base_rng(None)) + assert gen.seed is not None diff --git a/tests/unit/common/test_random_generator.py b/tests/unit/common/test_random_generator.py index add6cc402f..e2a37e4717 100644 --- a/tests/unit/common/test_random_generator.py +++ b/tests/unit/common/test_random_generator.py @@ -198,6 +198,30 @@ def test_reset_and_reinitialize(self): assert val1 != val2 +class TestRootSeedAccessor: + """Test the read-only root_seed() module accessor.""" + + def test_root_seed_initialized_returns_init_seed(self): + """Test that root_seed() returns the seed passed to init().""" + rng.reset() + rng.init(42) + + assert rng.root_seed() == 42 + + def test_root_seed_after_reset_returns_none(self): + """Test that root_seed() returns None when the manager is uninitialized.""" + rng.reset() + + assert rng.root_seed() is None + + def test_root_seed_unseeded_manager_returns_none(self): + """Test that root_seed() returns None for an unseeded (init(None)) manager.""" + rng.reset() + rng.init(None) + + assert rng.root_seed() is None + + class TestSampleMethodEdgeCases: """Test edge cases for sample_normal and sample_positive_normal_integer.""" diff --git a/tests/unit/config/test_camel_case_config.py b/tests/unit/config/test_camel_case_config.py index 13bf56b49e..f88164bc18 100644 --- a/tests/unit/config/test_camel_case_config.py +++ b/tests/unit/config/test_camel_case_config.py @@ -731,8 +731,8 @@ class TestTemplateFilesLoad: """Verify all shipped config templates (now camelCase) load correctly.""" pytestmark = pytest.mark.skip( - reason="Wave 2: shipped templates still use phases-as-dict; will be migrated " - "in Task 7 of 2026-04-26-phases-list-with-name.md." + reason="shipped templates still use phases-as-dict; skipped until they " + "are migrated to the phases-as-list shape." ) @staticmethod diff --git a/tests/unit/config/test_cli_flags_graph_selection.py b/tests/unit/config/test_cli_flags_graph_selection.py new file mode 100644 index 0000000000..437d7e21a7 --- /dev/null +++ b/tests/unit/config/test_cli_flags_graph_selection.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CLI-flag plumbing + resolved-config landing for the graph-plane +dataset-selection knobs ``--max-context-length`` and ``--allow-dataset-wrap``. + +These are the config FIELDS the resolver and dispatch layers wire to: + +- ``CLIConfig.max_context_length`` / ``CLIConfig.allow_dataset_wrap`` + (flat cyclopts flags, INPUT group). +- ``SynthesisConfig.max_context_length`` / ``SynthesisConfig.allow_dataset_wrap`` + (raw explicit values carried through ``FileDataset.synthesis``; ``None`` + when unset so ``GraphDispatchResolver`` can distinguish unset from explicit). +- ``ResolvedConfig.allow_dataset_wrap`` (``GraphDispatchResolver`` derives the + default here; ``GraphIRReplayStrategy`` reads + ``run.resolved.allow_dataset_wrap``). +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from pytest import param + +from aiperf.config.dataset import FileDataset, SynthesisConfig +from aiperf.config.flags._converter_dataset import build_dataset +from aiperf.config.flags.cli_config import CLIConfig +from aiperf.config.flags.converter import convert_cli_to_aiperf +from aiperf.config.resolution.plan import ResolvedConfig + +# --- cyclopts parse harness (mirrors test_auto_plot_fields) ------------------ + + +def _parse_cli_args(argv: list[str]) -> CLIConfig: + """Parse ``argv`` through cyclopts into a ``CLIConfig`` (no execution).""" + from cyclopts import App + + captured: dict[str, CLIConfig] = {} + app = App(name="test_profile") + + @app.default + def _runner(*, cli_config: CLIConfig) -> None: # pragma: no cover - capture only + captured["uc"] = cli_config + + try: + app(argv, exit_on_error=False) + except SystemExit as exc: + if exc.code not in (0, None): + raise + return captured["uc"] + + +def _required_endpoint_args() -> list[str]: + """Minimal endpoint flags needed for any CLIConfig parse to succeed.""" + return ["--url", "http://localhost:8000/test", "--model", "test-model"] + + +# --- CLI flag plumbing ------------------------------------------------------- + + +def test_cyclopts_parses_max_context_length_and_allow_dataset_wrap() -> None: + """``--max-context-length 131072 --allow-dataset-wrap`` populates the DTO.""" + uc = _parse_cli_args( + [ + *_required_endpoint_args(), + "--max-context-length", + "131072", + "--allow-dataset-wrap", + ] + ) + assert uc.max_context_length == 131072 + assert uc.allow_dataset_wrap is True + assert "max_context_length" in uc.model_fields_set + assert "allow_dataset_wrap" in uc.model_fields_set + + +def test_cyclopts_parses_no_allow_dataset_wrap() -> None: + """``--no-allow-dataset-wrap`` records an explicit False (not unset).""" + uc = _parse_cli_args([*_required_endpoint_args(), "--no-allow-dataset-wrap"]) + assert uc.allow_dataset_wrap is False + assert "allow_dataset_wrap" in uc.model_fields_set + + +def test_cli_defaults_are_none() -> None: + """Unset flags default to None and stay out of ``model_fields_set`` so the + resolver can distinguish unset from explicit.""" + uc = _parse_cli_args(_required_endpoint_args()) + assert uc.max_context_length is None + assert uc.allow_dataset_wrap is None + assert "max_context_length" not in uc.model_fields_set + assert "allow_dataset_wrap" not in uc.model_fields_set + + +@pytest.mark.parametrize( + "value", + [param(None, id="none"), param(True, id="true"), param(False, id="false")], +) # fmt: skip +def test_cli_allow_dataset_wrap_tristate(value: bool | None) -> None: + """``CLIConfig.allow_dataset_wrap`` accepts None / True / False unchanged.""" + uc = CLIConfig(allow_dataset_wrap=value) + assert uc.allow_dataset_wrap is value + + +# --- SynthesisConfig field --------------------------------------------------- + + +def test_synthesis_config_defaults_none() -> None: + """Both new SynthesisConfig fields default to None.""" + cfg = SynthesisConfig() + assert cfg.max_context_length is None + assert cfg.allow_dataset_wrap is None + + +def test_synthesis_config_accepts_values() -> None: + cfg = SynthesisConfig(max_context_length=131072, allow_dataset_wrap=True) + assert cfg.max_context_length == 131072 + assert cfg.allow_dataset_wrap is True + + +def test_synthesis_config_rejects_non_positive_max_context_length() -> None: + with pytest.raises(ValueError): + SynthesisConfig(max_context_length=0) + + +# --- ResolvedConfig field ---------------------------------------------------- + + +def test_resolved_config_allow_dataset_wrap_default_none() -> None: + """The resolved field defaults to None until ``GraphDispatchResolver`` derives it.""" + assert ResolvedConfig().allow_dataset_wrap is None + + +# --- values land on the resolved config -------------------------------------- + + +def _trace_file(tmp_path: Path) -> str: + """CLIConfig.input_file validates existence, so materialize a trace file.""" + path = tmp_path / "trace.jsonl" + path.write_text('{"timestamp": 0, "input_length": 8, "output_length": 4}\n') + return str(path) + + +def _file_user(tmp_path: Path, **overrides) -> CLIConfig: + kwargs: dict = { + "model_names": ["test-model"], + "input_file": _trace_file(tmp_path), + "max_context_length": 131072, + "allow_dataset_wrap": True, + } + kwargs.update(overrides) + return CLIConfig(**kwargs) + + +def test_values_land_on_dataset_synthesis_dict(tmp_path: Path) -> None: + """Converter routes both flags into the FileDataset ``synthesis`` sub-dict.""" + ds = build_dataset(_file_user(tmp_path)) + assert ds["synthesis"]["max_context_length"] == 131072 + assert ds["synthesis"]["allow_dataset_wrap"] is True + + +def test_values_land_on_full_aiperf_config(tmp_path: Path) -> None: + """Full CLI -> AIPerfConfig path preserves the values on the FileDataset.""" + cfg = convert_cli_to_aiperf(_file_user(tmp_path)) + main = cfg.benchmark.get_default_dataset() + assert isinstance(main, FileDataset) + assert main.synthesis is not None + assert main.synthesis.max_context_length == 131072 + assert main.synthesis.allow_dataset_wrap is True + + +def test_explicit_false_wrap_is_carried_through(tmp_path: Path) -> None: + """``allow_dataset_wrap=False`` is a non-None explicit value and is carried.""" + ds = build_dataset(_file_user(tmp_path, allow_dataset_wrap=False)) + assert ds["synthesis"]["allow_dataset_wrap"] is False + + +def test_unset_wrap_omits_synthesis_key(tmp_path: Path) -> None: + """Unset flags are not carried, so a resolver sees None (unset).""" + ds = build_dataset( + CLIConfig(model_names=["test-model"], input_file=_trace_file(tmp_path)) + ) + synthesis = ds.get("synthesis") or {} + assert "max_context_length" not in synthesis + assert "allow_dataset_wrap" not in synthesis diff --git a/tests/unit/config/test_converter_dataset_entries.py b/tests/unit/config/test_converter_dataset_entries.py new file mode 100644 index 0000000000..bafd6a2c56 --- /dev/null +++ b/tests/unit/config/test_converter_dataset_entries.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Dataset-kind gating for entry-count resolution (graph #1106 signal fix). + +``--request-count`` is a recycle count on the graph/file plane, not a corpus +cap: a single trace can emit many requests. So for a **file** dataset it must +NOT back-fill ``FileDataset.entries``. An explicit ``--num-dataset-entries`` / +``--num-conversations`` still caps a file dataset (single-pass semantics), and +synthetic/public datasets keep the historical ``--request-count`` back-fill. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from aiperf.config.flags._converter_dataset import _resolve_entries +from aiperf.config.flags.cli_config import CLIConfig + + +def _loadgen(**kwargs: object) -> dict: + """Model-dump the load-generator fields (request_count) as explicit-set.""" + return CLIConfig(**kwargs).model_dump(exclude_unset=True) + + +def _file_cli(tmp_path: Path, **kwargs: object) -> CLIConfig: + """Build a real FILE-dataset CLIConfig backed by an existing input file.""" + input_file = tmp_path / "trace.jsonl" + input_file.write_text('{"text": "hi"}\n') + return CLIConfig(model_names=["test-model"], input_file=str(input_file), **kwargs) + + +def test_request_count_does_not_backfill_file_dataset_entries( + tmp_path: Path, +) -> None: + """File dataset + only --request-count -> entries None (no corpus cap).""" + cli = _file_cli(tmp_path, **_loadgen(request_count=500)) + assert cli.input_file is not None + assert "request_count" in cli.model_fields_set + assert _resolve_entries(cli) is None + + +def test_explicit_num_dataset_entries_sets_file_entries(tmp_path: Path) -> None: + """Explicit --num-dataset-entries caps a file dataset even with --request-count.""" + cli = _file_cli( + tmp_path, + conversation_num_dataset_entries=50, + **_loadgen(request_count=500), + ) + assert _resolve_entries(cli) == 50 + + +def test_explicit_num_conversations_sets_file_entries(tmp_path: Path) -> None: + """Explicit --num-conversations caps a file dataset (single-pass semantics).""" + cli = _file_cli( + tmp_path, + conversation_num=25, + **_loadgen(request_count=500), + ) + assert _resolve_entries(cli) == 25 + + +def test_synthetic_still_backfills_from_request_count() -> None: + """Synthetic dataset keeps the --request-count -> entries back-fill.""" + cli = CLIConfig(model_names=["test-model"], **_loadgen(request_count=500)) + assert cli.input_file is None + assert _resolve_entries(cli) == 500 + + +def test_public_still_backfills_from_request_count() -> None: + """Public dataset keeps the --request-count -> entries back-fill.""" + cli = CLIConfig( + model_names=["test-model"], + public_dataset="sharegpt", + **_loadgen(request_count=500), + ) + assert cli.input_file is None + assert _resolve_entries(cli) == 500 + + +@pytest.mark.parametrize("request_count", [None, 500]) +def test_file_dataset_no_entry_source_returns_none( + tmp_path: Path, request_count: int | None +) -> None: + """File dataset with no explicit entry source resolves to None regardless of request_count.""" + extra = _loadgen(request_count=request_count) if request_count is not None else {} + cli = _file_cli(tmp_path, **extra) + assert _resolve_entries(cli) is None diff --git a/tests/unit/config/test_converter_graph_single_pass.py b/tests/unit/config/test_converter_graph_single_pass.py new file mode 100644 index 0000000000..3aa6d2d08a --- /dev/null +++ b/tests/unit/config/test_converter_graph_single_pass.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Bare graph run = single corpus pass (no auto-``--request-count 10`` truncation). + +A bare ``aiperf profile --graph X.json`` with none of ``--request-count`` / +``--num-conversations`` / ``--benchmark-duration`` must NOT get the plain-aiperf +auto-10 stop condition (which truncates the benchmark to 10 node dispatches). +Instead the profiling ``ConcurrencyPhase`` stays UNBOUNDED and validates -- the +graph strategy then does a single corpus pass (each loaded trace once) driven by +``_recycle_has_stop_condition() is False``. Non-graph bare runs keep the auto-10. + +The seam: +* ``_converter_profiling.build_profiling`` detects a graph workload from the CLI + (``--graph-format`` set, or ``--input-file`` sniffed as a graph adapter) and, + for a bare run, leaves the phase UNBOUNDED (no requests/duration/sessions) + instead of injecting ``requests=10``. +* ``check_phase_dataset_compatibility`` (run from ``BenchmarkConfig``) accepts a + no-stop concurrency phase ONLY when its dataset is a graph workload (mirrors + how ``FixedSchedulePhase`` infers its stop from the dataset). A no-stop phase + against a non-graph dataset is rejected there. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from aiperf.config.config import BenchmarkConfig +from aiperf.config.flags._converter_profiling import build_profiling +from aiperf.config.flags.cli_config import CLIConfig +from aiperf.config.phases import ConcurrencyPhase +from aiperf.plugin.enums import PhaseType + +_WEKA_MIN = Path(__file__).parents[1] / "graph" / "fixtures" / "weka_min.json" + + +def _config_with(datasets: list[dict], phases: list[dict]) -> dict: + """Minimal valid BenchmarkConfig dict for the given datasets + phases.""" + return { + "models": ["test-model"], + "endpoint": {"urls": ["http://localhost:8000/v1/chat/completions"]}, + "datasets": datasets, + "phases": phases, + } + + +def _validate_profiling_phase(prof: dict) -> ConcurrencyPhase: + """Validate the profiling prof-dict as the real ConcurrencyPhase. + + Mirrors ``converter._assemble_envelope_dict`` which prepends + ``name="profiling"`` before handing the dict to the phase model. + """ + return ConcurrencyPhase(name="profiling", **prof) + + +# --------------------------------------------------------------------------- +# Bare graph run: NO auto-10, phase unbounded but valid (single corpus pass) +# --------------------------------------------------------------------------- + + +def test_bare_graph_run_autodetected_has_no_stop_condition() -> None: + """Bare ``--input-file `` (autodetected graph) -> no auto-10 bound.""" + cli = CLIConfig(model_names=["test-model"], input_file=str(_WEKA_MIN)) + prof = build_profiling(cli) + + assert prof["type"] == PhaseType.CONCURRENCY + assert prof.get("requests") is None + assert prof.get("duration") is None + assert prof.get("sessions") is None + # The stop_condition_inferred flag is gone; graph-ness now drives validity. + assert "stop_condition_inferred" not in prof + + +def test_bare_graph_run_autodetected_phase_validates_unbounded() -> None: + """The resulting unbounded concurrency phase VALIDATES at phase level.""" + cli = CLIConfig(model_names=["test-model"], input_file=str(_WEKA_MIN)) + phase = _validate_profiling_phase(build_profiling(cli)) + + assert phase.requests is None + assert phase.duration is None + assert phase.sessions is None + + +def test_bare_graph_run_via_graph_format_flag_has_no_stop_condition() -> None: + """Explicit ``--graph-format`` forces graph mode -> no auto-10 bound.""" + cli = CLIConfig( + model_names=["test-model"], + input_file=str(_WEKA_MIN), + graph_format="weka_trace", + ) + prof = build_profiling(cli) + + assert prof.get("requests") is None + assert "stop_condition_inferred" not in prof + _validate_profiling_phase(prof) # must not raise + + +# --------------------------------------------------------------------------- +# Regression: NON-graph bare run still gets the auto-10 (unchanged behavior) +# --------------------------------------------------------------------------- + + +def test_non_graph_bare_run_still_gets_auto_10() -> None: + """A synthetic (non-graph) bare run keeps the plain-aiperf ``requests=10``.""" + cli = CLIConfig(model_names=["test-model"]) + prof = build_profiling(cli) + + assert prof["type"] == PhaseType.CONCURRENCY + assert prof["requests"] == 10 + assert "stop_condition_inferred" not in prof + # And it validates with the auto-10 bound. + assert _validate_profiling_phase(prof).requests == 10 + + +def test_non_graph_file_bare_run_still_gets_auto_10(tmp_path: Path) -> None: + """A plain (non-graph) single-turn file bare run keeps the auto-10.""" + input_file = tmp_path / "plain.jsonl" + input_file.write_text('{"text": "hi"}\n') + cli = CLIConfig(model_names=["test-model"], input_file=str(input_file)) + prof = build_profiling(cli) + + assert prof["requests"] == 10 + assert "stop_condition_inferred" not in prof + + +# --------------------------------------------------------------------------- +# Bounded graph run: an explicit bound is preserved (no inferred flag) +# --------------------------------------------------------------------------- + + +def test_bounded_graph_run_preserves_request_count() -> None: + """``--request-count 25`` on a graph run keeps requests==25 (no inferred flag).""" + cli = CLIConfig( + model_names=["test-model"], + input_file=str(_WEKA_MIN), + request_count=25, + ) + prof = build_profiling(cli) + + assert prof["requests"] == 25 + assert "stop_condition_inferred" not in prof + + +def test_bounded_graph_run_num_conversations_preserved() -> None: + """``--num-conversations 7`` on a graph run keeps sessions==7 (no inferred flag).""" + cli = CLIConfig( + model_names=["test-model"], + input_file=str(_WEKA_MIN), + conversation_num=7, + ) + prof = build_profiling(cli) + + assert prof["sessions"] == 7 + assert "stop_condition_inferred" not in prof + + +# --------------------------------------------------------------------------- +# Phase-model seam: the required-stop check no longer lives on the phase +# --------------------------------------------------------------------------- + + +def test_stop_condition_inferred_field_is_gone() -> None: + """The user-facing opt-out flag was deleted entirely (no bypass foot-gun).""" + assert "stop_condition_inferred" not in ConcurrencyPhase.model_fields + + +def test_bare_concurrency_phase_validates_at_phase_level() -> None: + """A no-stop concurrency phase no longer raises at phase construction. + + The required-stop check moved to ``check_phase_dataset_compatibility`` (run + from ``BenchmarkConfig``), which needs the dataset to decide -- so the phase + model alone accepts a no-stop concurrency phase. + """ + phase = ConcurrencyPhase(name="profiling", type=PhaseType.CONCURRENCY) + assert phase.requests is None + assert phase.duration is None + assert phase.sessions is None + + +# --------------------------------------------------------------------------- +# BenchmarkConfig validation seam: graph exempts no-stop, non-graph rejects it +# --------------------------------------------------------------------------- + + +def test_graph_config_no_stop_condition_validates_via_graph_format() -> None: + """A graph dataset (forced via ``graph_format``) + no-stop phase VALIDATES.""" + cfg = BenchmarkConfig( + **_config_with( + datasets=[ + { + "name": "main", + "type": "file", + "path": str(_WEKA_MIN), + "graph_format": "weka_trace", + } + ], + phases=[{"name": "profiling", "type": "concurrency", "concurrency": 8}], + ) + ) + profiling = next(p for p in cfg.phases if p.name == "profiling") + assert profiling.requests is None + assert profiling.duration is None + assert profiling.sessions is None + + +def test_graph_config_no_stop_condition_validates_via_autodetect() -> None: + """A graph dataset (auto-detected from the weka path) + no-stop phase VALIDATES.""" + cfg = BenchmarkConfig( + **_config_with( + datasets=[{"name": "main", "type": "file", "path": str(_WEKA_MIN)}], + phases=[{"name": "profiling", "type": "concurrency", "concurrency": 8}], + ) + ) + profiling = next(p for p in cfg.phases if p.name == "profiling") + assert profiling.requests is None + + +def test_non_graph_config_no_stop_condition_raises() -> None: + """A non-graph (synthetic) dataset + no-stop concurrency phase RAISES.""" + with pytest.raises(ValidationError, match="at least one of"): + BenchmarkConfig( + **_config_with( + datasets=[ + { + "name": "main", + "type": "synthetic", + "entries": 100, + "prompts": {"isl": 128, "osl": 64}, + } + ], + phases=[{"name": "profiling", "type": "concurrency", "concurrency": 8}], + ) + ) diff --git a/tests/unit/config/test_distribution_validators.py b/tests/unit/config/test_distribution_validators.py index e213872fb0..260695eb5e 100644 --- a/tests/unit/config/test_distribution_validators.py +++ b/tests/unit/config/test_distribution_validators.py @@ -2,11 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 """Distribution parameter validators: reject non-positive / non-finite values. -Round-2 R2-M14 reproduced ``isl: {mean: -512, stddev: 0}`` validating +Without these validators ``isl: {mean: -512, stddev: 0}`` validates silently. NormalDistribution truncates samples below 0, so a non-positive mean produces a degenerate (or impossible) distribution that crashes -later in synthesis. The fix adds ``gt=0`` to the mean field plus a -finite-value validator covering NaN/inf. +later in synthesis. ``gt=0`` on the mean field plus a +finite-value validator covering NaN/inf reject it at config time. """ from __future__ import annotations diff --git a/tests/unit/config/test_dump_config_roundtrip.py b/tests/unit/config/test_dump_config_roundtrip.py index ebed1d1f67..83c9f19860 100644 --- a/tests/unit/config/test_dump_config_roundtrip.py +++ b/tests/unit/config/test_dump_config_roundtrip.py @@ -2,9 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 """Round-trip regression: every bundled template loads -> dumps -> reloads. -Covers R2-H2 from round-2 adversarial review: ``GridSweep.type="grid"`` is +``GridSweep.type="grid"`` is a default Field, so ``model_dump(exclude_defaults=True)`` strips the -discriminator and reload fails with ``union_tag_not_found``. The fix in +discriminator and reload fails with ``union_tag_not_found``. ``dump_config`` re-injects ``sweep.type`` so the round-trip is stable for every template, not just sweep templates. @@ -82,7 +82,7 @@ def test_dump_reload_roundtrip_for_every_template( def test_sweep_discriminator_survives_dump() -> None: """The grid sweep ``type:`` discriminator is preserved on dump. - Targeted regression for R2-H2: ``exclude_defaults=True`` would strip + Targeted regression: ``exclude_defaults=True`` would strip ``type: grid`` because it's the default; the discriminated union rejects the dumped YAML on reload. ``dump_config`` re-injects the discriminator so this scenario round-trips. diff --git a/tests/unit/config/test_graph_dispatch_resolver.py b/tests/unit/config/test_graph_dispatch_resolver.py new file mode 100644 index 0000000000..8316011a5b --- /dev/null +++ b/tests/unit/config/test_graph_dispatch_resolver.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for GraphDispatchResolver: derives allow_dataset_wrap + publishes sampling. + +Runs the real ScenarioResolver (so the scenario's cache-bust auto-fill lands on +``endpoint.cache_bust``) followed by the new GraphDispatchResolver, against real +config objects and a real weka graph fixture. +""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +from aiperf.common.enums import CacheBustTarget +from aiperf.config import BenchmarkConfig +from aiperf.config.resolution.plan import BenchmarkRun +from aiperf.config.resolution.resolvers import ( + DatasetResolver, + GraphDispatchResolver, + ScenarioResolver, + build_default_resolver_chain, +) + +_WEKA_FIXTURE = ( + Path(__file__).parents[2] / "unit/graph/fixtures/weka_min.json" +).resolve() + + +def _graph_run( + *, + scenario: str | None = None, + cache_bust_unset: bool = True, + allow_wrap_unset: bool = True, + cache_bust: str | None = None, + allow_wrap: bool | None = None, +) -> BenchmarkRun: + """Build a graph-workload BenchmarkRun from real config objects. + + ``cache_bust``/``allow_wrap`` (when not None) are explicit user settings; + the ``*_unset`` flags document the default-path cases and are otherwise + unused (the None sentinels drive the actual wiring). + """ + endpoint: dict[str, object] = { + "urls": ["http://localhost:8000/v1/chat/completions"] + } + if cache_bust is not None: + endpoint["cache_bust"] = cache_bust + + dataset: dict[str, object] = { + "name": "profiling", + "type": "file", + "path": str(_WEKA_FIXTURE), + } + if allow_wrap is not None: + dataset["synthesis"] = {"allow_dataset_wrap": allow_wrap} + + cfg = BenchmarkConfig( + models=["claude-opus-4-5-20251101"], + endpoint=endpoint, + datasets=[dataset], + phases=[ + { + "name": "profiling", + "type": "concurrency", + "concurrency": 1, + "sessions": 5, + } + ], + scenario=scenario, + ) + return BenchmarkRun(benchmark_id="test-run", cfg=cfg, artifact_dir=Path("/tmp/x")) + + +def apply_chain(run: BenchmarkRun) -> SimpleNamespace: + """Run the real ScenarioResolver then GraphDispatchResolver over ``run``. + + Returns a view exposing the resolved endpoint (for cache_bust assertions) + and the resolved wrap/sampling fields. + """ + DatasetResolver().resolve(run) + ScenarioResolver().resolve(run) + GraphDispatchResolver().resolve(run) + return SimpleNamespace( + endpoint=run.cfg.endpoint, + allow_dataset_wrap=run.resolved.allow_dataset_wrap, + dataset_sampling_strategy=run.resolved.dataset_sampling_strategy, + ) + + +def test_wrap_default_true_when_scenario_forces_cache_bust(): + run = _graph_run( + scenario="inferencex-agentx-mvp", cache_bust_unset=True, allow_wrap_unset=True + ) + resolved = apply_chain( + run + ) # full default chain incl. ScenarioResolver then the new step + assert resolved.endpoint.cache_bust == CacheBustTarget.FIRST_TURN_PREFIX + assert resolved.allow_dataset_wrap is True + + +def test_wrap_default_false_when_cache_bust_none(): + run = _graph_run(scenario=None, cache_bust_unset=True, allow_wrap_unset=True) + assert apply_chain(run).allow_dataset_wrap is False + + +def test_explicit_allow_wrap_false_wins_even_with_cache_bust(): + run = _graph_run(cache_bust="first_turn_prefix", allow_wrap=False) + assert apply_chain(run).allow_dataset_wrap is False + + +def test_graph_dispatch_resolver_runs_after_scenario_resolver(): + chain = build_default_resolver_chain()._resolvers + types = [type(r) for r in chain] + assert GraphDispatchResolver in types + assert types.index(ScenarioResolver) < types.index(GraphDispatchResolver) + + +def test_publishes_dataset_sampling_strategy_for_graph_workload(): + from aiperf.plugin.enums import DatasetSamplingStrategy + + run = _graph_run(scenario=None) + resolved = apply_chain(run) + assert resolved.dataset_sampling_strategy == DatasetSamplingStrategy.SEQUENTIAL diff --git a/tests/unit/config/test_graph_format_flag.py b/tests/unit/config/test_graph_format_flag.py new file mode 100644 index 0000000000..37dcaa9a44 --- /dev/null +++ b/tests/unit/config/test_graph_format_flag.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from aiperf.common.enums import DatasetType +from aiperf.config.dataset.config import FileDataset + + +def test_filedataset_accepts_graph_format() -> None: + ds = FileDataset( + name="d", type=DatasetType.FILE, path="x.gz", graph_format="dynamo_trace" + ) + assert str(ds.graph_format) == "dynamo_trace" + + +def test_filedataset_graph_format_defaults_none() -> None: + ds = FileDataset(name="d", type=DatasetType.FILE, path="x.jsonl") + assert ds.graph_format is None diff --git a/tests/unit/config/test_graph_workload_resolution.py b/tests/unit/config/test_graph_workload_resolution.py new file mode 100644 index 0000000000..e81f7f503a --- /dev/null +++ b/tests/unit/config/test_graph_workload_resolution.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Single-source memoized graph workload resolution. + +Graph-ness is derived AT MOST ONCE per process via +``workload_detect.resolve_graph_workload``: the config resolver chain +populates ``run.resolved.graph_workload`` (+ the ``graph_workload_resolved`` +marker) eagerly in single-run mode, and the accessor derives-and-memoizes +lazily for runs that never pass the chain (``aiperf service`` processes, +test-built runs). These tests pin: + +* (a) a weka file resolves a ``GraphWorkloadRef`` at resolver-chain time; +* (b) a weka HF ``org/name`` id resolves a ref with the RAW path (never + ``.resolve()``\\d -- HF ids are not filesystem paths) via the resolver's + HF early-return branch; +* (c) a synthetic run resolves ``None`` WITH the marker set (distinguishes + "not a graph run" from "never checked"); +* (d) the accessor on a chain-less run derives exactly once across repeated + calls (memoization); +* (e) an explicit ``--graph-format`` override wins with detection never + called; +* (f) the timing/config veto: a graph input plus an explicit non-graph + ``--custom-dataset-type`` stays NON-graph for timing (the veto composes + OVER the accessor; pinned so the migration cannot absorb or drop it). +""" + +from __future__ import annotations + +from pathlib import Path + +from aiperf.config.dataset.resolver import DatasetResolver +from aiperf.config.flags.cli_config import CLIConfig +from tests.unit.conftest import make_run_from_cli + +WEKA_MIN = Path(__file__).resolve().parents[1] / "graph" / "fixtures" / "weka_min.json" + +# Weka-marked HF repo id (org/name shape, non-existent local path) -- routed +# through the resolver's HF early-return branch, never the file-existence gate. +HF_WEKA_ID = "semianalysisai/cc-traces-weka-000000" + + +def _run(**cli_overrides): + """Build an un-chained ``BenchmarkRun`` from CLI flags (no resolver chain).""" + cfg = CLIConfig(model_names=["test-model"], **cli_overrides) + return make_run_from_cli(cfg) + + +def test_weka_file_resolves_ref_at_resolver_chain_time() -> None: + run = _run(input_file=str(WEKA_MIN)) + DatasetResolver().resolve(run) + assert run.resolved.graph_workload_resolved is True + ref = run.resolved.graph_workload + assert ref is not None + assert ref.format == "weka_trace" + assert ref.path == WEKA_MIN + + +def test_hf_id_resolves_ref_with_raw_path() -> None: + run = _run(input_file=HF_WEKA_ID) + DatasetResolver().resolve(run) + assert run.resolved.graph_workload_resolved is True + ref = run.resolved.graph_workload + assert ref is not None + assert ref.format == "weka_trace" + # RAW org/name id verbatim: never .resolve()d into a cwd-anchored path. + assert ref.path == Path(HF_WEKA_ID) + assert not ref.path.is_absolute() + + +def test_synthetic_run_resolves_none_with_marker() -> None: + run = _run() + DatasetResolver().resolve(run) + assert run.resolved.graph_workload_resolved is True + assert run.resolved.graph_workload is None + + +def test_accessor_derives_once_and_memoizes(monkeypatch) -> None: + from aiperf.dataset.graph import workload_detect + + run = _run(input_file=str(WEKA_MIN)) + calls: list[Path] = [] + real_detect = workload_detect._detect_graph_workload_format + + def spy(path: Path): + calls.append(path) + return real_detect(path) + + monkeypatch.setattr(workload_detect, "_detect_graph_workload_format", spy) + first = workload_detect.resolve_graph_workload(run) + second = workload_detect.resolve_graph_workload(run) + assert first is not None + assert first.format == "weka_trace" + assert second is first + assert len(calls) == 1, "the second accessor call must read the memo" + + +def test_graph_format_override_wins_without_detection(monkeypatch, tmp_path) -> None: + from aiperf.dataset.graph import workload_detect + + plain = tmp_path / "plain.jsonl" + plain.write_text('{"messages": [{"role": "user", "content": "hi"}]}\n') + run = _run(input_file=str(plain), graph_format="native") + calls: list[Path] = [] + monkeypatch.setattr( + workload_detect, "_detect_graph_workload_format", lambda p: calls.append(p) + ) + ref = workload_detect.resolve_graph_workload(run) + assert ref is not None + assert ref.format == "native" + assert ref.path == plain + assert calls == [], "--graph-format must short-circuit registry detection" + + +def test_explicit_non_graph_format_vetoes_graph_timing() -> None: + from aiperf.common.enums import DatasetFormat + from aiperf.plugin.enums import TimingMode + from aiperf.timing.config import TimingConfig + + run = _run( + input_file=str(WEKA_MIN), + custom_dataset_type=DatasetFormat.MULTI_TURN, + warmup_request_count=2, + request_count=3, + ) + tc = TimingConfig.from_run(run) + assert tc.phase_configs + assert all(p.timing_mode != TimingMode.GRAPH_IR for p in tc.phase_configs), ( + "an explicit graph-incompatible --custom-dataset-type must veto graph " + "timing even though the input file sniffs as a graph workload" + ) diff --git a/tests/unit/config/test_plot_envelope.py b/tests/unit/config/test_plot_envelope.py index a9e3e916cf..184483b23b 100644 --- a/tests/unit/config/test_plot_envelope.py +++ b/tests/unit/config/test_plot_envelope.py @@ -9,7 +9,7 @@ from pydantic import ValidationError from pytest import param -from aiperf.config import AIPerfConfig # noqa: F401 (used by Task 3 tests below) +from aiperf.config import AIPerfConfig # noqa: F401 (used by resolution tests below) from aiperf.config.loader.core import load_config_from_string from aiperf.config.loader.errors import ConfigurationError from aiperf.config.plot import ( @@ -212,8 +212,8 @@ def test_load_plot_envelope_top_level_must_be_mapping(tmp_path: Path): def test_load_config_from_string_with_file_path_threads_source_dir(tmp_path): """When file_path is given, source_dir is the file's parent. No assertion - on plot here — that's Task 3. We just confirm the call shape still works - after threading the context.""" + on plot here — envelope resolution is covered by the tests below. We just + confirm the call shape still works after threading the context.""" minimal_yaml = """ benchmark: models: [llama-3-8b] @@ -229,8 +229,8 @@ def test_load_config_from_string_with_file_path_threads_source_dir(tmp_path): assert config.benchmark.models.items[0].name == "llama-3-8b" -# Minimum viable AIPerfConfig YAML — same shape used by the threading test in -# Task 2. Field paths/keys mirror tests/unit/config/test_envelope_restructure.py. +# Minimum viable AIPerfConfig YAML — same shape used by the threading test +# above. Field paths/keys mirror tests/unit/config/test_envelope_restructure.py. _BASE_BENCHMARK_YAML = """ benchmark: models: [llama-3-8b] diff --git a/tests/unit/config/test_resolver_edge_cases.py b/tests/unit/config/test_resolver_edge_cases.py index b74c5bf88e..d89806ab5c 100644 --- a/tests/unit/config/test_resolver_edge_cases.py +++ b/tests/unit/config/test_resolver_edge_cases.py @@ -470,6 +470,86 @@ def test_detect_type_returns_none_none_when_file_unreadable(self, tmp_path) -> N assert detected is None assert first_record is None + def test_local_graph_gz_workload_resolves_without_crash(self, tmp_path) -> None: + """A local graph trace (.gz binary) must NOT crash type detection. + + The graph-store build/schedule planes own record counting for graph + workloads; the resolver must short-circuit before ``_detect_type`` / + ``_count_records_and_sessions`` text-mode read the gzip binary (which + previously raised ``UnicodeDecodeError``). It records only the path, + mirroring the weka HF-id branch. + """ + dynamo_fixture = ( + Path(__file__).resolve().parents[1] + / "dataset/graph/adapters/fixtures/dynamo_nested/nested_2_level.jsonl.gz" + ) + + config = _make_config( + datasets=[ + {"name": "graph", "type": "file", "path": str(dynamo_fixture)}, + ], + phases=[ + { + "name": "profiling", + "type": "concurrency", + "requests": 2, + "concurrency": 1, + } + ], + ) + run = _make_run(config, artifact_dir=tmp_path / "out") + + # Previously raised UnicodeDecodeError from the text-mode open in + # _count_records_and_sessions; must now resolve cleanly. + DatasetResolver().resolve(run) + + # Path recorded; record counting skipped (owned by the graph planes). + assert run.resolved.dataset_file_paths is not None + assert "graph" in run.resolved.dataset_file_paths + assert run.resolved.dataset_total_records is None + + def test_graph_format_override_short_circuits_plain_jsonl(self, tmp_path) -> None: + """A plain `.jsonl` with `graph_format` set must short-circuit resolution. + + `--graph-format native` (and any forced adapter) marks the file a graph + workload that the graph-store planes own. The path-level content sniff + (`is_graph_workload_path`) excludes `native`, so without honoring the + config override the resolver would route this plain single-turn JSONL to + the custom-dataset loader: type-detected and record-counted. With the + override honored, only the path is recorded; type/record-count stay unset. + """ + plain = tmp_path / "hand_authored.jsonl" + plain.write_text('{"prompt": "hello"}\n{"prompt": "world"}\n') + + config = _make_config( + datasets=[ + { + "name": "graph", + "type": "file", + "path": str(plain), + "graph_format": "native", + }, + ], + phases=[ + { + "name": "profiling", + "type": "concurrency", + "requests": 2, + "concurrency": 1, + } + ], + ) + run = _make_run(config, artifact_dir=tmp_path / "out") + + DatasetResolver().resolve(run) + + # Path recorded; type detection and record counting skipped (owned by + # the graph planes), mirroring the local-graph-workload short-circuit. + assert run.resolved.dataset_file_paths is not None + assert "graph" in run.resolved.dataset_file_paths + assert run.resolved.dataset_types is None + assert run.resolved.dataset_total_records is None + def test_burst_gpt_csv_auto_detected_with_timing(self, tmp_path) -> None: """End-to-end resolver pass: BurstGPT CSV with no explicit format is recognized and reports has_timing=True so fixed_schedule can run. diff --git a/tests/unit/config/test_resolvers.py b/tests/unit/config/test_resolvers.py index 7c4a21a58e..75858addad 100644 --- a/tests/unit/config/test_resolvers.py +++ b/tests/unit/config/test_resolvers.py @@ -190,9 +190,9 @@ def test_resolve_for_probe_skips_user_files_materialization(self, tmp_path): """Probe runs must NOT materialize user_files (they re-run per variation). ``cli_runner._estimate_and_log_duration`` clones the user's first config - into a probe run only to estimate duration. After Task 6 the resolver - also wrote user_files; that produced a stray artifact tree before the - actual benchmark and could bake in template values (e.g. ``epoch``) + into a probe run only to estimate duration. If the resolver + also wrote user_files there, it would produce a stray artifact tree before + the actual benchmark and could bake in template values (e.g. ``epoch``) that don't match the per-variation runs. """ from aiperf.config.loader import load_config_from_string @@ -730,9 +730,14 @@ class TestBuildDefaultResolverChain: def test_returns_chain_with_all_resolvers(self): chain = build_default_resolver_chain() assert isinstance(chain, ConfigResolverChain) - assert len(chain._resolvers) == 6 + assert len(chain._resolvers) == 8 def test_resolver_order(self): + from aiperf.config.resolution.resolvers import ( + GraphDispatchResolver, + ScenarioResolver, + ) + chain = build_default_resolver_chain() types = [type(r) for r in chain._resolvers] assert types == [ @@ -741,6 +746,8 @@ def test_resolver_order(self): GpuMetricsResolver, CommConfigResolver, DatasetResolver, + ScenarioResolver, + GraphDispatchResolver, TimingResolver, ] diff --git a/tests/unit/config/test_scenario_cli_flags.py b/tests/unit/config/test_scenario_cli_flags.py new file mode 100644 index 0000000000..7bd9dda1b1 --- /dev/null +++ b/tests/unit/config/test_scenario_cli_flags.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""T3 + Q6 coverage for the scenario-lock CLI flags. + +T3: ``--scenario`` / ``--unsafe-override`` parse through cyclopts onto the +flat ``CLIConfig`` and thread to ``BenchmarkConfig`` via the converter. + +Q6: the converter writes endpoint fields only when the user explicitly set +them (``cli.model_fields_set & ENDPOINT_FIELDS``), so the resolved +``EndpointConfig.model_fields_set`` faithfully reflects user-explicitness for +``streaming`` / ``cache_bust`` -- the scenario validator can use membership +directly (no explicit-set sentinel needed on this branch's converter path). +""" + +from __future__ import annotations + +import pytest +from pytest import param + +from aiperf.config.flags.cli_config import CLIConfig +from aiperf.config.flags.converter import convert_cli_to_aiperf + + +def _parse_cli_args(argv: list[str]) -> CLIConfig: + """Parse ``argv`` through cyclopts into a ``CLIConfig`` (no execution).""" + from cyclopts import App + + captured: dict[str, CLIConfig] = {} + app = App(name="test_profile") + + @app.default + def _runner(*, cli_config: CLIConfig) -> None: # pragma: no cover - capture only + captured["uc"] = cli_config + + try: + app(argv, exit_on_error=False) + except SystemExit as exc: + if exc.code not in (0, None): + raise + return captured["uc"] + + +def _required_endpoint_args() -> list[str]: + return ["--url", "http://localhost:8000/test", "--model", "test-model"] + + +# --------------------------------------------------------------------------- +# T3: parse + converter pass-through +# --------------------------------------------------------------------------- + + +def test_cyclopts_parses_scenario_flag() -> None: + uc = _parse_cli_args( + [*_required_endpoint_args(), "--scenario", "inferencex-agentx-mvp"] + ) + assert uc.scenario == "inferencex-agentx-mvp" + assert "scenario" in uc.model_fields_set + + +def test_cyclopts_parses_unsafe_override_flag() -> None: + uc = _parse_cli_args([*_required_endpoint_args(), "--unsafe-override"]) + assert uc.unsafe_override is True + assert "unsafe_override" in uc.model_fields_set + + +def test_cyclopts_scenario_defaults_unset() -> None: + uc = _parse_cli_args(_required_endpoint_args()) + assert uc.scenario is None + assert uc.unsafe_override is False + assert "scenario" not in uc.model_fields_set + assert "unsafe_override" not in uc.model_fields_set + + +def test_converter_threads_scenario_fields_to_benchmark() -> None: + uc = _parse_cli_args( + [ + *_required_endpoint_args(), + "--scenario", + "inferencex-agentx-mvp", + "--unsafe-override", + ] + ) + benchmark = convert_cli_to_aiperf(uc).benchmark + assert benchmark.scenario == "inferencex-agentx-mvp" + assert benchmark.unsafe_override is True + + +def test_converter_scenario_fields_default_when_omitted() -> None: + uc = _parse_cli_args(_required_endpoint_args()) + benchmark = convert_cli_to_aiperf(uc).benchmark + assert benchmark.scenario is None + assert benchmark.unsafe_override is False + assert "scenario" not in benchmark.model_fields_set + assert "unsafe_override" not in benchmark.model_fields_set + + +# --------------------------------------------------------------------------- +# Q6: model_fields_set faithfully reflects user-explicitness +# --------------------------------------------------------------------------- + + +def _benchmark_for_cli(**kwargs): + cli = CLIConfig( + model_names=["test-model"], + urls=["http://localhost:8000/test"], + **kwargs, + ) + return convert_cli_to_aiperf(cli).benchmark + + +@pytest.mark.parametrize( + ("kwargs", "field", "present"), + [ + param({}, "streaming", False, id="streaming-unset-absent"), + param({"streaming": True}, "streaming", True, id="streaming-set-present"), + param( + {"streaming": False}, + "streaming", + True, + id="streaming-explicit-false-present", + ), + param({}, "cache_bust", False, id="cache_bust-unset-absent"), + param( + {"cache_bust": "first_turn_prefix"}, + "cache_bust", + True, + id="cache_bust-set-present", + ), + ], +) # fmt: skip +def test_endpoint_model_fields_set_reflects_explicitness( + kwargs: dict, field: str, present: bool +) -> None: + """Q6: ``EndpointConfig.model_fields_set`` carries ``streaming`` / + ``cache_bust`` only when the user explicitly set the field. + + An explicit ``streaming=False`` is still treated as set, distinguishing it + from the default -- exactly the auto-fill-vs-validate signal the scenario + validator needs. The converter writes endpoint keys only for fields in + ``cli.model_fields_set & ENDPOINT_FIELDS``, so membership is faithful and + no explicit-set sentinel is required on this converter path. + """ + endpoint = _benchmark_for_cli(**kwargs).endpoint + assert (field in endpoint.model_fields_set) is present diff --git a/tests/unit/config/test_scenario_resolver.py b/tests/unit/config/test_scenario_resolver.py new file mode 100644 index 0000000000..cc1a28adcc --- /dev/null +++ b/tests/unit/config/test_scenario_resolver.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the ScenarioResolver step and its place in the default chain.""" + +from __future__ import annotations + +from pathlib import Path + +from aiperf.common.enums import CacheBustTarget +from aiperf.config import BenchmarkConfig +from aiperf.config.resolution.plan import BenchmarkRun +from aiperf.config.resolution.resolvers import ( + DatasetResolver, + ScenarioResolver, + TimingResolver, + build_default_resolver_chain, +) + +_WEKA_FIXTURE = ( + Path(__file__).parents[2] / "unit/graph/fixtures/weka_min.json" +).resolve() + + +def _make_graph_run(*, scenario: str | None) -> BenchmarkRun: + cfg = BenchmarkConfig( + models=["claude-opus-4-5-20251101"], + endpoint={"urls": ["http://localhost:8000/v1/chat/completions"]}, + datasets=[{"name": "profiling", "type": "file", "path": str(_WEKA_FIXTURE)}], + phases=[ + { + "name": "profiling", + "type": "concurrency", + "concurrency": 1, + "sessions": 5, + } + ], + scenario=scenario, + ) + return BenchmarkRun(benchmark_id="test-run", cfg=cfg, artifact_dir=Path("/tmp/x")) + + +def test_scenario_resolver_populates_outcome_when_scenario_set() -> None: + run = _make_graph_run(scenario="inferencex-agentx-mvp") + ScenarioResolver().resolve(run) + outcome = run.resolved.scenario_outcome + assert outcome is not None + assert outcome.submission_valid is True + assert outcome.scenario_name == "inferencex-agentx-mvp" + # The resolver's auto-fills landed on the live config. + assert run.cfg.endpoint.cache_bust == CacheBustTarget.FIRST_TURN_PREFIX + + +def test_scenario_resolver_noop_when_scenario_unset() -> None: + run = _make_graph_run(scenario=None) + ScenarioResolver().resolve(run) + assert run.resolved.scenario_outcome.submission_valid is None + # No auto-fill when no scenario. + assert run.cfg.endpoint.cache_bust == CacheBustTarget.NONE + + +def test_scenario_resolver_runs_between_dataset_and_timing() -> None: + chain = build_default_resolver_chain()._resolvers + types = [type(r) for r in chain] + assert ScenarioResolver in types + assert types.index(DatasetResolver) < types.index(ScenarioResolver) + assert types.index(ScenarioResolver) < types.index(TimingResolver) + + +def test_scenario_resolver_autofilled_duration_visible_to_timing() -> None: + """ScenarioResolver runs before TimingResolver, so the auto-filled + 1800s duration is summed into total_expected_duration.""" + run = _make_graph_run(scenario="inferencex-agentx-mvp") + ScenarioResolver().resolve(run) + TimingResolver().resolve(run) + assert run.resolved.total_expected_duration == 1800.0 diff --git a/tests/unit/config/test_session_routing_config.py b/tests/unit/config/test_session_routing_config.py new file mode 100644 index 0000000000..83c3f22c7a --- /dev/null +++ b/tests/unit/config/test_session_routing_config.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from pytest import param + +from aiperf.config.endpoint import EndpointConfig +from aiperf.config.flags._converter_endpoint import _parse_routing_opts + + +def _config(**kwargs) -> EndpointConfig: + return EndpointConfig(urls=["http://localhost:8000"], **kwargs) + + +def test_defaults_off(): + config = _config() + assert config.session_routing is None + assert config.session_routing_opts == {} + + +def test_mode_with_valid_opts(): + config = _config( + session_routing="dynamo_nvext", + session_routing_opts={"timeout_seconds": "600"}, + ) + assert str(config.session_routing) == "dynamo_nvext" + + +def test_opts_canonicalized_to_typed_values(): + config = _config( + session_routing="dynamo_nvext", + session_routing_opts={"timeout_seconds": "600"}, + ) + assert config.session_routing_opts == {"timeout_seconds": 600} + assert isinstance(config.session_routing_opts["timeout_seconds"], int) + + +def test_canonicalization_idempotent(): + config = _config( + session_routing="dynamo_nvext", + session_routing_opts={"timeout_seconds": 600}, + ) + assert config.session_routing_opts == {"timeout_seconds": 600} + + +def test_opts_without_mode_rejected(): + with pytest.raises(ValueError, match="session-routing-opt"): + _config(session_routing_opts={"header_name": "X-A"}) + + +def test_unknown_opt_key_rejected(): + with pytest.raises(ValueError, match="timeout_secs"): + _config( + session_routing="dynamo_nvext", + session_routing_opts={"timeout_secs": "600"}, + ) + + +def test_invalid_opt_value_rejected(): + with pytest.raises(ValueError): + _config( + session_routing="dynamo_nvext", + session_routing_opts={"timeout_seconds": "0"}, + ) + + +def test_parameterless_mode_rejects_any_opt(): + with pytest.raises(ValueError): + _config( + session_routing="smg_routing_key", + session_routing_opts={"anything": "x"}, + ) + + +def test_parse_routing_opts_duplicate_key_rejected(): + with pytest.raises(ValueError, match="Duplicate"): + _parse_routing_opts(["header_name=X-A", "header_name=X-B"]) + + +@pytest.mark.parametrize( + "item", + [ + param("noequals", id="no_separator"), + param("key=", id="empty_value"), + ], +) # fmt: skip +def test_parse_routing_opts_malformed_pair_rejected(item): + with pytest.raises(ValueError, match="expected non-empty key=value"): + _parse_routing_opts([item]) diff --git a/tests/unit/config/test_sweep_round2_adversarial.py b/tests/unit/config/test_sweep_round2_adversarial.py index e0bda8be12..7051a81be1 100644 --- a/tests/unit/config/test_sweep_round2_adversarial.py +++ b/tests/unit/config/test_sweep_round2_adversarial.py @@ -1,13 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Round-2 adversarial regressions for the sweep + config system. +"""Adversarial regressions for the sweep + config system. Each class is named by hypothesis ID (H9-H14, continuing from ``test_sweep_qmc_adversarial.py``). Each test pins a behavior that -silently misbehaved or crashed unhelpfully before the round-2 fixes: +would otherwise silently misbehave or crash unhelpfully: -- H9: QMC sweeps now body-root paths (no more phantom top-level keys - that ``build_benchmark_plan`` silently discards). +- H9: QMC sweeps body-root paths (phantom top-level keys would be + silently discarded by ``build_benchmark_plan``). - H10: Path validators reject non-sweepable top-level prefixes (multi_run, random_seed, benchmark) and the redundant ``benchmark.`` prefix. @@ -128,8 +128,8 @@ def test_sobol_envelope_variables_paired_sweep(self, base: dict) -> None: assert variant["variables"]["foo"] == sampled def test_e2e_build_benchmark_plan_applies_sampled_concurrency(self) -> None: - # End-to-end: confirm the bug repro from the round-2 report passes. - # Before the fix, all 8 variants ended up with concurrency=8. + # End-to-end: without per-variant application of the sampled value, + # all 8 variants would end up with concurrency=8. from aiperf.config.config import AIPerfConfig from aiperf.config.loader.plan import build_benchmark_plan diff --git a/tests/unit/config/test_sweep_round3_adversarial.py b/tests/unit/config/test_sweep_round3_adversarial.py index f2aa8290bd..57e353bc77 100644 --- a/tests/unit/config/test_sweep_round3_adversarial.py +++ b/tests/unit/config/test_sweep_round3_adversarial.py @@ -1,11 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Round-3 adversarial regressions for the sweep + config system. +"""Adversarial regressions for the sweep + config system. Continues the H-series after ``test_sweep_round2_adversarial.py``: - H15: ``_set_nested_value`` raises a clear path-aware ValueError when a - dotted path crosses a scalar/None mid-traversal (was: opaque + dotted path crosses a scalar/None mid-traversal (otherwise: opaque ``TypeError: 'X' object does not support item assignment``). - H16: Grid sweep parameter keys go through the shared dotted-path validator -- empty/leading-dot/``..``/non-sweepable-prefix paths are diff --git a/tests/unit/config/test_sweep_round3_fixes.py b/tests/unit/config/test_sweep_round3_fixes.py index 6e46d64af4..a12871274c 100644 --- a/tests/unit/config/test_sweep_round3_fixes.py +++ b/tests/unit/config/test_sweep_round3_fixes.py @@ -1,16 +1,16 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Round-3 regression tests for round-2 adversarial findings. +"""Sweep-sampling edge-case regression tests. Covers: -- R2-H8/S1: ``SamplingDimension.choices`` rejects NaN/inf numeric entries. -- R2-L1/S4: ``_map_dim`` clamps int output to respect declared ``[lo, hi]`` +- ``SamplingDimension.choices`` rejects NaN/inf numeric entries. +- ``_map_dim`` clamps int output to respect declared ``[lo, hi]`` (banker's rounding could otherwise produce values below ``lo`` for ``kind=int, scale=log, lo<1``). -- R2-L2/S6 cross-cut: ``is_finite_value`` correctly handles - ``numpy.float32`` NaN/inf (the stdlib-float-only check would have - missed it). +- ``is_finite_value`` correctly handles + ``numpy.float32`` NaN/inf (a stdlib-float-only check would + miss it). """ from __future__ import annotations @@ -29,9 +29,9 @@ class TestChoicesFiniteValidation: """SamplingDimension.choices must reject non-finite numeric entries. - Round-1's ``cd953cb37`` added ``_validate_finite_bounds`` for - lo/hi only; ``choices`` slipped through and crashed mid-flight at - the orchestrator's defensive guard. Now blocks at validation time. + ``_validate_finite_bounds`` covers lo/hi only; without the same check on + ``choices``, a non-finite entry slips through and crashes mid-flight at + the orchestrator's defensive guard instead of blocking at validation time. """ @pytest.mark.parametrize( @@ -75,7 +75,7 @@ def test_choices_with_nonfinite_among_strings_still_rejected(self) -> None: class TestIntLogScaleClamping: """``_map_dim`` clamps int output so it never violates declared bounds. - Pre-existing bug surfaced by round-1's narrow-int-range warning: + Without the clamp, ``int(round(exp(log(0.5))))=0`` for ``lo=0.5, u=0.0``, which is below the user-declared ``lo``. """ @@ -110,8 +110,8 @@ def test_int_mapping_never_below_lo_or_above_hi( assert v <= hi_floor, f"u={u}: produced {v}, above floor(hi)={hi_floor}" def test_int_log_lo_half_u_zero_clamps_to_one(self) -> None: - """Direct repro of the round-2 finding: lo=0.5, log scale, - u=0.0 used to map to 0; must now clamp to 1.""" + """Direct repro: lo=0.5, log scale, + u=0.0 would map to 0 without the clamp; must clamp to 1.""" with pytest.warns(UserWarning, match="too narrow for sampling diversity"): dim = SamplingDimension(path="x", lo=0.5, hi=1.0, kind="int", scale="log") assert _map_dim(0.0, dim) == 1 @@ -120,7 +120,7 @@ def test_int_log_lo_half_u_zero_clamps_to_one(self) -> None: class TestIsFiniteValueNumpyFloat32: """Cross-cutting smoke for ``is_finite_value``. - Round-2 S6 noted that ``isinstance(x, float)`` misses + ``isinstance(x, float)`` misses ``numpy.float32``. ``is_finite_value`` duck-types via ``math.isfinite(float(x))`` so it handles all numpy float widths uniformly. diff --git a/tests/unit/config/test_sweep_round4_adversarial.py b/tests/unit/config/test_sweep_round4_adversarial.py index 7b7b6936cd..c5810e0588 100644 --- a/tests/unit/config/test_sweep_round4_adversarial.py +++ b/tests/unit/config/test_sweep_round4_adversarial.py @@ -1,18 +1,18 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Round-4 adversarial regressions for the sweep + config system. +"""Adversarial regressions for the sweep + config system. Continues the H-series after ``test_sweep_round3_adversarial.py``: - H20: ``SLAFilter`` rejects NaN/inf threshold and empty/whitespace - ``metric_tag`` (was: silent feasibility false-fail when threshold is + ``metric_tag`` (otherwise: silent feasibility false-fail when threshold is NaN; meaningless filter when metric_tag is empty). - H21: ``PostProcessSpec`` rejects empty handler, NUL bytes in ``output_filename``, and dot-only stems like ``..json`` / ``.json`` - (was: hidden-file confusion + OS-level open failures at write time). + (otherwise: hidden-file confusion + OS-level open failures at write time). - H22: ``detect_sweep_fields`` is scoped to phase-rooted paths only -- - a list at ``datasets.X.prompts.isl.mean = [100, 200]`` no longer - silently auto-sweeps and produces validation-failing variants. + a list at ``datasets.X.prompts.isl.mean = [100, 200]`` must not + silently auto-sweep and produce validation-failing variants. """ from __future__ import annotations diff --git a/tests/unit/credit/test_first_token_event_flag.py b/tests/unit/credit/test_first_token_event_flag.py new file mode 100644 index 0000000000..04515501da --- /dev/null +++ b/tests/unit/credit/test_first_token_event_flag.py @@ -0,0 +1,271 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-credit first-token event flag + graph first-token observer. + +Credit-pipeline wiring for post-TTFT first-token anchoring: + +- ``TurnToSend.first_token_event`` copies into the issued ``Credit`` on BOTH + the linear (``issue_credit``) and graph (``issue_graph_credit``) paths, and + propagates across turns via ``TurnToSend.from_previous_credit``. The known + three-touch credit-pipeline trap: without the explicit copy in the + ``Credit(...)`` construction the flag silently drops while tests that only + assert on the ``TurnToSend`` still pass. +- ``CreditCallbackHandler.set_graph_first_token_observer`` fires for every + ``FirstToken`` carrying a ``trace_id`` (a graph credit), independent of + phase-handler gating -- mirroring the unconditional graph-return observer. +- ``on_first_token`` must NOT advance the prefill-released counter when prefill + limiting is inactive: a graph first-token event arrives even then, and + counting it would inflate the stat. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +from aiperf.common.enums import CreditPhase +from aiperf.credit.callback_handler import CreditCallbackHandler +from aiperf.credit.issuer import CreditIssuer +from aiperf.credit.messages import FirstToken +from aiperf.credit.structs import Credit, TurnToSend + +# asyncio_mode = "auto" (pyproject) runs the async tests without per-test marks. + + +# ============================================================================= +# Issuer path — flag copies TurnToSend -> Credit +# ============================================================================= + + +class RecordingConcurrency: + """Grants every slot; records prefill acquisitions (mirrors the graph-path + issuer test's recording fake).""" + + def __init__(self) -> None: + self.prefill_acquires = 0 + + async def acquire_session_slot(self, phase, can_proceed_fn) -> bool: + return can_proceed_fn() + + def release_session_slot(self, phase) -> None: ... + + async def acquire_prefill_slot(self, phase, can_proceed_fn) -> bool: + self.prefill_acquires += 1 + return can_proceed_fn() + + def release_prefill_slot(self, phase) -> None: ... + + +class FakeProgress: + def increment_sent(self, turn) -> tuple[int, bool]: + return (0, False) + + def freeze_sent_counts(self) -> None: ... + + +class FakeStopChecker: + def can_start_new_session(self) -> bool: + return True + + def can_send_any_turn(self) -> bool: + return True + + def can_send_dag_child_turn(self) -> bool: + return True + + +class FakeRouter: + def __init__(self) -> None: + self.sent: list = [] + + async def send_credit(self, *, credit) -> None: + self.sent.append(credit) + + +class FakeCancellation: + def next_cancellation_delay_ns(self, turn, phase) -> int | None: + return None + + +class FakeLifecycle: + started_at_ns = 0 + started_at_perf_ns = 0 + + +def _issuer(router: FakeRouter) -> CreditIssuer: + return CreditIssuer( + phase=CreditPhase.PROFILING, + stop_checker=FakeStopChecker(), + progress=FakeProgress(), + concurrency_manager=RecordingConcurrency(), + credit_router=router, + cancellation_policy=FakeCancellation(), + lifecycle=FakeLifecycle(), + ) + + +class TestFirstTokenEventFlagCopy: + """The per-credit ``first_token_event`` flag must survive the + TurnToSend -> Credit hand-off on every issuance path.""" + + async def test_graph_credit_copies_flag_true(self): + router = FakeRouter() + turn = TurnToSend( + conversation_id="t0", + x_correlation_id="x-t0", + turn_index=0, + num_turns=1, + trace_id="t0#0", + node_ordinal=0, + first_token_event=True, + ) + + await _issuer(router).issue_graph_credit(turn) + + assert router.sent[0].first_token_event is True + + async def test_graph_credit_flag_defaults_false(self): + router = FakeRouter() + turn = TurnToSend( + conversation_id="t0", + x_correlation_id="x-t0", + turn_index=0, + num_turns=1, + trace_id="t0#0", + node_ordinal=0, + ) + + await _issuer(router).issue_graph_credit(turn) + + assert router.sent[0].first_token_event is False + + async def test_linear_credit_copies_flag_true(self): + router = FakeRouter() + turn = TurnToSend( + conversation_id="c", + x_correlation_id="x", + turn_index=0, + num_turns=1, + first_token_event=True, + ) + + await _issuer(router).issue_credit(turn) + + assert router.sent[0].first_token_event is True + + def test_from_previous_credit_propagates_flag(self): + credit = Credit( + id=0, + phase=CreditPhase.PROFILING, + conversation_id="c", + x_correlation_id="x", + turn_index=0, + num_turns=2, + issued_at_ns=1, + first_token_event=True, + ) + + nxt = TurnToSend.from_previous_credit(credit) + + assert nxt.first_token_event is True + + +# ============================================================================= +# Callback handler — graph first-token observer + prefill-stat guard +# ============================================================================= + + +def _make_handler(prefill_enabled: bool): + """Build a handler with one registered PROFILING phase. ``concurrency`` and + ``progress`` are MagicMocks so callers can assert on the prefill counters.""" + concurrency = MagicMock() + concurrency.prefill_limiting_enabled = prefill_enabled + concurrency.release_prefill_slot = MagicMock() + + progress = MagicMock() + progress.increment_prefill_released = MagicMock() + + handler = CreditCallbackHandler(concurrency) + handler.register_phase( + phase=CreditPhase.PROFILING, + progress=progress, + lifecycle=MagicMock(is_complete=False), + stop_checker=MagicMock(), + strategy=MagicMock(handle_credit_return=AsyncMock()), + ) + return handler, concurrency, progress + + +def _graph_first_token(**overrides) -> FirstToken: + kwargs = dict( + credit_id=1, + phase=CreditPhase.PROFILING, + ttft_ns=5, + trace_id="t0#0", + x_correlation_id="x-t0", + turn_index=0, + ) + kwargs.update(overrides) + return FirstToken(**kwargs) + + +class TestGraphFirstTokenObserver: + async def test_observer_fires_for_graph_first_token(self): + handler, _, _ = _make_handler(prefill_enabled=True) + seen: list[FirstToken] = [] + handler.set_graph_first_token_observer(seen.append) + + ft = _graph_first_token() + await handler.on_first_token(ft) + + assert seen == [ft] + + async def test_observer_not_fired_for_non_graph_first_token(self): + handler, _, _ = _make_handler(prefill_enabled=True) + seen: list[FirstToken] = [] + handler.set_graph_first_token_observer(seen.append) + + await handler.on_first_token(_graph_first_token(trace_id=None)) + + assert seen == [] + + async def test_observer_none_is_noop(self): + handler, _, _ = _make_handler(prefill_enabled=True) + handler.set_graph_first_token_observer(None) + + # Must not raise even for a graph-carrying first token. + await handler.on_first_token(_graph_first_token()) + + async def test_observer_fires_even_for_unregistered_phase(self): + """Mirrors the return observer: fires independent of phase-handler + registration so a graph first-token event is never stranded.""" + handler = CreditCallbackHandler(MagicMock()) + seen: list[FirstToken] = [] + handler.set_graph_first_token_observer(seen.append) + + ft = _graph_first_token(phase=CreditPhase.WARMUP) + await handler.on_first_token(ft) + + assert seen == [ft] + + +class TestFirstTokenPrefillStatGuard: + async def test_prefill_counter_not_touched_when_limiting_disabled(self): + handler, concurrency, progress = _make_handler(prefill_enabled=False) + seen: list[FirstToken] = [] + handler.set_graph_first_token_observer(seen.append) + + ft = _graph_first_token() + await handler.on_first_token(ft) + + progress.increment_prefill_released.assert_not_called() + # Observer still fires; slot release is internally guarded, so it is + # harmless if invoked, but the counter must not advance. + assert seen == [ft] + + async def test_prefill_counter_touched_when_limiting_enabled(self): + handler, concurrency, progress = _make_handler(prefill_enabled=True) + + await handler.on_first_token(_graph_first_token()) + + progress.increment_prefill_released.assert_called_once() + concurrency.release_prefill_slot.assert_called_once_with(CreditPhase.PROFILING) diff --git a/tests/unit/credit/test_graph_sticky_lifecycle.py b/tests/unit/credit/test_graph_sticky_lifecycle.py new file mode 100644 index 0000000000..e7e9e7e30e --- /dev/null +++ b/tests/unit/credit/test_graph_sticky_lifecycle.py @@ -0,0 +1,284 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Instance-keyed graph sticky routing + GraphTraceEnd lifecycle. + +Linear sessions key on ``x_correlation_id`` (legacy contract); graph credits +(``trace_id`` set) key their session on the trace INSTANCE id, so EVERY +trajectory of one instance shares one session -- and therefore one worker -- +because dynamic-slot capture/splice pools and spawn state are worker-local +and keyed by the instance. Graph credits mint ``turn_index`` per node, so a +trajectory's recorded final turn can complete while spawned work is still in +flight -- turn counting cannot express trace lifecycle. These tests pin the +router rules that make graph stickiness work: + +1. graph credits create their instance session on first sight and are NOT + torn down by the final-turn cleanup (without the gate, the session is + destroyed on the same credit that creates it -- invisible to + accounting-only assertions, hence the multi-credit PLACEMENT assertions); +2. every trajectory (fresh corr, same instance) joins the instance session; +3. ``end_graph_trace`` owns the close: ONE call per instance, synchronous + state mutation, idempotent, dead-worker safe, worker-forward only when a + session existed; +4. linear (no ``trace_id``) credits keep byte-identical behavior. + +Plus the issuer side: the ``end_graph_trace`` passthrough to the router. +(Template-hash URL affinity is pinned in ``test_issuer_graph_path.py``.) +""" + +import asyncio +import time +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from aiperf.common.enums import CreditPhase +from aiperf.credit.issuer import CreditIssuer +from aiperf.credit.messages import GraphTraceEnd +from aiperf.credit.sticky_router import StickyCreditRouter +from aiperf.credit.structs import Credit + + +def _graph_credit( + credit_id: int, + *, + corr: str, + node_turn: int = 0, + trace_id: str = "t-1::inst0", +) -> Credit: + """A graph credit as the adapter mints it: per-trajectory corr, per-node + turn_index, num_turns=1 (every fire looks final to turn counting).""" + return Credit( + id=credit_id, + phase=CreditPhase.PROFILING, + conversation_id=trace_id.split("::", 1)[0], + x_correlation_id=corr, + turn_index=node_turn, + num_turns=1, + issued_at_ns=0, + trace_id=trace_id, + node_ordinal=node_turn, + ) + + +def _linear_credit(credit_id: int, *, corr: str, turn: int, num_turns: int) -> Credit: + return Credit( + id=credit_id, + phase=CreditPhase.PROFILING, + conversation_id="conv", + x_correlation_id=corr, + turn_index=turn, + num_turns=num_turns, + issued_at_ns=0, + ) + + +def _router(benchmark_run, workers: list[str]) -> StickyCreditRouter: + router = StickyCreditRouter(run=benchmark_run, service_id="test-router") + router._router_client.send_to = AsyncMock() + for w in workers: + router._register_worker(w) + return router + + +class TestGraphStickyPlacement: + async def test_instance_pins_one_worker_despite_final_turns( + self, benchmark_run + ) -> None: + """The placement assertion: N final-looking credits, ONE worker, + ONE instance-keyed session.""" + router = _router(benchmark_run, ["worker-A", "worker-B"]) + + for turn in range(3): + await router.send_credit( + _graph_credit(turn, corr="t-1::c1", node_turn=turn) + ) + + placements = { + call[0][0] for call in router._router_client.send_to.call_args_list + } + assert len(placements) == 1 + assert router._sticky_sessions.get("t-1::inst0") in placements + worker_id = router._sticky_sessions["t-1::inst0"] + assert router._workers[worker_id].active_sessions == 1 + assert "t-1::inst0" in router._workers[worker_id].active_session_ids + + async def test_distinct_instances_balance_across_workers( + self, benchmark_run + ) -> None: + router = _router(benchmark_run, ["worker-A", "worker-B"]) + + await router.send_credit(_graph_credit(1, corr="t-1::c1")) + await router.send_credit( + _graph_credit(2, corr="t-2::c2", trace_id="t-2::inst0") + ) + + assert ( + router._sticky_sessions["t-1::inst0"] + != router._sticky_sessions["t-2::inst0"] + ) + + async def test_trajectories_of_one_instance_co_place(self, benchmark_run) -> None: + """INSTANCE co-placement: a child trajectory (own corr, same + credit.trace_id) shares the instance session -- worker-local + dynamic-slot/spawn state must be reachable from every trajectory.""" + router = _router(benchmark_run, ["worker-A", "worker-B"]) + + await router.send_credit(_graph_credit(1, corr="t-1::root-c")) + await router.send_credit(_graph_credit(2, corr="t-1::child-c")) + + placements = { + call[0][0] for call in router._router_client.send_to.call_args_list + } + assert len(placements) == 1 + worker_id = router._sticky_sessions["t-1::inst0"] + assert placements == {worker_id} + # ONE session for the whole instance, not one per trajectory. + assert router._workers[worker_id].active_sessions == 1 + + async def test_instance_reroutes_after_worker_death(self, benchmark_run) -> None: + """A dead sticky worker releases the instance session: the next + credit re-places on a survivor instead of raising or routing dead.""" + router = _router(benchmark_run, ["worker-A", "worker-B"]) + await router.send_credit(_graph_credit(1, corr="t-1::root-c")) + dead_worker = router._sticky_sessions["t-1::inst0"] + + router._unregister_worker(dead_worker) + assert "t-1::inst0" not in router._sticky_sessions + + await router.send_credit(_graph_credit(2, corr="t-1::child-c")) + survivor = ({"worker-A", "worker-B"} - {dead_worker}).pop() + assert router._sticky_sessions["t-1::inst0"] == survivor + + async def test_linear_final_turn_creates_no_session(self, benchmark_run) -> None: + """Existing linear behavior is byte-identical (no trace_id).""" + router = _router(benchmark_run, ["worker-A"]) + + await router.send_credit(_linear_credit(1, corr="c1", turn=0, num_turns=1)) + + assert router._sticky_sessions == {} + assert router._workers["worker-A"].active_sessions == 0 + + async def test_linear_multi_turn_lifecycle_unchanged(self, benchmark_run) -> None: + router = _router(benchmark_run, ["worker-A"]) + + await router.send_credit(_linear_credit(1, corr="c1", turn=0, num_turns=2)) + assert "c1" in router._sticky_sessions + await router.send_credit(_linear_credit(2, corr="c1", turn=1, num_turns=2)) + assert "c1" not in router._sticky_sessions + assert router._workers["worker-A"].active_sessions == 0 + + +class TestGraphTraceEndLifecycle: + async def test_end_closes_instance_session_and_forwards_to_worker( + self, benchmark_run + ) -> None: + """ONE end call closes the whole instance (all trajectories) and the + worker receives exactly ONE GraphTraceEnd.""" + router = _router(benchmark_run, ["worker-A"]) + await router.send_credit(_graph_credit(1, corr="t-1::root-c")) + await router.send_credit(_graph_credit(2, corr="t-1::child-c")) + router._router_client.send_to.reset_mock() + + await router.end_graph_trace("t-1::inst0", "profiling") + + assert "t-1::inst0" not in router._sticky_sessions + assert router._workers["worker-A"].active_sessions == 0 + assert "t-1::inst0" not in router._workers["worker-A"].active_session_ids + router._router_client.send_to.assert_called_once() + worker_id, message = router._router_client.send_to.call_args[0] + assert worker_id == "worker-A" + assert message == GraphTraceEnd( + trace_id="t-1::inst0", phase_variant="profiling" + ) + + async def test_end_is_idempotent_and_no_session_is_noop( + self, benchmark_run + ) -> None: + router = _router(benchmark_run, ["worker-A"]) + await router.send_credit(_graph_credit(1, corr="t-1::c1")) + await router.end_graph_trace("t-1::inst0", "profiling") + router._router_client.send_to.reset_mock() + + await router.end_graph_trace("t-1::inst0", "profiling") + await router.end_graph_trace("t-9::never", "profiling") + + router._router_client.send_to.assert_not_called() + assert router._workers["worker-A"].active_sessions == 0 + + async def test_end_after_worker_unregistered_is_safe(self, benchmark_run) -> None: + router = _router(benchmark_run, ["worker-A", "worker-B"]) + await router.send_credit(_graph_credit(1, corr="t-1::c1")) + worker_id = router._sticky_sessions["t-1::inst0"] + router._unregister_worker(worker_id) + router._router_client.send_to.reset_mock() + + await router.end_graph_trace("t-1::inst0", "profiling") + + assert "t-1::inst0" not in router._sticky_sessions + router._router_client.send_to.assert_not_called() + + async def test_unregister_cleans_graph_sessions(self, benchmark_run) -> None: + router = _router(benchmark_run, ["worker-A"]) + await router.send_credit(_graph_credit(1, corr="t-1::c1")) + + router._unregister_worker("worker-A") + + assert "t-1::inst0" not in router._sticky_sessions + + async def test_same_instance_after_end_creates_fresh_session( + self, benchmark_run + ) -> None: + router = _router(benchmark_run, ["worker-A"]) + await router.send_credit(_graph_credit(1, corr="t-1::c1")) + await router.end_graph_trace("t-1::inst0", "profiling") + + await router.send_credit(_graph_credit(2, corr="t-1::c2", node_turn=1)) + + assert router._sticky_sessions["t-1::inst0"] == "worker-A" + assert router._workers["worker-A"].active_sessions == 1 + + +def _issuer(router: MagicMock) -> CreditIssuer: + progress = MagicMock() + progress.increment_sent = MagicMock(return_value=(1, False)) + progress.all_credits_sent_event = asyncio.Event() + lifecycle = MagicMock() + lifecycle.started_at_ns = time.time_ns() + lifecycle.started_at_perf_ns = time.perf_counter_ns() + cancellation = MagicMock() + cancellation.next_cancellation_delay_ns = MagicMock(return_value=None) + return CreditIssuer( + phase=CreditPhase.PROFILING, + stop_checker=MagicMock(), + progress=progress, + concurrency_manager=MagicMock(), + credit_router=router, + cancellation_policy=cancellation, + lifecycle=lifecycle, + url_selection_strategy=None, + ) + + +class TestIssuerGraphTraceEndPassthrough: + async def test_end_graph_trace_forwards_to_router(self) -> None: + router = MagicMock() + router.send_credit = AsyncMock() + router.end_graph_trace = AsyncMock() + issuer = _issuer(router) + + await issuer.end_graph_trace("t-1::inst0", "profiling") + + router.end_graph_trace.assert_awaited_once_with("t-1::inst0", "profiling") + + +@pytest.mark.asyncio +async def test_graph_trace_end_survives_msgpack_roundtrip() -> None: + import msgspec + + from aiperf.credit.messages import RouterToWorkerMessage + + msg = GraphTraceEnd(trace_id="t-1::inst0", phase_variant="profiling") + decoded = msgspec.msgpack.Decoder(RouterToWorkerMessage).decode( + msgspec.msgpack.Encoder().encode(msg) + ) + assert decoded == msg diff --git a/tests/unit/credit/test_issuer_finality.py b/tests/unit/credit/test_issuer_finality.py new file mode 100644 index 0000000000..4306b0006e --- /dev/null +++ b/tests/unit/credit/test_issuer_finality.py @@ -0,0 +1,232 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Guards the credit issuer's lineage-finality stamp. + +Touch #1 (the ``Credit`` struct fields) and touch #3 (worker -> ``RequestInfo``, +``test_create_request_info_plumbs_finality_from_credit``) are covered elsewhere. +This file covers the issuer touch: ``CreditIssuer`` reading per-tree state from +a REAL ``SessionTreeRegistry`` (``_finality_for_issue``) AND stamping the result +onto the emitted ``Credit`` at its sole construction site +(``_issue_credit_internal``). + +Both the registry and the emitted ``Credit`` are real objects here on purpose -- +a MagicMock would auto-create ``is_parent_final`` / ``is_tree_final`` and pass +even if the ``Credit(...)`` kwargs were deleted, defeating the guard. +""" + +import time +from unittest.mock import MagicMock + +import pytest + +from aiperf.common.enums import CreditPhase +from aiperf.credit.issuer import CreditIssuer +from aiperf.credit.structs import Credit, TurnToSend +from aiperf.timing.session_tree import SessionTreeRegistry + + +class _FakeConcurrency: + """Slots always granted; releases are no-ops. + + The main ``SessionTreeRegistry`` needs no concurrency manager at all; the + issuer's acquire path needs the two coroutines. Neither the registry nor the + emitted ``Credit`` is a mock -- that is the point of this file. + """ + + async def acquire_session_slot(self, phase: CreditPhase, can_proceed) -> bool: + return True + + async def acquire_prefill_slot(self, phase: CreditPhase, can_proceed) -> bool: + return True + + def release_session_slot(self, phase: CreditPhase) -> None: + pass + + +class _CapturingRouter: + """Captures the emitted ``Credit`` so tests can assert its stamped finality.""" + + def __init__(self) -> None: + self.sent: list[Credit] = [] + + async def send_credit(self, *, credit: Credit) -> None: + self.sent.append(credit) + + +def _make_registry() -> SessionTreeRegistry: + return SessionTreeRegistry() + + +def _make_issuer( + registry: SessionTreeRegistry | None, +) -> tuple[CreditIssuer, _CapturingRouter]: + """Minimal real issuer: mocked scalars/lifecycle, REAL registry + router. + + ``session_tree_registry_enabled=True`` engages ``registry`` regardless of + phase; passing ``None`` for the registry yields the non-DAG path. + """ + progress = MagicMock() + progress.increment_sent = MagicMock(return_value=(1, False)) + progress.freeze_sent_counts = MagicMock() + progress.all_credits_sent_event = MagicMock() + + stop_checker = MagicMock() + stop_checker.can_send_any_turn = MagicMock(return_value=True) + stop_checker.can_start_new_session = MagicMock(return_value=True) + stop_checker.can_send_dag_child_turn = MagicMock(return_value=True) + + cancellation = MagicMock() + cancellation.next_cancellation_delay_ns = MagicMock(return_value=None) + + lifecycle = MagicMock() + lifecycle.started_at_ns = time.time_ns() + lifecycle.started_at_perf_ns = time.perf_counter_ns() + + router = _CapturingRouter() + issuer = CreditIssuer( + phase=CreditPhase.PROFILING, + stop_checker=stop_checker, + progress=progress, + concurrency_manager=_FakeConcurrency(), + credit_router=router, + cancellation_policy=cancellation, + lifecycle=lifecycle, + session_tree_registry=registry, + session_tree_registry_enabled=True, + ) + return issuer, router + + +def _root_turn(root_id: str = "root-1") -> TurnToSend: + """Depth-0 root, single-turn (final), no forks.""" + return TurnToSend( + conversation_id="conv-1", + x_correlation_id=root_id, + turn_index=0, + num_turns=1, + ) + + +def _child_turn(root_id: str = "root-1", child_id: str = "child-1") -> TurnToSend: + """Child whose parent IS the root, single-turn (final), no forks.""" + return TurnToSend( + conversation_id="conv-1", + x_correlation_id=child_id, + turn_index=0, + num_turns=1, + agent_depth=1, + parent_correlation_id=root_id, + root_correlation_id=root_id, + ) + + +# ============================================================================= +# _finality_for_issue: reads REAL SessionTreeRegistry state +# ============================================================================= + + +def test_finality_root_final_turn_no_descendants_is_tree_final(): + """Scenario 1: root, final turn, no descendants, no forks.""" + registry = _make_registry() + registry.open_tree("root-1", CreditPhase.PROFILING, root_pending=True) + issuer, _ = _make_issuer(registry) + + is_parent_final, is_tree_final = issuer._finality_for_issue(_root_turn()) + + assert is_parent_final is None + assert is_tree_final is True + + +def test_finality_root_with_outstanding_descendant_not_tree_final(): + """Scenario 2: root, final turn, one outstanding descendant -> not last.""" + registry = _make_registry() + registry.open_tree("root-1", CreditPhase.PROFILING, root_pending=True) + registry.register_descendants("root-1", n=1) + issuer, _ = _make_issuer(registry) + + is_parent_final, is_tree_final = issuer._finality_for_issue(_root_turn()) + + assert is_parent_final is None + assert is_tree_final is False + + +def test_finality_sole_child_after_root_terminal_is_both_final(): + """Scenario 3: child whose parent is the root; root terminal; sole + outstanding child on its final turn -> both facts True.""" + registry = _make_registry() + registry.open_tree("root-1", CreditPhase.PROFILING, root_pending=True) + registry.register_descendants("root-1", n=1) + registry.on_root_terminal("root-1") # root_pending cleared; tree still live + issuer, _ = _make_issuer(registry) + + is_parent_final, is_tree_final = issuer._finality_for_issue(_child_turn()) + + assert is_parent_final is True + assert is_tree_final is True + + +def test_finality_no_registry_is_conservative_none_false(): + """Scenario 4: no registry engaged -> conservative ``(None, False)``.""" + issuer, _ = _make_issuer(None) + + assert issuer._finality_for_issue(_root_turn()) == (None, False) + + +def test_finality_spawning_final_turn_never_tree_final(): + """Scenario 5 (regression): a final turn declaring ANY branch is never + tree-final, even with nothing outstanding in the registry. + + SPAWN children register only at return-intercept, AFTER this issue-time + stamp, so ``has_branches`` (any-mode) must gate the query -- the FORK-only + ``has_forks`` flag stays False for a SPAWN-declaring turn and previously + produced a wrong ``is_tree_final=True``. + """ + registry = _make_registry() + registry.open_tree("root-1", CreditPhase.PROFILING, root_pending=True) + issuer, _ = _make_issuer(registry) + + # Same registry state as scenario 1 (which yields True) but the turn + # declares a branch: SPAWN-shaped, so has_forks stays False. + spawning_root_final = TurnToSend( + conversation_id="conv-1", + x_correlation_id="root-1", + turn_index=0, + num_turns=1, + has_forks=False, + has_branches=True, + ) + is_parent_final, is_tree_final = issuer._finality_for_issue(spawning_root_final) + + assert is_parent_final is None + assert is_tree_final is False + + +# ============================================================================= +# GUARD: the Credit(...) construction site must pass the helper's results through +# ============================================================================= + + +@pytest.mark.asyncio +async def test_issue_credit_stamps_finality_onto_emitted_credit(): + """RED if either ``is_parent_final=`` / ``is_tree_final=`` kwarg is removed + from the ``Credit(...)`` construction in ``_issue_credit_internal``. + + Uses scenario 3 (both facts True) so the stamped values differ from the + struct defaults (``None`` / ``False``): a dropped kwarg reverts the emitted + ``Credit`` to the default and this assertion fails. + """ + registry = _make_registry() + registry.open_tree("root-1", CreditPhase.PROFILING, root_pending=True) + registry.register_descendants("root-1", n=1) + registry.on_root_terminal("root-1") + issuer, router = _make_issuer(registry) + + # Child inherits the root's slot; dispatch_child_turn is the wire path that + # runs _issue_credit_internal for a DAG child (issue_credit's session-slot + # path is for roots). + await issuer.dispatch_child_turn(_child_turn()) + + assert len(router.sent) == 1 + credit = router.sent[0] + assert credit.is_parent_final is True + assert credit.is_tree_final is True diff --git a/tests/unit/dataset/generator/test_coding_content_hash_ids.py b/tests/unit/dataset/generator/test_coding_content_hash_ids.py new file mode 100644 index 0000000000..25f7504035 --- /dev/null +++ b/tests/unit/dataset/generator/test_coding_content_hash_ids.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DM5: ``CodingContentGenerator.generate(hash_ids=...)`` with a default config. + +``PromptConfig.block_size`` defaults to ``None``; the hash-id path must fall +back to ``InputTokensDefaults.BLOCK_SIZE`` exactly like +``PromptGenerator.generate`` does, instead of raising ``TypeError`` on +``(len(hash_ids) - 1) * None``. +""" + +from __future__ import annotations + +from aiperf.config import PromptConfig +from aiperf.config.dataset.defaults import InputTokensDefaults +from aiperf.dataset.generator.coding_content import CodingContentGenerator +from tests.harness.fake_tokenizer import FakeTokenizer + + +def test_generate_hash_ids_default_config_falls_back_to_default_block_size() -> None: + generator = CodingContentGenerator(config=PromptConfig(), tokenizer=FakeTokenizer()) + + result = generator.generate(mean=InputTokensDefaults.BLOCK_SIZE, hash_ids=[1]) + + assert isinstance(result, str) and result + assert len(generator._cache[1]) == InputTokensDefaults.BLOCK_SIZE + + +def test_generate_hash_ids_explicit_block_size_still_wins() -> None: + generator = CodingContentGenerator(config=PromptConfig(), tokenizer=FakeTokenizer()) + + generator.generate(mean=64, hash_ids=[7], block_size=64) + + assert len(generator._cache[7]) == 64 diff --git a/tests/unit/dataset/graph/__init__.py b/tests/unit/dataset/graph/__init__.py new file mode 100644 index 0000000000..e5725ea5a4 --- /dev/null +++ b/tests/unit/dataset/graph/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unit/dataset/graph/adapters/__init__.py b/tests/unit/dataset/graph/adapters/__init__.py new file mode 100644 index 0000000000..d1e9cccf25 --- /dev/null +++ b/tests/unit/dataset/graph/adapters/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for Graph input-format adapters.""" diff --git a/tests/unit/dataset/graph/adapters/dag_jsonl/__init__.py b/tests/unit/dataset/graph/adapters/dag_jsonl/__init__.py new file mode 100644 index 0000000000..e5725ea5a4 --- /dev/null +++ b/tests/unit/dataset/graph/adapters/dag_jsonl/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unit/dataset/graph/adapters/dag_jsonl/test_lowering.py b/tests/unit/dataset/graph/adapters/dag_jsonl/test_lowering.py new file mode 100644 index 0000000000..545eda02d1 --- /dev/null +++ b/tests/unit/dataset/graph/adapters/dag_jsonl/test_lowering.py @@ -0,0 +1,796 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lowering tests: DagTree -> unified-segment-store ParsedGraph. + +Every case drives the REAL loader + tree expansion against a fixture file +(tmp_path line file or ``tests/fixtures/dag/*.dag.jsonl``) and then asserts +the byte-parity core contract: verbatim raw-message interning with prefix +dedup, live-reply slots for lineage producers, legacy-ordered dispatch +overrides, AND-fan-in gates, and START/END edge conventions. Structural +soundness is proven through ``validator.validate()`` (repo gotcha: node-level +assertions alone let structural bugs slip). +""" + +from pathlib import Path + +import orjson +import pytest +from pytest import param + +from aiperf.dataset.graph.adapters.dag_jsonl.lowering import lower_dag_trees +from aiperf.dataset.graph.adapters.dag_jsonl.tree import ( + expand_trees, + load_dag_conversations, +) +from aiperf.dataset.graph.models import ChannelRequirement, LlmNode, ParsedGraph +from aiperf.dataset.graph.segment_ir.envelope import read_prompt_segment_ids +from aiperf.dataset.graph.validator import validate + +FIXTURES_DIR = Path(__file__).parents[5] / "fixtures" / "dag" + +ALL_DAG_FIXTURES = [ + "background_fork.dag.jsonl", + "bg_fork_fanout.dag.jsonl", + "bg_fork_nested.dag.jsonl", + "bg_fork_with_spawn_join.dag.jsonl", + "full.dag.jsonl", + "multi_root_single_turn.dag.jsonl", + "small.dag.jsonl", + "spawn_minimal.dag.jsonl", +] + + +def _write_dag(tmp_path: Path, lines: list[dict]) -> Path: + path = tmp_path / "dag.jsonl" + path.write_bytes(b"\n".join(orjson.dumps(line) for line in lines)) + return path + + +def _lower( + path: Path, + *, + default_model: str | None = None, + run_streaming: bool = True, + endpoint_extra: list[tuple[str, object]] | None = None, +) -> ParsedGraph: + trees = expand_trees(load_dag_conversations(path, delay_cap_seconds=None)) + return lower_dag_trees( + trees, + default_model=default_model, + run_streaming=run_streaming, + endpoint_extra=endpoint_extra, + ) + + +def _trie(node: LlmNode) -> dict: + return node.metadata["trie"] + + +# (a) FORK fixture: parsed shape, node order, edges, inputs, record fields ---- + + +class TestForkFixtureStructure: + def test_parsed_shape_per_tree_graph_and_trace(self): + parsed = _lower(FIXTURES_DIR / "small.dag.jsonl") + assert set(parsed.graphs) == {"root"} + assert parsed.graph is parsed.graphs["root"] + assert [(t.id, t.graph_ref) for t in parsed.traces] == [("root", "root")] + assert parsed.segment_pool is not None + + def test_record_conventions_version_provenance_state_outputs(self): + graph = _lower(FIXTURES_DIR / "small.dag.jsonl").graphs["root"] + assert graph.version == "2.0" + assert graph.provenance.source == "dag_jsonl" + assert graph.provenance.tool not in ("", "manual") + assert set(graph.state) == {f"{nid}_out" for nid in graph.nodes} + for nid, node in graph.nodes.items(): + assert node.output == f"{nid}_out" + assert node.arrival_offset_us == 0 + assert node.streaming is True + + def test_node_order_edges_and_fan_in_inputs(self): + graph = _lower(FIXTURES_DIR / "small.dag.jsonl").graphs["root"] + assert list(graph.nodes) == [ + "root:0", + "branchA:0", + "branchA:1", + "branchB:0", + "branchB:1", + ] + edges = {(e.source, e.target): e for e in graph.edges} + assert set(edges) == { + ("START", "root:0"), + ("root:0", "branchA:0"), + ("branchA:0", "branchA:1"), + ("root:0", "branchB:0"), + ("branchB:0", "branchB:1"), + ("branchA:1", "END"), + ("branchB:1", "END"), + } + assert all(e.delay_after_predecessor_us is None for e in graph.edges) + assert graph.nodes["root:0"].inputs == [] + assert graph.nodes["branchA:0"].inputs == [ + ChannelRequirement(channel="root:0_out", count=1) + ] + assert graph.nodes["branchA:1"].inputs == [ + ChannelRequirement(channel="root:0_out", count=1), + ChannelRequirement(channel="branchA:0_out", count=1), + ] + + def test_sequential_delay_ms_lowered_to_edge_us(self, tmp_path): + path = _write_dag( + tmp_path, + [ + { + "session_id": "s", + "turns": [ + {"messages": [{"role": "user", "content": "a"}]}, + {"messages": [{"role": "user", "content": "b"}], "delay": 5.0}, + ], + } + ], + ) + graph = _lower(path).graphs["s"] + edges = {(e.source, e.target): e for e in graph.edges} + assert edges[("s:0", "s:1")].delay_after_predecessor_us == 5000.0 + + +# (b) assembly: verbatim segments, slot placement, prefix parent-chaining ----- + + +class TestAssembly: + def test_fork_child_assembly_slots_and_verbatim_segments(self): + parsed = _lower(FIXTURES_DIR / "full.dag.jsonl") + graph = parsed.graphs["root"] + node = graph.nodes["branch-a:0"] + trie = _trie(node) + sids = trie["prompt_segment_ids"] + assert trie["assembly"] == [ + {"seg": sids[0]}, + {"seg": sids[1]}, + {"s": {"src": "root:0"}}, + {"seg": sids[2]}, + {"seg": sids[3]}, + ] + assert parsed.segment_pool.materialize(sids) == [ + {"role": "system", "content": "root system prompt"}, + {"role": "user", "content": "root user prompt"}, + {"role": "user", "content": "branch-a turn-0 user message A"}, + {"role": "user", "content": "branch-a turn-0 user message B"}, + ] + assert ( + parsed.segment_pool.get(sids[0]).wire_json + == orjson.dumps( + {"role": "system", "content": "root system prompt"} + ).decode() + ), "raw segment must intern the authored message verbatim" + + def test_parent_chain_threads_through_the_slot_token(self): + parsed = _lower(FIXTURES_DIR / "full.dag.jsonl") + sids = _trie(parsed.graphs["root"].nodes["branch-a:0"])["prompt_segment_ids"] + pool = parsed.segment_pool + assert pool.get(sids[0]).parent_id is None + assert pool.get(sids[1]).parent_id == sids[0] + assert pool.get(sids[2]).parent_id == sids[1] + assert pool.get(sids[3]).parent_id == sids[2] + + def test_shared_root_prefix_dedups_across_sibling_children(self): + parsed = _lower(FIXTURES_DIR / "full.dag.jsonl") + graph = parsed.graphs["root"] + sids_a = _trie(graph.nodes["branch-a:0"])["prompt_segment_ids"] + sids_b = _trie(graph.nodes["branch-b:0"])["prompt_segment_ids"] + assert sids_a[:2] == sids_b[:2] + assert sids_a[2:] != sids_b[2:] + + def test_lineage_free_node_has_no_assembly_key(self): + parsed = _lower(FIXTURES_DIR / "full.dag.jsonl") + trie = _trie(parsed.graphs["root"].nodes["root:0"]) + assert "assembly" not in trie + assert len(trie["prompt_segment_ids"]) == 2 + + def test_spawn_child_starts_fresh_no_assembly_no_lineage_inputs(self): + parsed = _lower(FIXTURES_DIR / "spawn_minimal.dag.jsonl") + graph = parsed.graphs["root"] + child = graph.nodes["spawned-child:0"] + assert "assembly" not in _trie(child) + assert child.inputs == [] + assert parsed.segment_pool.materialize(read_prompt_segment_ids(child)) == [ + {"role": "system", "content": "spawn-sys"}, + {"role": "user", "content": "spawn-u"}, + ] + + +# (c) capture stamps ----------------------------------------------------------- + + +class TestCapture: + def test_capture_marks_lineage_producers_only(self): + graph = _lower(FIXTURES_DIR / "small.dag.jsonl").graphs["root"] + assert _trie(graph.nodes["root:0"])["capture"] is True + assert _trie(graph.nodes["branchA:0"])["capture"] is True + assert _trie(graph.nodes["branchB:0"])["capture"] is True + assert "capture" not in _trie(graph.nodes["branchA:1"]) + assert "capture" not in _trie(graph.nodes["branchB:1"]) + + def test_join_gating_leaf_is_not_captured(self, tmp_path): + path = _write_dag( + tmp_path, + [ + { + "session_id": "root", + "turns": [ + { + "messages": [{"role": "user", "content": "t0"}], + "spawns": [{"children": ["kid"], "join_at": 1}], + }, + {"messages": [{"role": "user", "content": "t1"}]}, + ], + }, + { + "session_id": "kid", + "turns": [{"messages": [{"role": "user", "content": "k"}]}], + }, + ], + ) + graph = _lower(path).graphs["root"] + # The join gate needs kid:0's channel COMMIT, never its response text. + assert "capture" not in _trie(graph.nodes["kid:0"]) + assert _trie(graph.nodes["root:0"])["capture"] is True + + +# (d) dispatch overrides + native body fields ------------------------------------ + + +class TestDispatchOverrides: + def test_native_fields_and_overrides_all_fields(self, tmp_path): + tools = [{"type": "function", "function": {"name": "lookup"}}] + path = _write_dag( + tmp_path, + [ + { + "session_id": "solo", + "turns": [ + { + "model": "authored-model", + "messages": [{"role": "user", "content": "hi"}], + "tools": tools, + "max_tokens": 7, + "extra": {"temperature": 0.5, "seed": 3}, + } + ], + } + ], + ) + node = _lower(path, default_model="fallback").graphs["solo"].nodes["solo:0"] + # Model / stream / token cap / tools ride the NATIVE fields (Turn + # naming); extra_body carries only the merged vendor keys. + assert node.model == "authored-model" + assert node.streaming is True + assert node.max_tokens == 7 + assert node.raw_tools == tools + assert node.extra_body == {"temperature": 0.5, "seed": 3} + + def test_overrides_default_model_and_omitted_optionals(self, tmp_path): + path = _write_dag( + tmp_path, + [ + { + "session_id": "solo", + "turns": [{"messages": [{"role": "user", "content": "hi"}]}], + } + ], + ) + parsed = _lower(path, default_model="fallback-model", run_streaming=False) + node = parsed.graphs["solo"].nodes["solo:0"] + assert node.model == "fallback-model" + assert node.streaming is False + assert node.max_tokens is None + assert node.extra_body == {} + + def test_endpoint_extra_lands_before_turn_extra(self, tmp_path): + # Legacy precedence: ``payload.update(endpoint.extra)`` runs BEFORE + # ``payload.update(turn extra)``, so the turn value wins on overlap. + path = _write_dag( + tmp_path, + [ + { + "session_id": "solo", + "turns": [ + { + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 7, + "extra": {"top_p": 0.9, "seed": 3}, + } + ], + } + ], + ) + parsed = _lower( + path, + default_model="fallback", + endpoint_extra=[("min_p", 0.05), ("vendor_tag", "run")], + ) + node = parsed.graphs["solo"].nodes["solo:0"] + assert node.max_tokens == 7 + overrides = node.extra_body + assert list(overrides) == [ + "min_p", + "vendor_tag", + "top_p", + "seed", + ] + assert overrides["min_p"] == 0.05 + assert overrides["vendor_tag"] == "run" + + def test_endpoint_extra_overlap_turn_extra_wins_at_endpoint_position( + self, tmp_path + ): + # OVERLAP: the same key in endpoint_extra AND the turn's ``extra`` keeps + # the endpoint_extra POSITION (dict.update first insertion) but takes + # the TURN value (later update wins) -- exactly the legacy merge. + path = _write_dag( + tmp_path, + [ + { + "session_id": "solo", + "turns": [ + { + "messages": [{"role": "user", "content": "hi"}], + "extra": {"temperature": 0.5, "top_p": 0.9}, + } + ], + } + ], + ) + parsed = _lower( + path, + default_model="fallback", + endpoint_extra=[("min_p", 0.05), ("temperature", 0.9)], + ) + overrides = parsed.graphs["solo"].nodes["solo:0"].extra_body + assert list(overrides) == [ + "min_p", + "temperature", + "top_p", + ], "overlap key must keep its endpoint_extra insertion position" + assert overrides["temperature"] == 0.5, "turn extra value must win" + assert overrides["min_p"] == 0.05 + assert overrides["top_p"] == 0.9 + + @pytest.mark.parametrize( + "endpoint_extra", + [ + param(None, id="no-extras"), + param([("min_p", 0.05)], id="with-extras"), + ], + ) # fmt: skip + def test_every_node_stamped_endpoint_extra_applied(self, endpoint_extra): + # Parse-time folding is authoritative even when the run has NO extras: + # every dag node carries the dispatch stamp so the worker never + # re-merges ``endpoint.extra`` over the adapter-owned overrides. + parsed = _lower(FIXTURES_DIR / "small.dag.jsonl", endpoint_extra=endpoint_extra) + for graph in parsed.graphs.values(): + for node in graph.nodes.values(): + assert node.metadata["dispatch"]["endpoint_extra_applied"] is True + + +# (e) tools inheritance ---------------------------------------------------------- + + +class TestToolsInheritance: + def test_nearest_lineage_tools_win(self, tmp_path): + t0 = [{"type": "function", "function": {"name": "t0"}}] + t1 = [{"type": "function", "function": {"name": "t1"}}] + path = _write_dag( + tmp_path, + [ + { + "session_id": "s", + "turns": [ + {"messages": [{"role": "user", "content": "a"}], "tools": t0}, + {"messages": [{"role": "user", "content": "b"}], "tools": t1}, + {"messages": [{"role": "user", "content": "c"}]}, + ], + } + ], + ) + graph = _lower(path).graphs["s"] + assert graph.nodes["s:0"].raw_tools == t0 + assert graph.nodes["s:1"].raw_tools == t1 + assert graph.nodes["s:2"].raw_tools == t1 + + def test_fork_child_inherits_spawn_child_does_not(self, tmp_path): + tools = [{"type": "function", "function": {"name": "root-tool"}}] + path = _write_dag( + tmp_path, + [ + { + "session_id": "root", + "turns": [ + { + "messages": [{"role": "user", "content": "r0"}], + "tools": tools, + "forks": ["forked"], + "spawns": ["spawned"], + } + ], + }, + { + "session_id": "forked", + "turns": [{"messages": [{"role": "user", "content": "f0"}]}], + }, + { + "session_id": "spawned", + "turns": [{"messages": [{"role": "user", "content": "s0"}]}], + }, + ], + ) + graph = _lower(path).graphs["root"] + assert graph.nodes["forked:0"].raw_tools == tools + assert graph.nodes["spawned:0"].raw_tools is None + + +# (f) structural soundness via the real validator -------------------------------- + + +class TestValidate: + @pytest.mark.parametrize( + "fixture_name", + [param(name, id=name) for name in ALL_DAG_FIXTURES], + ) # fmt: skip + def test_validate_no_issues_for_every_fixture(self, fixture_name): + parsed = _lower(FIXTURES_DIR / fixture_name) + assert validate(parsed) == [] + + def test_validate_no_issues_prespawn_and_join(self, tmp_path): + path = _write_dag( + tmp_path, + [ + { + "session_id": "root", + "pre_session_spawns": ["helper"], + "turns": [ + { + "messages": [{"role": "user", "content": "r0"}], + "spawns": [{"children": ["kidA", "kidB"], "join_at": 1}], + }, + {"messages": [{"role": "user", "content": "r1"}]}, + ], + }, + { + "session_id": "helper", + "turns": [{"messages": [{"role": "user", "content": "h0"}]}], + }, + { + "session_id": "kidA", + "turns": [{"messages": [{"role": "user", "content": "a"}]}], + }, + { + "session_id": "kidB", + "turns": [{"messages": [{"role": "user", "content": "b"}]}], + }, + ], + ) + assert validate(_lower(path)) == [] + + +# (g) spawn-join gating ---------------------------------------------------------- + + +class TestSpawnJoin: + def test_join_turn_inputs_gate_both_leaves_and_sequential_edge_exists( + self, tmp_path + ): + path = _write_dag( + tmp_path, + [ + { + "session_id": "root", + "turns": [ + { + "messages": [{"role": "user", "content": "t0"}], + "spawns": [{"children": ["kidA", "kidB"], "join_at": 1}], + }, + {"messages": [{"role": "user", "content": "t1"}], "delay": 5.0}, + ], + }, + { + "session_id": "kidA", + "turns": [{"messages": [{"role": "user", "content": "a"}]}], + }, + { + "session_id": "kidB", + "turns": [{"messages": [{"role": "user", "content": "b"}]}], + }, + ], + ) + parsed = _lower(path) + graph = parsed.graphs["root"] + assert graph.nodes["root:1"].inputs == [ + ChannelRequirement(channel="kidA:0_out", count=1), + ChannelRequirement(channel="kidB:0_out", count=1), + ChannelRequirement(channel="root:0_out", count=1), + ] + edges = {(e.source, e.target): e for e in graph.edges} + assert edges[("root:0", "root:1")].delay_after_predecessor_us == 5000.0 + assert edges[("root:0", "kidA:0")].delay_after_predecessor_us is None + assert edges[("root:0", "kidB:0")].delay_after_predecessor_us is None + assert validate(parsed) == [] + + def test_cross_turn_spawn_fan_in_gates_both_leaves_on_join_turn(self, tmp_path): + # Two SPAWN branches on two DIFFERENT turns (turn 0 spawns "a", turn 1 + # spawns "b"), both joining on turn 2. The join turn's AND-fan-in must + # carry both post-spawn leaf channels plus its own lineage producers, + # and the sequential edge from the previous turn must still exist. + path = _write_dag( + tmp_path, + [ + { + "session_id": "root", + "turns": [ + { + "messages": [{"role": "user", "content": "t0"}], + "spawns": [{"children": ["a"], "join_at": 2}], + }, + { + "messages": [{"role": "user", "content": "t1"}], + "spawns": [{"children": ["b"], "join_at": 2}], + "delay": 3.0, + }, + { + "messages": [{"role": "user", "content": "t2"}], + "delay": 5.0, + }, + ], + }, + { + "session_id": "a", + "turns": [{"messages": [{"role": "user", "content": "a0"}]}], + }, + { + "session_id": "b", + "turns": [{"messages": [{"role": "user", "content": "b0"}]}], + }, + ], + ) + parsed = _lower(path) + graph = parsed.graphs["root"] + # Both spawn leaves gate the join turn, followed by its sequential + # lineage producers (root:0, root:1). + assert graph.nodes["root:2"].inputs == [ + ChannelRequirement(channel="a:0_out", count=1), + ChannelRequirement(channel="b:0_out", count=1), + ChannelRequirement(channel="root:0_out", count=1), + ChannelRequirement(channel="root:1_out", count=1), + ] + edges = {(e.source, e.target): e for e in graph.edges} + # The sequential edge into the join turn carries the authored delay; the + # two spawn-dispatch edges fire immediately (delay-free). + assert edges[("root:1", "root:2")].delay_after_predecessor_us == 5000.0 + assert edges[("root:0", "a:0")].delay_after_predecessor_us is None + assert edges[("root:1", "b:0")].delay_after_predecessor_us is None + assert validate(parsed) == [] + + def test_pre_session_spawn_child_gets_start_edge(self, tmp_path): + path = _write_dag( + tmp_path, + [ + { + "session_id": "root", + "pre_session_spawns": ["helper"], + "turns": [{"messages": [{"role": "user", "content": "r0"}]}], + }, + { + "session_id": "helper", + "turns": [{"messages": [{"role": "user", "content": "h0"}]}], + }, + ], + ) + parsed = _lower(path) + graph = parsed.graphs["root"] + pairs = {(e.source, e.target) for e in graph.edges} + assert pairs == { + ("START", "root:0"), + ("START", "helper:0"), + ("root:0", "END"), + ("helper:0", "END"), + } + assert validate(parsed) == [] + + +# (g2) repeated SPAWN instances: distinct #n nodes, shared segment dedup --------- + + +class TestRepeatedSpawnInstances: + def test_same_template_spawned_from_two_turns_gets_distinct_instance_nodes( + self, tmp_path + ): + """One SPAWN template fired from two different turns => ``kid`` + ``kid#2``. + + The first instance is bare, the second carries the ``#2`` suffix (tree + encounter order). Both are fresh-context spawn children with identical + authored bytes, so their prompt segments dedup to the SAME pool ids -- the + deferred ``#n``-channel proof for the multiset parity comparator. + """ + path = _write_dag( + tmp_path, + [ + { + "session_id": "root", + "turns": [ + { + "model": "Qwen3-0.6B", + "messages": [{"role": "user", "content": "plan"}], + "max_tokens": 32, + "spawns": ["kid"], + }, + { + "model": "Qwen3-0.6B", + "messages": [{"role": "user", "content": "revise"}], + "max_tokens": 32, + "spawns": ["kid"], + }, + ], + }, + { + "session_id": "kid", + "turns": [ + { + "model": "Qwen3-0.6B", + "messages": [{"role": "user", "content": "work"}], + "max_tokens": 16, + } + ], + }, + ], + ) + parsed = _lower(path) + graph = parsed.graphs["root"] + + # Both spawn instances exist as distinct nodes. + assert "kid:0" in graph.nodes + assert "kid#2:0" in graph.nodes + assert set(graph.nodes) == {"root:0", "root:1", "kid:0", "kid#2:0"} + + # Each instance owns a distinct per-node output channel keyed by node id. + assert graph.nodes["kid:0"].output == "kid:0_out" + assert graph.nodes["kid#2:0"].output == "kid#2:0_out" + assert "kid:0_out" in graph.state + assert "kid#2:0_out" in graph.state + + # Identical authored bytes + fresh context => the content-addressed pool + # hands both instances the SAME prompt segment ids (shared dedup). + kid_sids = read_prompt_segment_ids(graph.nodes["kid:0"]) + kid2_sids = read_prompt_segment_ids(graph.nodes["kid#2:0"]) + assert kid_sids == kid2_sids + assert parsed.segment_pool.materialize(kid_sids) == [ + {"role": "user", "content": "work"} + ] + + # Structural soundness through the real validator (node-level asserts + # alone let structural bugs slip -- repo gotcha). + assert validate(parsed) == [] + + +# (g3) dag identity metadata stamp ---------------------------------------------- + + +class TestDagIdentityMetadata: + """Lowering stamps ``metadata["dag"]`` with the tree's instance identity + (``agent_depth`` / ``parent_node``) on every node, coexisting with the + ``dispatch`` and ``trie`` metadata stamps.""" + + def test_fork_fixture_stamps_root_and_child_identity(self): + graph = _lower(FIXTURES_DIR / "small.dag.jsonl").graphs["root"] + assert graph.nodes["root:0"].metadata["dag"] == { + "agent_depth": 0, + "parent_node": None, + } + for nid in ("branchA:0", "branchA:1", "branchB:0", "branchB:1"): + assert graph.nodes[nid].metadata["dag"] == { + "agent_depth": 1, + "parent_node": "root:0", + }, nid + + def test_dag_stamp_coexists_with_dispatch_and_trie_stamps(self): + graph = _lower(FIXTURES_DIR / "small.dag.jsonl").graphs["root"] + for node in graph.nodes.values(): + assert node.metadata["dispatch"] == {"endpoint_extra_applied": True} + assert "prompt_segment_ids" in node.metadata["trie"] + assert "dag" in node.metadata + + def test_prespawn_child_stamped_depth_one_no_parent(self, tmp_path): + path = _write_dag( + tmp_path, + [ + { + "session_id": "root", + "pre_session_spawns": ["helper"], + "turns": [{"messages": [{"role": "user", "content": "r0"}]}], + }, + { + "session_id": "helper", + "turns": [{"messages": [{"role": "user", "content": "h0"}]}], + }, + ], + ) + graph = _lower(path).graphs["root"] + assert graph.nodes["helper:0"].metadata["dag"] == { + "agent_depth": 1, + "parent_node": None, + } + assert graph.nodes["root:0"].metadata["dag"] == { + "agent_depth": 0, + "parent_node": None, + } + + def test_grandchild_stamped_depth_two(self, tmp_path): + path = _write_dag( + tmp_path, + [ + { + "session_id": "root", + "turns": [ + { + "messages": [{"role": "user", "content": "r0"}], + "forks": ["A"], + } + ], + }, + { + "session_id": "A", + "turns": [ + { + "messages": [{"role": "user", "content": "a0"}], + "spawns": ["kid"], + }, + {"messages": [{"role": "user", "content": "a1"}]}, + ], + }, + { + "session_id": "kid", + "turns": [{"messages": [{"role": "user", "content": "k0"}]}], + }, + ], + ) + graph = _lower(path).graphs["root"] + assert graph.nodes["kid:0"].metadata["dag"] == { + "agent_depth": 2, + "parent_node": "A:0", + } + + +# (h) determinism + multi-root --------------------------------------------------- + + +class TestDeterminismAndMultiRoot: + def test_two_fresh_parses_lower_identically(self): + path = FIXTURES_DIR / "full.dag.jsonl" + first = _lower(path) + second = _lower(path) + ga, gb = first.graphs["root"], second.graphs["root"] + assert list(ga.nodes) == list(gb.nodes) + for nid in ga.nodes: + assert _trie(ga.nodes[nid]) == _trie(gb.nodes[nid]) + assert ga.nodes[nid].extra_body == gb.nodes[nid].extra_body + assert ga.edges == gb.edges + assert sorted(first.segment_pool.by_id) == sorted(second.segment_pool.by_id) + + def test_multi_root_per_tree_graphs_share_one_pool(self): + parsed = _lower(FIXTURES_DIR / "multi_root_single_turn.dag.jsonl") + assert list(parsed.graphs) == ["r1", "r2"] + assert [(t.id, t.graph_ref) for t in parsed.traces] == [ + ("r1", "r1"), + ("r2", "r2"), + ] + assert parsed.graph is parsed.graphs["r1"] + for graph in parsed.graphs.values(): + for node in graph.nodes.values(): + for sid in read_prompt_segment_ids(node): + parsed.segment_pool.get(sid) + assert validate(parsed) == [] + + def test_empty_trees_raise_loc_prefixed(self): + with pytest.raises(NotImplementedError, match=r"^dag_jsonl workload: "): + lower_dag_trees([], default_model=None, run_streaming=True) diff --git a/tests/unit/dataset/graph/adapters/dag_jsonl/test_trace.py b/tests/unit/dataset/graph/adapters/dag_jsonl/test_trace.py new file mode 100644 index 0000000000..efba5d1339 --- /dev/null +++ b/tests/unit/dataset/graph/adapters/dag_jsonl/test_trace.py @@ -0,0 +1,391 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Adapter registration, detection, and parse-seam tests for ``dag_jsonl``. + +Covers the adapter selection surface on top of the (already tested) tree +expansion + lowering core: + +* ``DagJsonlGraphAdapter.can_load`` -- a strict, bounded, never-raising sniff + that claims real dag files and rejects dynamo / native / mooncake / empty + inputs (kept mutually exclusive with dynamo's sniff). +* plugin-registry resolution of the new ``graph_adapter.dag_jsonl`` entry. +* autodetect exclusion (a dag file is NEVER auto-claimed) paired with explicit + ``--graph-format dag_jsonl`` forcing, driven through a REAL ``BenchmarkRun`` + (MagicMock auto-creates attribute paths and would hide config drift). +* determinism: repeated lowering of the same dag file produces byte-identical + build catalogs (the parse is a pure function of file + run config). +""" + +from __future__ import annotations + +from pathlib import Path + +import orjson +import pytest +from pytest import param + +from aiperf.config.flags.cli_config import CLIConfig +from aiperf.dataset.graph.adapters.dag_jsonl.trace import ( + DagJsonlGraphAdapter, + _assert_dag_zero_arrival_offsets, + from_dag_jsonl, +) +from aiperf.dataset.graph.graph_path_catalog import build_graph_path_catalog +from aiperf.dataset.graph.models import GraphRecord, LlmNode, ParsedGraph +from aiperf.dataset.graph.workload_detect import ( + _detect_graph_workload_format, + is_graph_workload_path, + parse_graph_workload, +) +from aiperf.plugin import plugins +from aiperf.plugin.enums import PluginType +from tests.unit.conftest import make_run_from_cli + +FIXTURES_DIR = Path(__file__).parents[5] / "fixtures" / "dag" + +ALL_DAG_FIXTURES = [ + "background_fork.dag.jsonl", + "bg_fork_fanout.dag.jsonl", + "bg_fork_nested.dag.jsonl", + "bg_fork_with_spawn_join.dag.jsonl", + "full.dag.jsonl", + "multi_root_single_turn.dag.jsonl", + "small.dag.jsonl", + "spawn_minimal.dag.jsonl", +] + + +def _write_jsonl(tmp_path: Path, name: str, lines: list[dict]) -> Path: + path = tmp_path / name + path.write_bytes(b"\n".join(orjson.dumps(line) for line in lines)) + return path + + +# --- (a) can_load sniff ----------------------------------------------------- + + +class TestCanLoad: + @pytest.mark.parametrize("fixture", ALL_DAG_FIXTURES) + def test_true_on_real_dag_fixtures(self, fixture: str) -> None: + assert DagJsonlGraphAdapter.can_load(FIXTURES_DIR / fixture) is True + + def test_true_on_bare_jsonl_suffix(self, tmp_path: Path) -> None: + path = _write_jsonl( + tmp_path, + "conv.jsonl", + [{"session_id": "root", "turns": [{"messages": [{"role": "user"}]}]}], + ) + assert DagJsonlGraphAdapter.can_load(path) is True + + def test_false_on_dynamo_first_line(self, tmp_path: Path) -> None: + # A genuine dynamo trace line -- claimed by DynamoTraceAdapter, never us. + path = _write_jsonl( + tmp_path, + "dynamo.jsonl", + [ + { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "agent_context": {}, + } + ], + ) + assert DagJsonlGraphAdapter.can_load(path) is False + + def test_false_on_dynamo_discriminator_even_with_dag_keys( + self, tmp_path: Path + ) -> None: + # Mutual-exclusivity guard: the dynamo discriminator wins even if the + # (adversarial) line also carries session_id/turns. + path = _write_jsonl( + tmp_path, + "hybrid.jsonl", + [ + { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "session_id": "root", + "turns": [], + } + ], + ) + assert DagJsonlGraphAdapter.can_load(path) is False + + def test_false_on_dynamo_sink_envelope(self, tmp_path: Path) -> None: + # Dynamo file sinks wrap each record in {"timestamp", "event"}; unwrap + # before matching the discriminator. + path = _write_jsonl( + tmp_path, + "wrapped.jsonl", + [ + { + "timestamp": 12, + "event": { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + }, + } + ], + ) + assert DagJsonlGraphAdapter.can_load(path) is False + + def test_false_on_native_yaml(self, tmp_path: Path) -> None: + # Native graph workloads are .yaml/.yml -- the dag sniff only claims + # .jsonl, so a native file (even one carrying graph-shaped content) is + # rejected on suffix. + path = tmp_path / "graph.yaml" + path.write_text("version: '2.0'\nnodes: {}\nedges: []\n") + assert DagJsonlGraphAdapter.can_load(path) is False + + def test_false_on_mooncake_style_jsonl(self, tmp_path: Path) -> None: + path = _write_jsonl( + tmp_path, + "mooncake.jsonl", + [{"timestamp": 0, "input_length": 10, "output_length": 5, "hash_ids": [1]}], + ) + assert DagJsonlGraphAdapter.can_load(path) is False + + @pytest.mark.parametrize( + "name, content", + [ + param("empty.jsonl", b"", id="empty"), + param("blank_lines.jsonl", b"\n\n \n", id="blank-lines-only"), + param("garbage.jsonl", b"not json at all\n", id="invalid-json"), + param( + "missing_turns.jsonl", + b'{"session_id": "root"}\n', + id="missing-turns-key", + ), + param( + "missing_sid.jsonl", + b'{"turns": [{"messages": []}]}\n', + id="missing-session-id-key", + ), + ], + ) # fmt: skip + def test_false_on_malformed_inputs( + self, tmp_path: Path, name: str, content: bytes + ) -> None: + path = tmp_path / name + path.write_bytes(content) + assert DagJsonlGraphAdapter.can_load(path) is False + + def test_false_on_missing_file(self, tmp_path: Path) -> None: + assert DagJsonlGraphAdapter.can_load(tmp_path / "nope.jsonl") is False + + def test_false_on_wrong_suffix(self, tmp_path: Path) -> None: + path = _write_jsonl( + tmp_path, + "conv.json", + [{"session_id": "root", "turns": [{"messages": []}]}], + ) + assert DagJsonlGraphAdapter.can_load(path) is False + + +# --- (b) plugin-registry resolution ---------------------------------------- + + +class TestRegistry: + def test_registry_resolves_adapter_class(self) -> None: + cls = plugins.get_class(PluginType.GRAPH_ADAPTER, "dag_jsonl") + assert cls is DagJsonlGraphAdapter + + def test_adapter_parse_produces_parsed_graph(self) -> None: + parsed = DagJsonlGraphAdapter.parse(FIXTURES_DIR / "small.dag.jsonl") + assert isinstance(parsed, ParsedGraph) + assert set(parsed.graphs) == {"root"} + + +# --- (c) autodetect exclusion + explicit --graph-format forcing ------------- + + +def _dag_run(fixture: str, **cli_overrides): + return _dag_run_from_path(FIXTURES_DIR / fixture, **cli_overrides) + + +def _dag_run_from_path(path: Path, **cli_overrides): + cfg = CLIConfig( + model_names=["test-model"], + input_file=str(path), + tokenizer_name="builtin", + **cli_overrides, + ) + return make_run_from_cli(cfg) + + +class TestSelectionSeam: + @pytest.mark.parametrize("fixture", ALL_DAG_FIXTURES) + def test_autodetect_never_claims_dag_file(self, fixture: str) -> None: + path = FIXTURES_DIR / fixture + # Excluded from autodetect: no adapter (dag_jsonl skipped, weka/dynamo + # reject .jsonl dag content) claims the file. + assert _detect_graph_workload_format(path) is None + assert is_graph_workload_path(path) is False + + def test_forced_graph_format_returns_parsed_graph(self) -> None: + run = _dag_run("small.dag.jsonl", graph_format="dag_jsonl") + parsed = parse_graph_workload(run, FIXTURES_DIR / "small.dag.jsonl") + assert isinstance(parsed, ParsedGraph) + assert set(parsed.graphs) == {"root"} + + def test_forced_graph_format_threads_run_streaming(self) -> None: + run = _dag_run("small.dag.jsonl", graph_format="dag_jsonl", streaming=True) + parsed = parse_graph_workload(run, FIXTURES_DIR / "small.dag.jsonl") + node = next(iter(parsed.graph.nodes.values())) + assert node.streaming is True + + def test_forced_graph_format_threads_streaming_false(self) -> None: + # Discriminating counterpart to the streaming=True case above: a + # hardcoded ``True`` would survive that test but fail here. Every node's + # body ``stream`` AND its ``LlmNode.streaming`` must track the resolved + # run flag. + run = _dag_run("small.dag.jsonl", graph_format="dag_jsonl", streaming=False) + parsed = parse_graph_workload(run, FIXTURES_DIR / "small.dag.jsonl") + nodes = [ + node + for record in (parsed.graph, *parsed.graphs.values()) + for node in record.nodes.values() + ] + assert nodes + for node in nodes: + assert node.streaming is False + + def test_default_model_stamped_when_turn_has_no_authored_model( + self, tmp_path: Path + ) -> None: + # Same-run determinism can't catch a wrong-but-consistent model knob. + # A turn with NO authored ``model`` must inherit the run's resolved + # primary model (``test-model`` here), not some hardcoded default. + path = _write_jsonl( + tmp_path, + "no_model.dag.jsonl", + [ + { + "session_id": "root", + "turns": [{"messages": [{"role": "user", "content": "hi"}]}], + } + ], + ) + run = _dag_run_from_path(path, graph_format="dag_jsonl") + parsed = parse_graph_workload(run, path) + node = next(iter(parsed.graph.nodes.values())) + assert node.model == "test-model" + + def test_forced_graph_format_threads_extra_inputs(self, tmp_path: Path) -> None: + # The dag branch threads ``endpoint.extra`` (--extra-inputs) into the + # lowering so run-level vendor keys are folded into extra_body + # at parse (endpoint extras first, turn extras last; turn ``extra`` + # wins on overlap) and every node carries the + # ``endpoint_extra_applied`` stamp for the worker's skip. + path = _write_jsonl( + tmp_path, + "extras.dag.jsonl", + [ + { + "session_id": "root", + "turns": [ + { + "messages": [{"role": "user", "content": "hi"}], + "extra": {"temperature": 0.5}, + } + ], + } + ], + ) + run = _dag_run_from_path( + path, + graph_format="dag_jsonl", + extra_inputs=["temperature:0.9", "min_p:0.05"], + ) + parsed = parse_graph_workload(run, path) + node = parsed.graph.nodes["root:0"] + assert list(node.extra_body) == [ + "temperature", + "min_p", + ] + assert node.extra_body["temperature"] == 0.5, ( + "turn extra must win over --extra-inputs on overlap" + ) + assert node.extra_body["min_p"] == 0.05 + assert node.metadata["dispatch"]["endpoint_extra_applied"] is True + + def test_inter_turn_delay_cap_clamps_sequential_edge(self, tmp_path: Path) -> None: + # ``--inter-turn-delay-cap-seconds`` (seconds) clamps the authored + # per-turn ``delay`` (milliseconds) before it becomes the sequential + # edge's ``delay_after_predecessor_us``. Cap 2s => 2_000_000us; the + # uncapped 60s delay would instead surface as 60_000_000us, so this + # discriminates a missing/ignored cap. + path = _write_jsonl( + tmp_path, + "big_delay.dag.jsonl", + [ + { + "session_id": "root", + "turns": [ + {"messages": [{"role": "user", "content": "t0"}]}, + { + "messages": [{"role": "user", "content": "t1"}], + "delay": 60000.0, + }, + ], + } + ], + ) + run = _dag_run_from_path( + path, graph_format="dag_jsonl", inter_turn_delay_cap_seconds=2.0 + ) + parsed = parse_graph_workload(run, path) + edge = next( + e + for e in parsed.graph.edges + if e.source == "root:0" and e.target == "root:1" + ) + assert edge.delay_after_predecessor_us == 2_000_000.0 + + +class TestArrivalOffsetGuard: + """The t*/dynamic-slot gate carves out graphs whose every node carries an + explicit-zero ``arrival_offset_us`` -- the shape dag lowering stamps. + ``DagJsonlGraphAdapter.parse`` wraps ``from_dag_jsonl`` in + ``_assert_dag_zero_arrival_offsets`` to hold that invariant at the one + seam production dispatches through; these tests exercise the guard + directly and through the full run-parse seam.""" + + def _one_node_graph(self, arrival_offset_us: int) -> ParsedGraph: + node = LlmNode(prompt=[], output="n_out", arrival_offset_us=arrival_offset_us) + record = GraphRecord(nodes={"n": node}) + return ParsedGraph(graph=record, graphs={"root": record}) + + def test_guard_passes_on_zero_offsets(self) -> None: + _assert_dag_zero_arrival_offsets(self._one_node_graph(0)) + + def test_guard_raises_on_nonzero_offset(self) -> None: + with pytest.raises(ValueError, match="arrival_offset_us"): + _assert_dag_zero_arrival_offsets(self._one_node_graph(5)) + + def test_parse_seam_invokes_guard(self, monkeypatch) -> None: + # Prove the seam actually calls the guard: force ``from_dag_jsonl`` to + # return a nonzero-offset graph and assert the adapter's parse rejects + # it. Through the registry seam the guard's ValueError surfaces as + # GraphParseError (a ValueError subclass; message text preserved). + bad = self._one_node_graph(7) + monkeypatch.setattr( + "aiperf.dataset.graph.adapters.dag_jsonl.trace.from_dag_jsonl", + lambda *args, **kwargs: bad, + ) + run = _dag_run("small.dag.jsonl", graph_format="dag_jsonl") + with pytest.raises(ValueError, match="arrival_offset_us"): + parse_graph_workload(run, FIXTURES_DIR / "small.dag.jsonl") + + +# --- (d) determinism -------------------------------------------------------- + + +class TestDeterminism: + def test_default_knobs_deterministic(self) -> None: + # ``from_dag_jsonl`` with protocol defaults is a pure function of file. + path = FIXTURES_DIR / "full.dag.jsonl" + first = from_dag_jsonl(str(path)) + second = from_dag_jsonl(str(path)) + assert build_graph_path_catalog(first) == build_graph_path_catalog(second) diff --git a/tests/unit/dataset/graph/adapters/dag_jsonl/test_tree.py b/tests/unit/dataset/graph/adapters/dag_jsonl/test_tree.py new file mode 100644 index 0000000000..63dd2fb817 --- /dev/null +++ b/tests/unit/dataset/graph/adapters/dag_jsonl/test_tree.py @@ -0,0 +1,569 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Topology expansion tests: legacy Conversations -> per-root DagTree. + +Every case drives the real ``dag_jsonl`` loader against a fixture file (either +a tmp_path line file or an existing ``tests/fixtures/dag/*.dag.jsonl``) — no +MagicMock for the loader or the Conversation objects. The trees are then +asserted against the lineage / predecessor / join_input contract from the +task brief. +""" + +import re +from pathlib import Path + +import orjson +import pytest +from pytest import param + +from aiperf.dataset.graph.adapters.dag_jsonl.tree import ( + DagTree, + expand_trees, + load_dag_conversations, +) + +FIXTURES_DIR = Path(__file__).parents[5] / "fixtures" / "dag" + + +def _write_dag(tmp_path: Path, lines: list[dict]) -> Path: + path = tmp_path / "dag.jsonl" + path.write_bytes(b"\n".join(orjson.dumps(line) for line in lines)) + return path + + +def _expand(path: Path) -> list[DagTree]: + return expand_trees(load_dag_conversations(path, delay_cap_seconds=None)) + + +def _one_tree(path: Path) -> DagTree: + trees = _expand(path) + assert len(trees) == 1, [t.trace_id for t in trees] + return trees[0] + + +# (a) linear 3-turn single session ------------------------------------------ + + +def test_linear_session_lineage_predecessors_and_delay_ms_to_us(tmp_path): + path = _write_dag( + tmp_path, + [ + { + "session_id": "solo", + "turns": [ + {"messages": [{"role": "user", "content": "t0"}], "delay": 100.0}, + {"messages": [{"role": "user", "content": "t1"}], "delay": 5.0}, + {"messages": [{"role": "user", "content": "t2"}], "delay": 10.0}, + ], + } + ], + ) + tree = _one_tree(path) + assert tree.trace_id == "solo" + assert list(tree.nodes) == ["solo:0", "solo:1", "solo:2"] + + n0 = tree.nodes["solo:0"] + # Root turn-0 delay is NOT honored by the legacy rate loop: no predecessor. + assert n0.predecessors == [] + assert n0.lineage == [] + + n1 = tree.nodes["solo:1"] + assert n1.predecessors == [("solo:0", 5000)] + assert n1.lineage == ["solo:0"] + assert n1.turn.delay == 5.0 + + n2 = tree.nodes["solo:2"] + assert n2.predecessors == [("solo:1", 10000)] + assert n2.lineage == ["solo:0", "solo:1"] + + assert all(not n.join_inputs for n in tree.nodes.values()) + + +# (b) dag.md minimal FORK example -------------------------------------------- + + +def test_fork_children_instantiated_with_parent_lineage(): + tree = _one_tree(FIXTURES_DIR / "small.dag.jsonl") + assert tree.trace_id == "root" + assert set(tree.nodes) == { + "root:0", + "branchA:0", + "branchA:1", + "branchB:0", + "branchB:1", + } + + assert tree.nodes["root:0"].predecessors == [] + assert tree.nodes["root:0"].lineage == [] + + a0 = tree.nodes["branchA:0"] + # FORK child turn-0: predecessor is the forking parent node, delay 0. + assert a0.predecessors == [("root:0", 0)] + # child lineage = parent node's lineage + [parent node id]. + assert a0.lineage == ["root:0"] + + b0 = tree.nodes["branchB:0"] + assert b0.predecessors == [("root:0", 0)] + assert b0.lineage == ["root:0"] + + a1 = tree.nodes["branchA:1"] + assert a1.predecessors == [("branchA:0", 0)] + assert a1.lineage == ["root:0", "branchA:0"] + + # Parent has a single turn -> terminal fork, no join gates anywhere. + assert all(not n.join_inputs for n in tree.nodes.values()) + + +# (c) background fork --------------------------------------------------------- + + +def test_background_fork_parent_chain_continues_no_join(): + tree = _one_tree(FIXTURES_DIR / "background_fork.dag.jsonl") + assert tree.trace_id == "root" + assert set(tree.nodes) == {"root:0", "root:1", "root:2", "side:0", "side:1"} + + # Parent chain continues past the fork turn, purely sequential. + assert tree.nodes["root:1"].predecessors == [("root:0", 0)] + assert tree.nodes["root:2"].predecessors == [("root:1", 0)] + + side0 = tree.nodes["side:0"] + assert side0.predecessors == [("root:0", 0)] + assert side0.lineage == ["root:0"] + + side1 = tree.nodes["side:1"] + assert side1.predecessors == [("side:0", 0)] + assert side1.lineage == ["root:0", "side:0"] + + # Background fork emits no SPAWN_JOIN gate. + assert all(not n.join_inputs for n in tree.nodes.values()) + + +# (d) SPAWN with explicit join_at -------------------------------------------- + + +def test_spawn_join_at_gates_only_join_turn(tmp_path): + path = _write_dag( + tmp_path, + [ + { + "session_id": "root", + "turns": [ + { + "messages": [{"role": "user", "content": "t0"}], + "spawns": [{"children": ["kid"], "join_at": 2}], + }, + {"messages": [{"role": "user", "content": "t1"}], "delay": 3.0}, + {"messages": [{"role": "user", "content": "t2"}], "delay": 7.0}, + ], + }, + { + "session_id": "kid", + "turns": [{"messages": [{"role": "user", "content": "c0"}]}], + }, + ], + ) + tree = _one_tree(path) + assert set(tree.nodes) == {"root:0", "root:1", "root:2", "kid:0"} + + # Intermediate turn runs concurrently with the child: NOT gated, no child pred. + assert tree.nodes["root:1"].predecessors == [("root:0", 3000)] + assert tree.nodes["root:1"].join_inputs == [] + + # Join turn gates on the child leaf; sequential predecessor is still root:1. + assert tree.nodes["root:2"].predecessors == [("root:1", 7000)] + assert tree.nodes["root:2"].join_inputs == ["kid:0"] + + kid = tree.nodes["kid:0"] + assert kid.predecessors == [("root:0", 0)] + assert kid.lineage == [] + assert tree.nodes["root:0"].join_inputs == [] + + +# (e) SPAWN template referenced by two parents in one tree ------------------- + + +def test_spawn_template_referenced_twice_instances_distinctly(tmp_path): + path = _write_dag( + tmp_path, + [ + { + "session_id": "root", + "turns": [ + { + "messages": [{"role": "user", "content": "r0"}], + "forks": ["A", "B"], + } + ], + }, + { + "session_id": "A", + "turns": [ + { + "messages": [{"role": "user", "content": "a0"}], + "spawns": ["kid"], + }, + {"messages": [{"role": "user", "content": "a1"}]}, + ], + }, + { + "session_id": "B", + "turns": [ + { + "messages": [{"role": "user", "content": "b0"}], + "spawns": ["kid"], + }, + {"messages": [{"role": "user", "content": "b1"}]}, + ], + }, + { + "session_id": "kid", + "turns": [{"messages": [{"role": "user", "content": "k0"}]}], + }, + ], + ) + tree = _one_tree(path) + # Fresh instance per SPAWN reference: first "kid", second "kid#2". + assert "kid:0" in tree.nodes + assert "kid#2:0" in tree.nodes + + assert tree.nodes["kid:0"].predecessors == [("A:0", 0)] + assert tree.nodes["kid#2:0"].predecessors == [("B:0", 0)] + assert tree.nodes["kid:0"].lineage == [] + assert tree.nodes["kid#2:0"].lineage == [] + + # Auto-join on the spawning parent's next turn resolves to the right instance. + assert tree.nodes["A:1"].join_inputs == ["kid:0"] + assert tree.nodes["B:1"].join_inputs == ["kid#2:0"] + + +# (f) pre_session_spawns ------------------------------------------------------ + + +def test_pre_session_spawn_owner_and_empty_root_predecessors(tmp_path): + path = _write_dag( + tmp_path, + [ + { + "session_id": "root", + "pre_session_spawns": ["helper"], + "turns": [{"messages": [{"role": "user", "content": "r0"}]}], + }, + { + "session_id": "helper", + "turns": [{"messages": [{"role": "user", "content": "h0"}]}], + }, + ], + ) + tree = _one_tree(path) + assert tree.trace_id == "root" + assert set(tree.nodes) == {"root:0", "helper:0"} + + helper = tree.nodes["helper:0"] + # Owner is the root: its turn-0 has no predecessors, so the copy is empty. + assert helper.predecessors == [] + assert helper.lineage == [] + + assert tree.nodes["root:0"].predecessors == [] + + +# (g) session_id collides with correlation-id framing ------------------------ + + +@pytest.mark.parametrize("bad_char", ["|", "#"]) +def test_session_id_with_framing_char_raises_loc_prefixed(tmp_path, bad_char): + sid = f"root{bad_char}x" + path = _write_dag( + tmp_path, + [ + { + "session_id": sid, + "turns": [{"messages": [{"role": "user", "content": "r0"}]}], + } + ], + ) + conversations = load_dag_conversations(path, delay_cap_seconds=None) + with pytest.raises(NotImplementedError, match=rf"conversation '{re.escape(sid)}':"): + expand_trees(conversations) + + +# (h) cross-turn SPAWN fan-in: two spawn turns join on one gated turn ---------- + + +def test_two_spawn_turns_fan_in_to_one_join_turn(tmp_path): + # root turn 0 spawns "a" (join_at 2) and turn 1 spawns "b" (join_at 2), so the + # gated turn root:2 fans in BOTH post-spawn leaves through two SPAWN_JOIN + # prerequisites resolved from distinct branches. + path = _write_dag( + tmp_path, + [ + { + "session_id": "root", + "turns": [ + { + "messages": [{"role": "user", "content": "t0"}], + "spawns": [{"children": ["a"], "join_at": 2}], + }, + { + "messages": [{"role": "user", "content": "t1"}], + "spawns": [{"children": ["b"], "join_at": 2}], + "delay": 3.0, + }, + {"messages": [{"role": "user", "content": "t2"}], "delay": 5.0}, + ], + }, + { + "session_id": "a", + "turns": [{"messages": [{"role": "user", "content": "a0"}]}], + }, + { + "session_id": "b", + "turns": [{"messages": [{"role": "user", "content": "b0"}]}], + }, + ], + ) + tree = _one_tree(path) + assert set(tree.nodes) == {"root:0", "a:0", "root:1", "b:0", "root:2"} + + # Each spawn child fires at its own spawning parent's credit return (delay 0). + assert tree.nodes["a:0"].predecessors == [("root:0", 0)] + assert tree.nodes["b:0"].predecessors == [("root:1", 0)] + + # Only the join turn gates; the intermediate spawn turn does not. + assert tree.nodes["root:0"].join_inputs == [] + assert tree.nodes["root:1"].join_inputs == [] + + # Join turn fans in both leaves (branch/prerequisite declaration order) and + # still chains sequentially off root:1 carrying its authored delay. + root2 = tree.nodes["root:2"] + assert root2.join_inputs == ["a:0", "b:0"] + assert root2.predecessors == [("root:1", 5000)] + assert root2.lineage == ["root:0", "root:1"] + + +# (i) legacy identity: agent_depth / parent_node_id --------------------------- + + +class TestDagIdentity: + """Legacy identity semantics mirrored EXACTLY from ``conversation_source``: + roots ``(0, None)``; FORK/SPAWN child instances ``(owner depth + 1, + triggering parent node)``; pre-session spawn children ``(1, None)``. + Every turn of one child instance carries the instance's identity.""" + + def test_root_turns_depth_zero_no_parent(self, tmp_path): + path = _write_dag( + tmp_path, + [ + { + "session_id": "solo", + "turns": [ + {"messages": [{"role": "user", "content": "t0"}]}, + {"messages": [{"role": "user", "content": "t1"}]}, + ], + } + ], + ) + tree = _one_tree(path) + for node in tree.nodes.values(): + assert node.agent_depth == 0 + assert node.parent_node_id is None + + def test_fork_children_all_turns_depth_one_parent_is_forking_node(self): + tree = _one_tree(FIXTURES_DIR / "small.dag.jsonl") + assert tree.nodes["root:0"].agent_depth == 0 + assert tree.nodes["root:0"].parent_node_id is None + for nid in ("branchA:0", "branchA:1", "branchB:0", "branchB:1"): + assert tree.nodes[nid].agent_depth == 1, nid + assert tree.nodes[nid].parent_node_id == "root:0", nid + + def test_spawn_child_depth_one_parent_is_spawning_node(self, tmp_path): + path = _write_dag( + tmp_path, + [ + { + "session_id": "root", + "turns": [ + { + "messages": [{"role": "user", "content": "t0"}], + "spawns": [{"children": ["kid"], "join_at": 1}], + }, + {"messages": [{"role": "user", "content": "t1"}]}, + ], + }, + { + "session_id": "kid", + "turns": [{"messages": [{"role": "user", "content": "k0"}]}], + }, + ], + ) + tree = _one_tree(path) + assert tree.nodes["kid:0"].agent_depth == 1 + assert tree.nodes["kid:0"].parent_node_id == "root:0" + assert tree.nodes["root:1"].agent_depth == 0 + assert tree.nodes["root:1"].parent_node_id is None + + def test_grandchild_depth_two_parent_is_triggering_child_node(self, tmp_path): + # root forks A; A's turn 0 spawns kid => kid is a grandchild instance. + path = _write_dag( + tmp_path, + [ + { + "session_id": "root", + "turns": [ + { + "messages": [{"role": "user", "content": "r0"}], + "forks": ["A"], + } + ], + }, + { + "session_id": "A", + "turns": [ + { + "messages": [{"role": "user", "content": "a0"}], + "spawns": ["kid"], + }, + {"messages": [{"role": "user", "content": "a1"}]}, + ], + }, + { + "session_id": "kid", + "turns": [ + {"messages": [{"role": "user", "content": "k0"}]}, + {"messages": [{"role": "user", "content": "k1"}]}, + ], + }, + ], + ) + tree = _one_tree(path) + assert tree.nodes["root:0"].agent_depth == 0 + for nid in ("A:0", "A:1"): + assert tree.nodes[nid].agent_depth == 1, nid + assert tree.nodes[nid].parent_node_id == "root:0", nid + # ALL turns of the grandchild instance share the instance identity. + for nid in ("kid:0", "kid:1"): + assert tree.nodes[nid].agent_depth == 2, nid + assert tree.nodes[nid].parent_node_id == "A:0", nid + + def test_pre_session_spawn_child_depth_one_no_parent(self, tmp_path): + path = _write_dag( + tmp_path, + [ + { + "session_id": "root", + "pre_session_spawns": ["helper"], + "turns": [{"messages": [{"role": "user", "content": "r0"}]}], + }, + { + "session_id": "helper", + "turns": [{"messages": [{"role": "user", "content": "h0"}]}], + }, + ], + ) + tree = _one_tree(path) + helper = tree.nodes["helper:0"] + assert helper.agent_depth == 1 + assert helper.parent_node_id is None + assert tree.nodes["root:0"].agent_depth == 0 + assert tree.nodes["root:0"].parent_node_id is None + + def test_repeated_spawn_instances_each_carry_their_own_spawning_parent( + self, tmp_path + ): + path = _write_dag( + tmp_path, + [ + { + "session_id": "root", + "turns": [ + { + "messages": [{"role": "user", "content": "r0"}], + "spawns": ["kid"], + }, + { + "messages": [{"role": "user", "content": "r1"}], + "spawns": ["kid"], + }, + ], + }, + { + "session_id": "kid", + "turns": [{"messages": [{"role": "user", "content": "k0"}]}], + }, + ], + ) + tree = _one_tree(path) + assert tree.nodes["kid:0"].agent_depth == 1 + assert tree.nodes["kid:0"].parent_node_id == "root:0" + assert tree.nodes["kid#2:0"].agent_depth == 1 + assert tree.nodes["kid#2:0"].parent_node_id == "root:1" + + +# (j) non-root pre_session_spawns rejected at load ---------------------------- + + +@pytest.mark.parametrize( + "owner_lines", + [ + param( + [ + { + "session_id": "root", + "turns": [ + { + "messages": [{"role": "user", "content": "r0"}], + "forks": ["owner"], + } + ], + }, + { + "session_id": "owner", + "pre_session_spawns": ["bg"], + "turns": [{"messages": [{"role": "user", "content": "o0"}]}], + }, + { + "session_id": "bg", + "turns": [{"messages": [{"role": "user", "content": "b0"}]}], + }, + ], + id="fork-child-owner", + ), + param( + [ + { + "session_id": "root", + "turns": [ + { + "messages": [{"role": "user", "content": "r0"}], + "spawns": ["owner"], + } + ], + }, + { + "session_id": "owner", + "pre_session_spawns": ["bg"], + "turns": [{"messages": [{"role": "user", "content": "o0"}]}], + }, + { + "session_id": "bg", + "turns": [{"messages": [{"role": "user", "content": "b0"}]}], + }, + ], + id="spawn-child-owner", + ), + ], +) # fmt: skip +def test_non_root_pre_session_spawns_rejected_at_load(tmp_path, owner_lines): + # A non-root conversation declaring ``pre_session_spawns`` is rejected by the + # v1 orchestrator gate ("pre-session dispatch requires a root conversation") + # during load. This is why a pre-session child's copied predecessors are + # always empty: only a root can own pre-session spawns and its turn-0 has no + # predecessors, so the generic non-root-owner branch of ``_child_entry_context`` + # is unreachable through the real loader. + path = _write_dag(tmp_path, owner_lines) + with pytest.raises( + NotImplementedError, + match=r"conversation 'owner' branch 'owner:pre': " + r"pre-session dispatch requires a root conversation", + ): + load_dag_conversations(path, delay_cap_seconds=None) diff --git a/tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/_build_fixtures.py b/tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/_build_fixtures.py new file mode 100644 index 0000000000..d657f66032 --- /dev/null +++ b/tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/_build_fixtures.py @@ -0,0 +1,400 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Reproducibly emit synthetic Dynamo nested-subagent trace fixtures. + +Run as a script to regenerate every ``*.jsonl.gz`` fixture next to this file:: + + uv run python tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/_build_fixtures.py + +Each fixture is a gzipped JSONL stream of ``dynamo.request.trace.v1`` records +modelling a different nesting topology. The schema mirrors +``src/aiperf/dataset/graph/adapters/dynamo/trace_reader.py`` +(``AgentTraceRecord``); session identity lives in +``agent_context.session_id`` / ``agent_context.parent_session_id``. + +Time conventions used in these fixtures: + - Each session's request_end events are spaced 100 ms apart starting at + the session's base offset. + - A subagent session's first request_end falls inside the parent's + K-th window, where the K-th window is the interval + ``(request_end[K-1], request_end[K])`` (1-based; window K=1 is + ``(start, request_end[1])``). + - Tool events on the parent session use timestamps that fall just before + the next ``request_end`` of that session. +""" + +from __future__ import annotations + +import gzip +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import orjson + +MODEL = "synthetic-llm" + +FIXTURES_DIR = Path(__file__).resolve().parent + + +# --- record builders ------------------------------------------------------ + + +def _ctx(*, session_id: str, parent_session_id: str | None = None) -> dict[str, Any]: + out: dict[str, Any] = { + "session_id": session_id, + } + if parent_session_id is not None: + out["parent_session_id"] = parent_session_id + return out + + +def request_end( + *, + ts: int, + session_id: str, + parent_session_id: str | None = None, + request_id: str | None = None, + input_tokens: int = 32, + output_tokens: int = 16, + cached_tokens: int = 0, + ttft_ms: float = 50.0, +) -> dict[str, Any]: + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": _ctx( + session_id=session_id, parent_session_id=parent_session_id + ), + "request": { + "request_id": request_id or f"{session_id}-r{ts}", + "model": MODEL, + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "cached_tokens": cached_tokens, + "ttft_ms": ttft_ms, + }, + } + + +def tool_start( + *, + ts: int, + session_id: str, + parent_session_id: str | None = None, + tool_call_id: str = "tc", + tool_class: str = "search", +) -> dict[str, Any]: + return { + "schema": "dynamo.request.trace.v1", + "event_type": "tool_start", + "event_time_unix_ms": ts, + "event_source": "harness", + "agent_context": _ctx( + session_id=session_id, parent_session_id=parent_session_id + ), + "tool": { + "tool_call_id": tool_call_id, + "tool_class": tool_class, + "started_at_unix_ms": ts, + "status": "running", + }, + } + + +def tool_end( + *, + ts: int, + session_id: str, + parent_session_id: str | None = None, + tool_call_id: str = "tc", + tool_class: str = "search", + duration_ms: float = 30.0, + status: str = "succeeded", +) -> dict[str, Any]: + return { + "schema": "dynamo.request.trace.v1", + "event_type": "tool_end", + "event_time_unix_ms": ts, + "event_source": "harness", + "agent_context": _ctx( + session_id=session_id, parent_session_id=parent_session_id + ), + "tool": { + "tool_call_id": tool_call_id, + "tool_class": tool_class, + "ended_at_unix_ms": ts, + "duration_ms": duration_ms, + "status": status, + }, + } + + +# --- declarative chain helpers -------------------------------------------- + + +@dataclass(frozen=True) +class SessionSpec: + """Declarative description of one session's request_end chain.""" + + session_id: str + parent_session_id: str | None + base_ts: int + n_turns: int + + +def request_ends_for(spec: SessionSpec) -> list[dict[str, Any]]: + """Emit ``n_turns`` request_end records 100 ms apart from ``base_ts``.""" + return [ + request_end( + ts=spec.base_ts + 100 * i, + session_id=spec.session_id, + parent_session_id=spec.parent_session_id, + ) + for i in range(spec.n_turns) + ] + + +def write_fixture(name: str, records: list[dict[str, Any]]) -> Path: + """Sort by event_time (Dynamo doesn't guarantee global order) and write.""" + out = FIXTURES_DIR / name + records_sorted = sorted(records, key=lambda r: r["event_time_unix_ms"]) + with gzip.open(out, "wb") as f: + for r in records_sorted: + f.write(orjson.dumps(r)) + f.write(b"\n") + return out + + +# --- fixtures ------------------------------------------------------------- + + +def build_nested_2_level() -> Path: + """A: 3 turns. B (parent=A): 2 turns; B's first request_end inside A.K=2 window. + + Timeline: + A.r1 @ 1000 (A window K=1 starts here) + A.r2 @ 1100 (A window K=2 starts here) + B.r1 @ 1150 (inside (A.r2=1100, A.r3=1200)) -> resolves to A.K=2 + B.r2 @ 1170 + A.r3 @ 1200 + """ + a = SessionSpec( + session_id="sess_A", + parent_session_id=None, + base_ts=1000, + n_turns=3, + ) + b_records = [ + request_end(ts=1150, session_id="sess_B", parent_session_id="sess_A"), + request_end(ts=1170, session_id="sess_B", parent_session_id="sess_A"), + ] + records = request_ends_for(a) + b_records + return write_fixture("nested_2_level.jsonl.gz", records) + + +def build_nested_3_level() -> Path: + """A -> B -> C; B at A.K=2, C at B.K=1. + + Timeline: + A.r1 @ 1000 + A.r2 @ 1100 (A window K=2 starts) + B.r1 @ 1150 -> A.K=2 + C.r1 @ 1160 (inside (B.r1=1150, B.r2=1180)) -> B.K=1 + B.r2 @ 1180 + A.r3 @ 1200 + """ + a_records = request_ends_for( + SessionSpec( + session_id="sess_A", + parent_session_id=None, + base_ts=1000, + n_turns=3, + ) + ) + b_records = [ + request_end(ts=1150, session_id="sess_B", parent_session_id="sess_A"), + request_end(ts=1180, session_id="sess_B", parent_session_id="sess_A"), + ] + c_records = [ + request_end(ts=1160, session_id="sess_C", parent_session_id="sess_B"), + ] + return write_fixture("nested_3_level.jsonl.gz", a_records + b_records + c_records) + + +def build_mixed_turn() -> Path: + """A.K=2 has BOTH A's own tool events AND a subagent B's first request_end. + + The parent's tool_call_id does NOT mention any child session_id, so the + splice resolution must come from the time-window heuristic, not from the + tool_call_id contains-sid heuristic. + """ + a_records = request_ends_for( + SessionSpec( + session_id="sess_A", + parent_session_id=None, + base_ts=1000, + n_turns=3, + ) + ) + # A.K=2 window = (1100, 1200) + parent_tool_events = [ + tool_start( + ts=1110, + session_id="sess_A", + tool_call_id="local_search_call_xyz", + tool_class="local_search", + ), + tool_end( + ts=1140, + session_id="sess_A", + tool_call_id="local_search_call_xyz", + tool_class="local_search", + duration_ms=30.0, + ), + ] + b_records = [ + request_end(ts=1150, session_id="sess_B", parent_session_id="sess_A"), + request_end(ts=1170, session_id="sess_B", parent_session_id="sess_A"), + ] + records = a_records + parent_tool_events + b_records + return write_fixture("mixed_turn.jsonl.gz", records) + + +def build_cycle_AB_A() -> Path: + """Malformed: A claims parent=B, B claims parent=A. + + Both sessions declare each other as parent; the adapter must reject this. + """ + records: list[dict[str, Any]] = [] + for i in range(3): + records.append( + request_end( + ts=1000 + i * 100, + session_id="sess_A", + parent_session_id="sess_B", + ) + ) + for i in range(2): + records.append( + request_end( + ts=1150 + i * 30, + session_id="sess_B", + parent_session_id="sess_A", + ) + ) + return write_fixture("cycle_AB_A.jsonl.gz", records) + + +def build_parallel_subagents() -> Path: + """A invokes B and C from the same parent turn (K=2). + + Both B's and C's first request_end fall inside A.K=2's window + (1100, 1200). + """ + a_records = request_ends_for( + SessionSpec( + session_id="sess_A", + parent_session_id=None, + base_ts=1000, + n_turns=3, + ) + ) + b_records = [ + request_end(ts=1130, session_id="sess_B", parent_session_id="sess_A"), + request_end(ts=1160, session_id="sess_B", parent_session_id="sess_A"), + ] + c_records = [ + request_end(ts=1140, session_id="sess_C", parent_session_id="sess_A"), + request_end(ts=1170, session_id="sess_C", parent_session_id="sess_A"), + ] + return write_fixture( + "parallel_subagents.jsonl.gz", a_records + b_records + c_records + ) + + +def build_tool_call_id_linkage() -> Path: + """A invokes B, with parent A's tool_start.tool_call_id naming B's session_id. + + Encodes the substring as ``"subagent::invoke"`` per the spec + convention. Two tools are emitted by A (one bracketing B); the second's + tool_call_id contains B's session_id verbatim, which is the heuristic the + adapter uses to upgrade splice-K resolution from ``time_window`` to + ``tool_call_id``. + """ + a_records = request_ends_for( + SessionSpec( + session_id="sess_A", + parent_session_id=None, + base_ts=1000, + n_turns=3, + ) + ) + parent_tool_events = [ + tool_start( + ts=1110, + session_id="sess_A", + tool_call_id="subagent:sess_B:invoke", + tool_class="subagent_call", + ), + tool_end( + ts=1190, + session_id="sess_A", + tool_call_id="subagent:sess_B:invoke", + tool_class="subagent_call", + duration_ms=80.0, + ), + ] + b_records = [ + request_end(ts=1130, session_id="sess_B", parent_session_id="sess_A"), + request_end(ts=1160, session_id="sess_B", parent_session_id="sess_A"), + ] + return write_fixture( + "tool_call_id_linkage.jsonl.gz", + a_records + parent_tool_events + b_records, + ) + + +def build_parallel_two_root() -> Path: + """Two parentless (root) sessions in one file. + + Exercises the Task-8 multi-root gate: the subgraph-mode adapter supports + only one root session per file and must raise when it sees two roots. + """ + a_records = request_ends_for( + SessionSpec( + session_id="sess_A", + parent_session_id=None, + base_ts=1000, + n_turns=2, + ) + ) + b_records = request_ends_for( + SessionSpec( + session_id="sess_B", + parent_session_id=None, + base_ts=1010, + n_turns=2, + ) + ) + return write_fixture("parallel_two_root.jsonl.gz", a_records + b_records) + + +def build_all() -> list[Path]: + return [ + build_nested_2_level(), + build_nested_3_level(), + build_mixed_turn(), + build_cycle_AB_A(), + build_parallel_subagents(), + build_tool_call_id_linkage(), + build_parallel_two_root(), + ] + + +if __name__ == "__main__": + paths = build_all() + for p in paths: + print(p) diff --git a/tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/cycle_AB_A.jsonl.gz b/tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/cycle_AB_A.jsonl.gz new file mode 100644 index 0000000000000000000000000000000000000000..6681731283d9a9e74df6e8351f3dde027ea83f72 GIT binary patch literal 273 zcmV+s0q*`EiwFo$xc5w;5 z#Xlm1v}r^w$(2?i#`xbWg*Xiq+CyoLkDA#Lvk!KHAWvunw4g>Sj4M69M+{Vw7ZS^J z1qI+4EpZwLOk+2nqg9`lHprvZ7j6Qy6=Q61f{R|_{r8j*uUVJQ5?%<}2?BxK<0**rf@e*v#rA6~Uxc+KzNb>G1|{0Ohh@S4rS+piCA Xzg2kk`e$4J$Ju@WgzI3=F$Mqt>!f_D literal 0 HcmV?d00001 diff --git a/tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/mixed_turn.jsonl.gz b/tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/mixed_turn.jsonl.gz new file mode 100644 index 0000000000000000000000000000000000000000..109b4761ecb3f0cacbe47c8c6708a554e98eccb1 GIT binary patch literal 371 zcmV-(0gV11iwFo$xCWOmjwJMwxBLDkqum_pTA*|^-}GGd^dyi}O4 zO2`0LXh~4FK&I{L0&VqC(hZ7ebt{?x)3PXveI>Z)6&_CK4mnS_a+a_qn64pU2-;bp zE2gC%@eG@6uag*i$s@+VClx-faTS^bca~}-tp>(4jBC4yY|Oo)WlN)B1}@SuR6b>( zBoAL5Aa+gkLj4fqA}=V;=5QX<1@bi7`h4{+64drc)x@| literal 0 HcmV?d00001 diff --git a/tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/nested_3_level.jsonl.gz b/tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/nested_3_level.jsonl.gz new file mode 100644 index 0000000000000000000000000000000000000000..dfcfcc85127b8d47b0edd608421676e419bd686f GIT binary patch literal 288 zcmV+*0pI=~iwFo$xToQ5QXx z?JUBh|$<%u>Hgt8u=Si8rfjfm1z%cSro+*@pf6; zJQz#u6`=vTQe@5JnMC%@Yl`lN@_}@T{-#5Q4wX_1lT#Ti-KWbBbF8>ck?F?JO^_*P zpQ83oWc8unU~ZN@Rk5smP;7^~MaG5XsITxC`9$pPw(o=BKAwvhOCJ=)(#x)aRQM6X z*9NSSN`LTa`-7|W2iNG&Ck^t; I?SKRT0E`KLMgRZ+ literal 0 HcmV?d00001 diff --git a/tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/tool_call_id_linkage.jsonl.gz b/tests/unit/dataset/graph/adapters/fixtures/dynamo_nested/tool_call_id_linkage.jsonl.gz new file mode 100644 index 0000000000000000000000000000000000000000..81ba16fe65eb5bf37520b8196f8c08391d902d01 GIT binary patch literal 382 zcmV-^0fGJ>iwFo$x4yZa; zdU%Tvs3fm6leY3PArySxN!k15lEjA%Bp3!ife}i z8xqYZ0)${3O4BkdJ>-+v7JHtA*i)V$CO)xnyUw*>Q+PmH0U159w&S?LKC&?P3doXX z(?YmNr=faY2(5Vd_7Gy#L@$(wxUb4;yS*DP;ZJ=Pf9l)xhn%xAkn;4YJ@vUgE9H>~ z`-VS>C#5r=R@UwS@i}*BGIG%vT=R6I*x{^Vk|Qb?ax&{>Clr0u$r*~N#79JfOL4KK znf-2Q77Xi^&-Fjps72bS*}uKG`QMv2zuj~E%DKmzJ4RM}c8W=58@-}Lo0 (0-based), the latest parent-session turn started +at or before a subagent's first turn, and ``None`` for chain roots / no earlier +parent turn. +""" + +from __future__ import annotations + +from aiperf.dataset.graph.adapters.dynamo.trace import _Chain, _Turn +from aiperf.dataset.graph.adapters.dynamo.trace_reader import AgentTraceRecord +from aiperf.dataset.graph.adapters.dynamo.trie_lowering import dynamo_trie_nodes + + +def _rec_obj(ts, sid, itok, otok, received=None, total=None): + req = { + "request_id": f"r{ts}", + "model": "m", + "input_tokens": itok, + "output_tokens": otok, + "cached_tokens": 0, + } + if received is not None: + req["request_received_ms"] = received + if total is not None: + req["total_time_ms"] = total + return AgentTraceRecord.model_validate( + { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": {"session_id": sid}, + "request": req, + } + ) + + +def _chain(sid, turns, parent=None): + return _Chain( + sid, + parent_session_id=parent, + turns=[_Turn(record=r) for r in turns], + ) + + +def _by_id(nodes): + return {n.node_id: n for n in nodes} + + +def test_second_turn_stamps_previous_turn_in_chain(): + """Turn k>0 (0-based) causal parent is the prior turn in the same chain.""" + chains = { + "s1": _chain( + "s1", + [ + _rec_obj(1000, "s1", 32, 8, received=1000), + _rec_obj(2000, "s1", 64, 8, received=2000), + ], + ) + } + nodes, _bs, _tags = dynamo_trie_nodes(chains) + by_id = _by_id(nodes) + assert by_id["s1:1"].causal_parent_id == "s1:0" + + +def test_child_first_turn_stamps_latest_parent_turn_at_or_before_start(): + """A subagent's first turn (start 5.0) picks the latest parent-session turn + whose start (0.0, 4.0) is <= its own start -- the 4.0 turn.""" + chains = { + "root": _chain( + "root", + [ + _rec_obj(0, "root", 32, 8, received=0), + _rec_obj(4000, "root", 64, 8, received=4000), + ], + ), + "child": _chain( + "child", + [_rec_obj(5000, "child", 16, 4, received=5000)], + parent="root", + ), + } + nodes, _bs, _tags = dynamo_trie_nodes(chains) + by_id = _by_id(nodes) + assert by_id["child:0"].causal_parent_id == "root:1" + + +def test_root_first_turn_has_no_causal_parent(): + """Turn k==0 of a root chain (no parent_session_id) is None.""" + chains = {"s1": _chain("s1", [_rec_obj(1000, "s1", 32, 8, received=1000)])} + nodes, _bs, _tags = dynamo_trie_nodes(chains) + assert _by_id(nodes)["s1:0"].causal_parent_id is None + + +def test_child_first_turn_before_any_parent_turn_has_no_causal_parent(): + """A subagent's first turn whose parent session has NO turn started at or + before it resolves to None.""" + chains = { + "root": _chain( + "root", + [_rec_obj(5000, "root", 32, 8, received=5000)], + ), + "child": _chain( + "child", + [_rec_obj(1000, "child", 16, 4, received=1000)], + parent="root", + ), + } + nodes, _bs, _tags = dynamo_trie_nodes(chains) + assert _by_id(nodes)["child:0"].causal_parent_id is None diff --git a/tests/unit/dataset/graph/adapters/test_dynamo_corpus_scale_memory.py b/tests/unit/dataset/graph/adapters/test_dynamo_corpus_scale_memory.py new file mode 100644 index 0000000000..4e31e6c457 --- /dev/null +++ b/tests/unit/dataset/graph/adapters/test_dynamo_corpus_scale_memory.py @@ -0,0 +1,928 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Corpus-scale build-plane memory measurement for the dynamo trie parse. + +Sizes the four build tiers on a synthetic corpus-scale dynamo capture +(``tests/harness/dynamo_synth_corpus.py``) and linearly extrapolates each to the +confirmed real-capture scale (~1M nodes), so a per-node memory regression in +any tier is caught at test scale (MBs) before it becomes tens of GB at corpus +scale. The build's memory posture rests on several deliberate optimizations +(offset-cached decode, flat resolution automaton, content-buffer spill, direct +store route, read-time hash-id interning, sid-string pool interning); each is +gated below by a calibrated ratio so a silent regression cannot land. + +All measurement tests are ``@pytest.mark.slow`` (deselected by default; +``pyproject.toml`` ``addopts`` carries ``not slow``), so their cost is opt-in: + + uv run pytest -m slow tests/unit/dataset/graph/adapters/test_dynamo_corpus_scale_memory.py + +The tiers, and where each is allocated in the parse: + +| Tier | Allocation site (filename) | What it is | +|------|----------------------------|------------| +| hash-id ints + lists | ``dynamo/trace_reader.py`` (u64 ints via ``orjson.loads`` + pydantic), ``dynamo/trie_lowering.py`` (``TrieRequest.hash_ids`` list containers + ``TrieNode``s), ``dynamo/trace.py`` | every turn re-lists its full prefix, so ~``G*(T+1)/2`` slots per node; ``_collect_records`` interns the recorded ints read-time, so each slot is one list pointer into a shared per-UNIQUE-value ``int`` (above the small-int cache), not a distinct object per slot | +| decode cache | ``adapters/shared/content.py`` (``_decode_block_tokens_offset_cached``: one ``int`` offset per unique hash id) | one cached offset per unique hash id, held for the whole build; blocks re-sliced from the corpus on demand | +| resolution automaton | ``segment_ir/trie_content.py`` (the flat int-state automaton built by ``_insert_flat``) | one automaton state per unique block position; a TRANSIENT freed before emission, so it is measured in its own phase isolate, not the full-parse snapshot | +| content pool | ``segment_ir/pool.py`` (``Segment`` + ``SegmentPool._by_id``) | one content-addressed entry per unique message segment | + +MEASURED TIER TABLE (this machine, CPython 3.12.10, pydantic 2.13.4, at the +committed default scale S=150 T=40 G=25 -> N=6000, H=3,075,000 hash slots, +U=150,000 unique blocks; tracemalloc ``nframes=1``; corpus pre-warmed OUTSIDE +the traced region so the fixed ~600k-token pool is excluded and the numbers +are per-node marginal; full parse runs in ~25 s): + +| Tier | bytes @ N=6000 | bytes/node | linear @ 1M | +|------------------------------------|---------------:|-----------:|------------:| +| hash-id ints + lists (full parse) | 42.7 MB | 7,124 B | 7.1 GB | +| decode cache (full-parse window) | 5.2 MB | 874 B | 0.9 GB | +| decode cache (isolate, all U) | 26.3 MB | 4,390 B | 4.4 GB | +| resolution automaton (isolate) | 20.4 MB | 3,400 B | 3.4 GB | +| trie build artifacts (full parse) | 1.8 MB | 294 B | 0.3 GB | +| content pool (pool.py, eager) | 3.0 MB | 501 B | 0.5 GB | +| sid addressing tier (segment_id) | 1.3 MB | 217 B | 0.2 GB | +| resident Segment graph (eager) | ~1.7 MB | ~1,283 B | ~0.3 GB | +| content pool (pool.py, direct) | 1.3 MB | 217 B | 0.2 GB | +| _sids list (store_backed_pool.py) | 0.2 MB | 26 B | 0.0 GB | +| MEASURED TOTAL PEAK (full parse) | 123.1 MB | 20,523 B | 20.5 GB | +| MEASURED TOTAL PEAK (direct route) | 113.0 MB | 18,833 B | 18.8 GB | + +THE TWO INTERNING LAYERS the ratio gates below protect, and why each matters: + +* Read-time hash-id interning (``_collect_records``, plus the mandatory + ``list()`` companion in ``dynamo_trie_nodes``): every turn re-lists its full + prefix, so without interning the recorded hash-id-int tier holds one + distinct ``int`` per re-listed slot (~149 MB at test scale, ~24.8 GB @1M on + this shape); interning shares one canonical ``int`` per UNIQUE value + (~42.7 MB / ~7.1 GB), leaving only the per-slot list pointers and per-node + list headers scaling with slot count. +* Sid-string pool interning (eager ``InterningSegmentPool`` / direct + ``StoreBackedSegmentPool._sids``): every node's ``prompt_segment_ids`` + re-lists its whole message chain, so without interning the ``segment_id`` + hexdigest ADDRESSING tier holds ~246k fresh strings for 17,850 unique + segments (~18 MB / ~3.0 GB @1M); interning returns the FIRST-BORN canonical + object (~1.3 MB / ~0.2 GB). Interning changes str identity only -- golden + digest and three-way parity are byte-unchanged. + +DIRECT-ROUTE CONTENT POOL (``test_direct_route_content_tier_collapses``): +threading the live unified store as ``direct_store`` makes the resident +``Segment`` graph in ``SegmentPool._by_id`` collapse 17,850 -> 0 segments (the +pool holds NOTHING; every segment is interned straight into the store): the +write-through removes the ~1.7 MB (~0.3 GB @1M) +resident Segment graph plus the ``content`` strings it pinned, while the +canonical ``segment_id`` sid strings that EVERY node's ``prompt_segment_ids`` +retains on BOTH routes (per-node ADDRESSING, not the content pool) stay at the +~1.3 MB interned floor. The direct route's ``_sids`` canonical list +adds ~0.2 MB in ``store_backed_pool.py`` (24 MB @1M) -- the strings themselves are +already retained as the store's ``_ids`` keys. + +(Exact bytes vary by +interpreter build -- the load-bearing conclusion is the tier ORDERING and +extrapolated MAGNITUDE, not the exact bytes; every assertion gates a +calibrated ratio, never absolute bytes.) + +CAVEAT ON THE @1M COLUMN: the linear extrapolation holds the synthetic chain +shape fixed (turns_per_session=40, new_blocks_per_turn=25 -> H/N ~= 512 +re-listed hash slots per node). Real captures with deeper chains grow the +hash-id tier SUPERLINEARLY in chain depth, so the absolute @1M magnitudes are +chain-shape-conditional and pending real-capture re-validation +(``AIPERF_TEST_DYNAMO_SCALE_NODES`` lane); the tier ORDERING the assertions +gate is robust to this. + +The remaining build-plane residue is the ~4.1 GB @1M hash-slot LIST-pointer +term, which interning cannot shrink (every re-listed slot still holds one +pointer regardless of value sharing). + +``AIPERF_TEST_DYNAMO_SCALE_NODES`` overrides the node count for a manual +corpus-scale run (following the ``AIPERF_TEST_WEKA_CORPUS_DIR`` precedent in +``tests/unit/graph/test_weka_trie_build_resolution.py``); the default scale is +calibrated to keep the module under ~120 s. +""" + +from __future__ import annotations + +import gc +import inspect +import math +import os +import sys +import time +import tracemalloc +from collections.abc import Callable + +import pytest + +from aiperf.common.tokenizer import BUILTIN_TOKENIZER_NAME +from tests.harness.dynamo_synth_corpus import ( + SyntheticCorpusShape, + synthetic_corpus_shape, + write_synthetic_dynamo_capture, +) + +# POSIX-only ru_maxrss RSS accounting; the whole module is a POSIX isolate. +resource = pytest.importorskip("resource") + +# --- calibrated default scale ---------------------------------------------- +# S x T x G chosen so N ~= 6k nodes with a realistic chain shape (G unique +# blocks/turn, T turns => H/N = G*(T+1)/2 ~= 512 hash slots/node, matching real +# corpus captures) and the whole slow module runs in well under ~120 s. +_DEFAULT_SESSIONS = 150 +_TURNS_PER_SESSION = 40 +_NEW_BLOCKS_PER_TURN = 25 +_BLOCK_SIZE = 16 +_CONTENT_SEED = 1234 +_EXTRAPOLATION_NODES = 1_000_000 + +# ``@pytest.mark.slow`` is deselected by default; opt in with ``-m slow``. The +# env override lets a developer size a manual corpus-scale run. +_SCALE_ENV = "AIPERF_TEST_DYNAMO_SCALE_NODES" +# nframes=1 keeps tracemalloc's per-allocation cost low (the top frame is the +# real allocation site, which is all ``statistics("filename")`` groups on) while +# the parse stays measurable in seconds rather than minutes. +_TRACE_NFRAMES = 1 +# Measured total peak must stay within this multiple of the analytic per-tier +# model computed from the generator parameters (a calibrated RATIO, never an +# absolute-bytes gate): it catches a gross per-node regression without being +# fragile to interpreter-version object-layout drift. +_BUDGET_RATIO = 1.5 +# Read-time interning shaves the recorded hash-id-int tier from one +# distinct int per re-listed slot to one per UNIQUE value. This calibrated RATIO +# gates that win permanently (following the ``_FLAT_BUDGET_RATIO`` precedent -- +# calibrated ratio, never absolute bytes): the measured hash-id tier must stay at +# most this fraction of the distinct-int cost (``h`` slots x (int + pointer)). +# Post-intern measures ~0.32; a regression back to distinct ints per slot +# measures ~1.0, so 0.55 catches the regression with interpreter-drift headroom. +_HASH_INTERN_RATIO = 0.55 +# Pool interning shaves the sid-string ADDRESSING tier (every node's +# prompt_segment_ids re-lists its whole chain) from one fresh hexdigest per slot to +# one canonical object per UNIQUE segment. This calibrated RATIO gates that win +# permanently (never absolute bytes): the retained ``segment_id`` hexdigest bytes +# must stay at most this fraction of ``n_unique_segments * sys.getsizeof(sid)``. +# Post-intern measures ~1.0-1.1x (the floor itself); a 246k-duplicate regression +# overshoots ~10x, so 2.0 catches it with headroom for the getsizeof-vs-tracemalloc +# base mismatch (see ``_SID_STR``). +_SID_INTERN_RATIO = 2.0 +# A 32-hex-char ``segment_id`` (blake2b digest_size=16 -> 32 hex chars). getsizeof +# reports ~81 B for this compact-ASCII str; tracemalloc's per-allocation accounting +# observes ~73 B/str, so the ratio-gate denominator (getsizeof) runs ~11% HIGHER +# than the measured numerator base -- the ``_SID_INTERN_RATIO`` headroom absorbs +# that unit mismatch so it is not silently load-bearing. +_SID_STR = sys.getsizeof("0" * 32) + +# --- CPython object-layout sizes for the analytic model -------------------- +# Measured at import (getsizeof) or documented structural constants -- per-object +# sizes scaled by param-derived element counts, NOT hardcoded byte totals. +_PTR = 8 +_INT64 = sys.getsizeof(1 << 63) # a u64-valued int object (above the small-int cache) +_LIST_HDR = sys.getsizeof([]) # empty list container header +_LIST_BLOCK = sys.getsizeof([0] * _BLOCK_SIZE) # one decoded block's token list +# Amortized combined-dict slot (index array + key/value/hash entry at the ~2/3 +# load factor CPython dicts grow to); covers the decode cache and pool dicts. +_DICT_ENTRY = 110 +# Per unique content block: a slots ``Segment`` plus its ``_by_id`` dict slot. +_SEGMENT_ENTRY = 150 +# Per node, the persistent non-hash overhead: the ``AgentTraceRecord`` pydantic +# object graph, the ``TrieNode``/``TrieRequest``, the assembled ``LlmNode`` and +# its metadata dict. +_NODE_FIXED = 3600 + + +def _resolve_scale() -> tuple[int, int, int]: + """``(sessions, turns_per_session, new_blocks_per_turn)`` for this run. + + ``AIPERF_TEST_DYNAMO_SCALE_NODES`` overrides the node count (sessions are + sized to hit it at the fixed turn/block shape); unset uses the calibrated + default. + """ + turns, g = _TURNS_PER_SESSION, _NEW_BLOCKS_PER_TURN + override = os.environ.get(_SCALE_ENV) + if override: + sessions = max(1, math.ceil(int(override) / turns)) + else: + sessions = _DEFAULT_SESSIONS + return sessions, turns, g + + +def _analytic_model_bytes(shape: SyntheticCorpusShape) -> int: + """Analytic peak-window per-tier model, computed from the generator counts. + + Sums the tiers that coexist at the emission peak (hash-id ints + list + containers, the decode cache, the content pool, and the per-node fixed + overhead) using the CPython object-layout sizes above. The resolution trie + is a transient freed before this window, so it is deliberately excluded. + """ + n, h, u = shape.n_nodes, shape.n_hash_slots, shape.n_unique_blocks + # Read-time interning canonicalizes the recorded hash ints in + # ``_collect_records``, so the int OBJECTS scale with UNIQUE values (``u``), + # not total slots (``h``); only the per-slot list POINTERS (``h * _PTR``) and + # the per-node list headers (``n * _LIST_HDR``) still scale with ``h``/``n``. + hash_tier = u * _INT64 + h * _PTR + n * _LIST_HDR + # Deliberately over-modeled: this term sizes a full per-hash list cache + # (one ``_LIST_BLOCK`` of block_size ints per unique + # hash), while the actual dynamo decode cache is one int OFFSET per unique hash + # (``_decode_block_tokens_offset_cached``), so the true structural size is + # ~``u * (_INT64 + _DICT_ENTRY)``. It is kept as a CONSERVATIVE per-tier + # ceiling: shrinking it to the offset shape drops the model ~22 MB and leaves + # the measured ~140 MB peak only ~3% under the 1.5x budget -- fragile to + # interpreter object-layout drift. The interning-specific regression is gated + # sharply by the ``_HASH_INTERN_RATIO`` hash-tier assertion instead, so this + # looseness does not weaken the interning proof. + decode_tier = u * (_LIST_BLOCK + _DICT_ENTRY) + content_tier = u * _SEGMENT_ENTRY + # Canonical sid-string addressing tier. Post-intern every node's + # prompt_segment_ids shares one str object per UNIQUE segment (``S*(2T-1)+N``), + # NOT per unique block (``u``) -- the floor scales with segments, not hashes. + sid_tier = shape.n_unique_segments * _SID_STR + fixed = n * _NODE_FIXED + return hash_tier + decode_tier + content_tier + sid_tier + fixed + + +def _attribute_by_filename(snapshot: tracemalloc.Snapshot) -> dict[str, int]: + """Map absolute source filename -> summed allocated bytes at the snapshot.""" + by: dict[str, int] = {} + for stat in snapshot.statistics("filename"): + fn = stat.traceback[0].filename + by[fn] = by.get(fn, 0) + stat.size + return by + + +def _tier_bytes(by_file: dict[str, int], *needles: str) -> int: + """Sum bytes for every source file whose path contains any of ``needles``.""" + return sum(sz for fn, sz in by_file.items() if any(n in fn for n in needles)) + + +def _segment_id_line_range() -> tuple[str, int, int]: + """``(filename, first_lineno, last_lineno)`` of ``pool.segment_id``. + + Derived dynamically via ``inspect.getsourcelines`` so a docstring-only edit to + the frozen ``pool.py`` cannot silently mis-target the addressing-tier split. + ``pool.py`` has THREE id-minting hexdigest sites (``segment_id``, + ``text_segment_id``, ``raw_segment_id``); the dynamo route mints only via + ``segment_id``, so the range is scoped to that function alone. + """ + from aiperf.dataset.graph.segment_ir.pool import segment_id + + lines, start = inspect.getsourcelines(segment_id) + return segment_id.__code__.co_filename, start, start + len(lines) - 1 + + +def _addressing_tier_bytes(snapshot: tracemalloc.Snapshot) -> int: + """Bytes retained by the ``segment_id`` hexdigest line -- the sid-string tier. + + ``statistics("lineno")`` groups allocations by (filename, lineno); summing the + lines inside ``segment_id``'s source range isolates the ``h.hexdigest()`` sid + strings that every node's ``prompt_segment_ids`` retains, excluding the sibling + ``text_segment_id`` / ``raw_segment_id`` mints and the resident ``Segment`` + graph that share the ``pool.py`` FILENAME tier. + """ + fname, lo, hi = _segment_id_line_range() + total = 0 + for stat in snapshot.statistics("lineno"): + frame = stat.traceback[0] + if frame.filename == fname and lo <= frame.lineno <= hi: + total += stat.size + return total + + +def _extrapolate_gb(byte_count: int, n_nodes: int) -> float: + """Linear per-node extrapolation of ``byte_count`` to ``_EXTRAPOLATION_NODES``.""" + return byte_count / n_nodes * _EXTRAPOLATION_NODES / 1e9 + + +def _fmt_tier(label: str, byte_count: int, n_nodes: int) -> str: + return ( + f" {label:<34} {byte_count / 1e6:8.1f} MB " + f"{byte_count / n_nodes:8.0f} B/node " + f"-> {_extrapolate_gb(byte_count, n_nodes):6.1f} GB @ {_EXTRAPOLATION_NODES:,} nodes" + ) + + +def _prewarm_synthesizer() -> int: + """Build the shared corpus synthesizer OUTSIDE any traced region. + + Returns the resolved content seed. Pre-warming means the traced parse reuses + the cached synthesizer, so the fixed ~600k-token corpus pool is excluded from + the per-node tier measurement (the tiers become marginal per-node costs). + """ + from aiperf.dataset.graph.adapters.shared.content import ( + CorpusContentSynthesizer, + get_or_build_synthesizer, + resolve_effective_root_seed, + ) + + CorpusContentSynthesizer.reset_worker_cache() + seed = resolve_effective_root_seed(_CONTENT_SEED) + get_or_build_synthesizer( + BUILTIN_TOKENIZER_NAME, prompt_corpus="coding", root_seed=seed + ) + return seed + + +@pytest.fixture(scope="module") +def capture_path(tmp_path_factory: pytest.TempPathFactory) -> str: + """Write the synthetic capture once for the module's slow tests.""" + sessions, turns, g = _resolve_scale() + path = tmp_path_factory.mktemp("dynamo_corpus") / "synth.jsonl" + write_synthetic_dynamo_capture( + path, + sessions=sessions, + turns_per_session=turns, + new_blocks_per_turn=g, + block_size=_BLOCK_SIZE, + seed=1, + ) + return str(path) + + +@pytest.fixture(scope="module") +def prebuilt_nodes(capture_path: str): + """Parse chains -> ``TrieNode``s ONCE (untraced) for the phase isolates. + + The isolates measure a single downstream structure (the resolution trie, the + decode cache), so the node set they run on is built outside their traced + region -- only the structure under test is attributed. + """ + from aiperf.common.environment import Environment + from aiperf.dataset.graph.adapters.dynamo.trace import _collect_chains + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import dynamo_trie_nodes + + chains = _collect_chains( + capture_path, None, max_depth=Environment.DYNAMO.MAX_SUBAGENT_DEPTH + ) + nodes, block_size, _ = dynamo_trie_nodes(chains, release_replay=True) + return nodes, block_size + + +def test_synthetic_capture_parses_through_from_dynamo_trace(tmp_path) -> None: + """Fast guard: a tiny synthetic capture parses cleanly end-to-end. + + Not slow-marked so the default unit run pins the generator contract: the + block-alignment gate (``_assert_block_aligned``) and the covered-count ISL + gate both accept the generated ``input_length`` / hash-count relationship, + and every node lowers to a ``prompt=[]`` trie ``LlmNode`` addressed by the + segment pool. Protects the slow measurement from silent generator drift. + """ + from aiperf.dataset.graph.adapters.dynamo.trace import from_dynamo_trace + from aiperf.dataset.graph.models import LlmNode + + path = tmp_path / "tiny.jsonl" + write_synthetic_dynamo_capture( + path, + sessions=3, + turns_per_session=4, + new_blocks_per_turn=2, + block_size=_BLOCK_SIZE, + seed=7, + ) + parsed = from_dynamo_trace( + path, content_root_seed=_CONTENT_SEED, content_tokenizer="builtin" + ) + assert parsed.segment_pool is not None and parsed.segment_pool.by_id + # Dynamo emits one GraphRecord per independent session-tree (multi-graph), so + # gather LlmNodes across every per-tree graph, not just the first (graph). + graphs = list(parsed.graphs.values()) or [parsed.graph] + llm_nodes = [n for g in graphs for n in g.nodes.values() if isinstance(n, LlmNode)] + assert len(llm_nodes) == 3 * 4 + assert all(n.prompt == [] for n in llm_nodes) + assert all(n.metadata["trie"]["prompt_segment_ids"] for n in llm_nodes), ( + "each node must address its prompt through the segment pool" + ) + + +@pytest.mark.slow +def test_full_parse_tier_measurement( + capture_path: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """Size the persistent build tiers on a full ``from_dynamo_trace`` parse. + + Runs the REAL ``from_dynamo_trace`` (no re-implementation) under + ``tracemalloc``; a spy on ``stamp_theoretical_prefix_cache`` (its last call + before the parse returns) snapshots the peak window while the recorded hash + lists, decode cache, content pool, and assembled nodes are ALL still live -- + a snapshot taken after the call returns would see every transient already + freed. Attributes the snapshot by allocation-site filename into the named + tiers, asserts each is non-zero, checks the calibrated-ratio budget, and logs + the extrapolated tier table. + """ + sessions, turns, g = _resolve_scale() + shape = synthetic_corpus_shape( + sessions=sessions, turns_per_session=turns, new_blocks_per_turn=g + ) + n_nodes = shape.n_nodes + + _prewarm_synthesizer() + + from aiperf.dataset.graph.adapters.dynamo.trace import from_dynamo_trace + from aiperf.dataset.graph.segment_ir import prefix_cache + + captured: dict[str, tracemalloc.Snapshot] = {} + orig_stamp = prefix_cache.stamp_theoretical_prefix_cache + + def _spy(llm_nodes, nodes) -> None: + orig_stamp(llm_nodes, nodes) + captured["peak_window"] = tracemalloc.take_snapshot() + + monkeypatch.setattr(prefix_cache, "stamp_theoretical_prefix_cache", _spy) + + gc.collect() + tracemalloc.start(_TRACE_NFRAMES) + t0 = time.perf_counter() + parsed = from_dynamo_trace( + capture_path, + content_root_seed=_CONTENT_SEED, + content_tokenizer="builtin", + release_replay=True, + ) + elapsed = time.perf_counter() - t0 + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + + assert parsed.segment_pool is not None + assert "peak_window" in captured, "spy never fired -- parse path changed?" + by_file = _attribute_by_filename(captured["peak_window"]) + + hash_id_bytes = _tier_bytes( + by_file, + "dynamo/trace_reader.py", + "dynamo/trie_lowering.py", + "dynamo/trace.py", + ) + decode_bytes = _tier_bytes(by_file, "adapters/shared/content.py") + content_bytes = _tier_bytes(by_file, "segment_ir/pool.py") + trie_bytes = _tier_bytes(by_file, "segment_ir/trie_content.py") + + table = "\n".join( + [ + f"[dynamo corpus-scale memory] full parse: N={n_nodes:,} nodes " + f"H={shape.n_hash_slots:,} slots U={shape.n_unique_blocks:,} unique " + f"blocks in {elapsed:.1f}s (nframes={_TRACE_NFRAMES})", + _fmt_tier("hash-id ints+lists (trace/lower)", hash_id_bytes, n_nodes), + _fmt_tier("decode cache (shared/content)", decode_bytes, n_nodes), + _fmt_tier("content pool (segment_ir/pool)", content_bytes, n_nodes), + _fmt_tier("trie_content (build artifacts)", trie_bytes, n_nodes), + f" {'MEASURED TOTAL PEAK':<34} {peak / 1e6:8.1f} MB " + f"{peak / n_nodes:8.0f} B/node " + f"-> {_extrapolate_gb(peak, n_nodes):6.1f} GB @ {_EXTRAPOLATION_NODES:,} nodes", + ] + ) + print("\n" + table) + + # Each named tier must attribute non-zero (a zero means the allocation site + # moved and the measurement is silently mis-attributing). + assert hash_id_bytes > 0, f"hash-id tier attributed 0 bytes\n{table}" + assert decode_bytes > 0, f"decode-cache tier attributed 0 bytes\n{table}" + assert content_bytes > 0, f"content-pool tier attributed 0 bytes\n{table}" + assert trie_bytes > 0, f"trie_content tier attributed 0 bytes\n{table}" + + # Calibrated-ratio budget (never absolute bytes): measured peak within + # 1.5x the analytic per-tier model derived from the generator parameters. + model = _analytic_model_bytes(shape) + assert peak <= _BUDGET_RATIO * model, ( + f"measured peak {peak / 1e6:.1f} MB exceeds " + f"{_BUDGET_RATIO} x analytic model {model / 1e6:.1f} MB " + f"(= {_BUDGET_RATIO * model / 1e6:.1f} MB) -- a per-node allocation " + f"regression, not interpreter drift\n{table}" + ) + + # Read-time interning gate: the recorded hash-id tier must stay + # well under the distinct-int cost (one int per re-listed slot). Measures + # ~0.32 post-intern; a regression to distinct ints per slot measures ~1.0. + hash_intern_budget = _HASH_INTERN_RATIO * shape.n_hash_slots * (_INT64 + _PTR) + assert hash_id_bytes <= hash_intern_budget, ( + f"hash-id tier {hash_id_bytes / 1e6:.1f} MB exceeds {_HASH_INTERN_RATIO} x " + f"the distinct-int cost {hash_intern_budget / 1e6:.1f} MB -- read-time " + f"interning in _collect_records regressed. CAVEAT: this tier sums the " + f"trace_reader / trie_lowering / trace needles, which also capture " + f"record-model allocations, so a future record-model growth can fire this " + f"assertion without an interning regression\n{table}" + ) + + +def _measure_pool_tier_at_peak( + capture_path: str, + real_stamp: Callable[[object, object], None], + *, + direct_store: object | None, +) -> tuple[int, int, int, tracemalloc.Snapshot]: + """Full ``from_dynamo_trace``; return ``(pool.py bytes, total peak, resident segments, snapshot)``. + + A local spy on ``stamp_theoretical_prefix_cache`` (its last call before the + parse returns) snapshots the peak window while every parse structure is still + live, so the content pool -- ``SegmentPool._by_id`` on the eager route, empty + on the ``StoreBackedSegmentPool`` write-through route -- is attributed to + ``segment_ir/pool.py`` at its resident maximum. ``real_stamp`` is the TRUE + original (captured before any patch) so repeated calls never wrap a prior spy. + + The third return value is ``len(segment_pool.by_id)`` -- 0 on the direct route + (the pool holds nothing), non-zero on the eager route. + """ + from aiperf.dataset.graph.adapters.dynamo.trace import from_dynamo_trace + from aiperf.dataset.graph.segment_ir import prefix_cache + + _prewarm_synthesizer() + captured: dict[str, tracemalloc.Snapshot] = {} + + def _spy(llm_nodes, nodes) -> None: + real_stamp(llm_nodes, nodes) + captured["peak_window"] = tracemalloc.take_snapshot() + + prev = prefix_cache.stamp_theoretical_prefix_cache + prefix_cache.stamp_theoretical_prefix_cache = _spy + try: + gc.collect() + tracemalloc.start(_TRACE_NFRAMES) + parsed = from_dynamo_trace( + capture_path, + content_root_seed=_CONTENT_SEED, + content_tokenizer="builtin", + release_replay=True, + direct_store=direct_store, # type: ignore[arg-type] + ) + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + finally: + prefix_cache.stamp_theoretical_prefix_cache = prev + + assert parsed.segment_pool is not None + assert "peak_window" in captured, "spy never fired -- parse path changed?" + snapshot = captured["peak_window"] + pool_bytes = _tier_bytes(_attribute_by_filename(snapshot), "segment_ir/pool.py") + return pool_bytes, peak, len(parsed.segment_pool.by_id), snapshot + + +@pytest.mark.slow +def test_direct_route_content_tier_collapses(capture_path: str, tmp_path) -> None: + """Direct write-through route: the resident content pool collapses to ZERO. + + The Task-6 gate evidence ("Task-2 re-run shows parse-window content tier + approx 0"). The default full-parse tier measurement above exercises the EAGER + route (no ``direct_store``): every unique message segment lives as a resident + ``Segment`` in ``SegmentPool._by_id``. This test threads a live + ``GraphSegmentUnifiedBackingStore`` as ``direct_store`` so + ``StoreBackedSegmentPool.add`` writes each segment THROUGH to the store's + on-disk ``content.blob`` at parse time and the pool holds NOTHING. + + THE DEFINITIVE COLLAPSE is the resident ``Segment`` graph: ``by_id`` goes from + thousands of live ``Segment`` objects on the eager route to EXACTLY 0 on the + direct route (asserted). The freed ``Segment`` objects also stop pinning their + ``content`` strings alive, so the measured total peak drops too. + + NUANCE (informational, not asserted tightly): the + filename-attributed ``segment_ir/pool.py`` tier does NOT drop to ~0 on the + direct route, because it still carries the ``segment_id`` hex-digest strings + (``pool.py``'s ``h.hexdigest()`` line) that EVERY node's + ``metadata["trie"]["prompt_segment_ids"]`` retains on BOTH routes -- per-node + ADDRESSING cost, not the content pool. Pool interning collapses + that addressing tier from ~246k re-listed duplicates to one canonical str per + unique segment (this test splits it out via ``_addressing_tier_bytes`` and + gates it with ``_SID_INTERN_RATIO``), so the direct route removes the + resident ``Segment`` object graph (``_by_id``) AND leaves only the interned + ~1.3 MB addressing floor -- the delta measured here is the Segment graph. + """ + from aiperf.dataset.graph.segment_ir import prefix_cache + from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + ) + + sessions, turns, g = _resolve_scale() + shape = synthetic_corpus_shape( + sessions=sessions, turns_per_session=turns, new_blocks_per_turn=g + ) + n_nodes = shape.n_nodes + + # Capture the TRUE original ONCE so neither leg wraps the other's spy. + real_stamp = prefix_cache.stamp_theoretical_prefix_cache + + eager_pool_bytes, eager_peak, eager_segs, eager_snap = _measure_pool_tier_at_peak( + capture_path, real_stamp, direct_store=None + ) + + store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id="direct-content-tier" + ) + try: + ( + direct_pool_bytes, + direct_peak, + direct_segs, + direct_snap, + ) = _measure_pool_tier_at_peak(capture_path, real_stamp, direct_store=store) + finally: + store.abort() # close the write handle + unlink the partial store files + + # Addressing tier: the segment_id hexdigest sid strings every node's + # prompt_segment_ids retains, split out of the pool.py FILENAME tier via a + # statistics("lineno") filter scoped to segment_id (see _addressing_tier_bytes). + # Interning collapses ~246k re-listed duplicates to one canonical + # per unique segment on BOTH routes. + eager_addr_bytes = _addressing_tier_bytes(eager_snap) + direct_addr_bytes = _addressing_tier_bytes(direct_snap) + # The direct route's canonical _sids pointer list allocates in store_backed_pool.py + # (a distinct FILENAME tier under nframes=1), NOT pool.py -- report it separately. + direct_sids_bytes = _tier_bytes( + _attribute_by_filename(direct_snap), "adapters/dynamo/store_backed_pool.py" + ) + + # eager_pool_bytes and direct_pool_bytes share the identical segment_id + # hex-string allocation (byte-identical parse), so their delta isolates the + # resident Segment object graph the write-through removes. + resident_graph = eager_pool_bytes - direct_pool_bytes + print( + "\n[dynamo corpus-scale memory] parse-window content pool, " + f"N={n_nodes:,} nodes U={shape.n_unique_blocks:,} unique blocks " + f"seg={shape.n_unique_segments:,} unique segments:" + ) + print( + f" resident Segment objects (by_id): EAGER={eager_segs:,} -> DIRECT={direct_segs}" + ) + print(_fmt_tier("pool.py filename tier EAGER", eager_pool_bytes, n_nodes)) + print(_fmt_tier("pool.py filename tier DIRECT", direct_pool_bytes, n_nodes)) + print(_fmt_tier("addressing tier (segment_id) EAGER", eager_addr_bytes, n_nodes)) + print(_fmt_tier("addressing tier (segment_id) DIRECT", direct_addr_bytes, n_nodes)) + print( + _fmt_tier("_sids list (store_backed_pool) DIRECT", direct_sids_bytes, n_nodes) + ) + print( + f" resident Segment-graph removed by write-through: {resident_graph / 1e6:.2f} MB; " + f"total peak {eager_peak / 1e6:.1f} MB (eager) -> {direct_peak / 1e6:.1f} MB (direct)" + ) + + # The eager route must actually build a content pool (else the test is vacuous). + assert eager_segs > 0, "eager route interned 0 segments -- parse path changed?" + # Cross-check the harness formula against the measured resident segment count so + # the analytic sid tier + the ratio-gate denominator use a validated count. + assert eager_segs == shape.n_unique_segments, ( + f"measured unique segments {eager_segs:,} != harness formula " + f"{shape.n_unique_segments:,} (S*(2T-1)+N) -- generator shape drifted" + ) + # Sid-interning gate: the retained sid-string addressing tier must stay + # within _SID_INTERN_RATIO of the interned floor (one canonical str per unique + # segment). A regression to a fresh hexdigest per re-listed slot overshoots ~10x. + # Denominator base is sys.getsizeof (~81 B) vs the tracemalloc-observed ~73 B/str + # in the numerator (see _SID_STR) -- the 2.0x headroom absorbs that unit mismatch. + sid_intern_budget = _SID_INTERN_RATIO * shape.n_unique_segments * _SID_STR + assert eager_addr_bytes <= sid_intern_budget, ( + f"addressing tier {eager_addr_bytes / 1e6:.2f} MB exceeds {_SID_INTERN_RATIO}x " + f"the interned floor {sid_intern_budget / 1e6:.2f} MB " + f"({shape.n_unique_segments:,} unique segments x {_SID_STR} B) -- " + "pool interning regressed to a fresh sid per re-listed prompt_segment_ids slot" + ) + # THE gate assertion: on the direct route the pool holds NOTHING. + assert direct_segs == 0, ( + f"direct route pool retained {direct_segs} segments; the write-through shim " + "must intern every segment into the store and hold none in _by_id" + ) + # The resident Segment graph is a real, removed allocation, and freeing it (plus + # its pinned content strings) lowers the measured peak. + assert resident_graph > 0, ( + f"eager pool.py tier {eager_pool_bytes / 1e6:.2f} MB was not larger than the " + f"direct route's {direct_pool_bytes / 1e6:.2f} MB -- the resident Segment " + "graph should attribute to pool.py on the eager route only" + ) + assert direct_peak < eager_peak, ( + f"direct-route peak {direct_peak / 1e6:.1f} MB was not below eager " + f"{eager_peak / 1e6:.1f} MB -- the write-through should free content earlier" + ) + + +@pytest.mark.slow +def test_phase_isolate_resolution_trie(prebuilt_nodes) -> None: + """Compositional cross-check: the resolution-trie transient in isolation. + + Builds the flat int-state automaton with the REAL frozen + ``trie_content._insert_flat`` (a ``transitions`` dict + parallel + ``terminal``/``passer`` lists, state 0 = + root) over the same node set, holding the structures alive at the snapshot so + the transient (freed inside ``build_trie_ir`` before emission, hence + invisible to the full-parse snapshot) is attributed to ``trie_content.py``. + Logs the extrapolation. + """ + from aiperf.dataset.graph.segment_ir import trie_content + from aiperf.dataset.graph.segment_ir.trie_content import TrieNode + + nodes, _ = prebuilt_nodes + n_nodes = len(nodes) + + gc.collect() + gc.disable() + try: + tracemalloc.start(_TRACE_NFRAMES) + tracemalloc.reset_peak() + transitions: dict[tuple[int, int], int] = {} + terminal: list[TrieNode | None] = [None] + passer: list[TrieNode | None] = [None] + for node in nodes: + if node.request.hash_ids: + trie_content._insert_flat( + transitions, terminal, passer, node.request.hash_ids, node + ) + snap = tracemalloc.take_snapshot() + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + finally: + gc.enable() + + trie_bytes = _tier_bytes(_attribute_by_filename(snap), "segment_ir/trie_content.py") + line = _fmt_tier("resolution automaton (isolate)", trie_bytes, n_nodes) + print( + f"\n[dynamo corpus-scale memory] resolution-automaton isolate " + f"(peak {peak / 1e6:.1f} MB):\n{line}" + ) + assert trie_bytes > 0, "resolution automaton attributed 0 bytes to trie_content.py" + # Keep the structures reachable until after the snapshot is attributed. + del transitions, terminal, passer + + +@pytest.mark.slow +def test_phase_isolate_decode_cache(prebuilt_nodes) -> None: + """Compositional cross-check: the decode cache over the unique hash set. + + Drives ``dynamo_recon_callbacks``' decoder over every unique hash id exactly + once (the block cache is content-addressed, so this is the full steady-state + cache) and attributes it to ``shared/content.py``. Logs the extrapolation. + """ + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + dynamo_recon_callbacks, + ) + + nodes, block_size = prebuilt_nodes + n_nodes = len(nodes) + seed = _prewarm_synthesizer() + unique_hashes = sorted({h for node in nodes for h in node.request.hash_ids}) + + callbacks = dynamo_recon_callbacks( + BUILTIN_TOKENIZER_NAME, "coding", seed, block_size=block_size, trace_scope="t" + ) + + gc.collect() + gc.disable() + try: + tracemalloc.start(_TRACE_NFRAMES) + tracemalloc.reset_peak() + decoded = callbacks.decode_block_tokens(unique_hashes) + snap = tracemalloc.take_snapshot() + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + finally: + gc.enable() + + assert len(decoded) == len(unique_hashes) * block_size + decode_bytes = _tier_bytes( + _attribute_by_filename(snap), "adapters/shared/content.py" + ) + line = _fmt_tier("decode cache (isolate)", decode_bytes, n_nodes) + print( + f"\n[dynamo corpus-scale memory] decode-cache isolate over " + f"{len(unique_hashes):,} unique hashes (peak {peak / 1e6:.1f} MB):\n{line}" + ) + assert decode_bytes > 0, "decode cache attributed 0 bytes to shared/content.py" + + +@pytest.mark.slow +def test_peak_rss_logged(capture_path: str) -> None: + """Log peak process RSS across a full parse; NEVER assert on it. + + RSS assertions are environment-fragile (allocator high-water retention, the + RLIMIT_AS unit-suite gotcha), so this only surfaces the number for a human + reading the slow-run output -- it is not a gate. + """ + from aiperf.dataset.graph.adapters.dynamo.trace import from_dynamo_trace + + _prewarm_synthesizer() + before = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + parsed = from_dynamo_trace( + capture_path, + content_root_seed=_CONTENT_SEED, + content_tokenizer="builtin", + release_replay=True, + ) + after = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + assert parsed.segment_pool is not None + # ru_maxrss is KiB on Linux, bytes on macOS; report both interpretations. + print( + f"\n[dynamo corpus-scale memory] peak RSS after full parse: " + f"maxrss={after} (delta {after - before}); " + f"~{after * 1024 / 1e9:.2f} GB if KiB (Linux)" + ) + + +# --- flat-automaton memory budget (NOT slow-marked; ~1-2 s) ------------------ + +# Chain-heavy position count for the resolution-automaton budget. A pure chain +# (one owner, strictly increasing ids) is the worst case for the deleted +# node-object trie -- every interior node held a one-entry ``children`` dict +# (the 224 B killer) -- and the best case for the flat automaton, which is +# exactly the shape the resolution trie takes on deep dynamo prefix chains. +_FLAT_BUDGET_POSITIONS = 200_000 +# The flat automaton must use at most this fraction of the deleted trie's peak. +# Measured 0.51-0.59 across dict-resize phases on CPython +# 3.12; 0.80 leaves interpreter-drift headroom while still proving the win. +# (Distinct name from ``_BUDGET_RATIO`` above: that 1.5x gate belongs to the +# slow full-parse tier measurement and is read at call time -- rebinding the +# module global here would silently tighten it.) +_FLAT_BUDGET_RATIO = 0.80 + + +def _measure_peak(build: Callable[[], object]) -> int: + """Return tracemalloc peak bytes for ``build()`` with the result held alive. + + ``build`` returns the constructed structure; it stays referenced until after + the peak is read, so the peak reflects the fully-built structure plus its + construction transients (dict/list overallocation), not a half-collected one. + """ + gc.collect() + gc.disable() + try: + tracemalloc.start(1) + tracemalloc.reset_peak() + built = build() + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + del built + finally: + gc.enable() + return peak + + +class _ReplicaTrieNode: + """~10-line in-test replica of the DELETED ``trie_content._PrefixTrieNode``. + + A ``@dataclass(slots=True)`` with a ``children`` dict + ``terminal``/``passer`` + slots; ``__slots__`` reproduces the exact per-node memory layout so the + budget ratio compares the flat automaton against the real thing it replaced. + """ + + __slots__ = ("children", "terminal", "passer") + + def __init__(self) -> None: + self.children: dict = {} + self.terminal = None + self.passer = None + + +def test_flat_automaton_memory_budget() -> None: + """The flat int automaton uses <= 80% of the deleted node-object trie's peak. + + Calibrated-RATIO budget (the §6 requirement that replaced the plan's rejected + 220 B absolute gate): an absolute-bytes assert is dict-resize-phase- and + CPython-version-fragile, so this measures BOTH structures at the same + position count in-process and asserts the ratio. It is the memory half of + the frozen-code rewrite's proof -- the differential property test proves the + output is byte-identical; this proves the representation is actually leaner. + """ + from aiperf.dataset.graph.segment_ir import trie_content + from aiperf.dataset.graph.segment_ir.trie_content import ( + TrieNode, + TrieRequest, + ) + + positions = list(range(_FLAT_BUDGET_POSITIONS)) + owner = TrieNode( + node_id="owner", + request=TrieRequest( + hash_ids=positions, + input_length=1, + output_length=1, + t=0.0, + api_time=0.0, + ), + order=0, + ) + + def build_replica() -> _ReplicaTrieNode: + root = _ReplicaTrieNode() + cur = root + for h in positions: + child = cur.children.get(h) + if child is None: + child = _ReplicaTrieNode() + cur.children[h] = child + cur = child + if cur.passer is None: + cur.passer = owner + cur.terminal = owner + return root + + def build_flat() -> tuple[ + dict[tuple[int, int], int], list[TrieNode | None], list[TrieNode | None] + ]: + transitions: dict[tuple[int, int], int] = {} + terminal: list[TrieNode | None] = [None] + passer: list[TrieNode | None] = [None] + trie_content._insert_flat(transitions, terminal, passer, positions, owner) + return transitions, terminal, passer + + replica_peak = _measure_peak(build_replica) + flat_peak = _measure_peak(build_flat) + + ratio = flat_peak / replica_peak + print( + f"\n[flat-automaton budget @ {_FLAT_BUDGET_POSITIONS:,} positions] " + f"replica(node-object trie)={replica_peak / 1e6:.1f} MB, " + f"flat(automaton)={flat_peak / 1e6:.1f} MB, " + f"ratio={ratio:.3f} (must be <= {_FLAT_BUDGET_RATIO})" + ) + assert flat_peak <= _FLAT_BUDGET_RATIO * replica_peak, ( + f"flat automaton peak {flat_peak:,} B is not <= {_FLAT_BUDGET_RATIO} x the " + f"deleted node-object trie peak {replica_peak:,} B (ratio {ratio:.3f}); " + "the flat-automaton representation shave regressed -- investigate before landing." + ) diff --git a/tests/unit/dataset/graph/adapters/test_dynamo_gates.py b/tests/unit/dataset/graph/adapters/test_dynamo_gates.py new file mode 100644 index 0000000000..79046f424a --- /dev/null +++ b/tests/unit/dataset/graph/adapters/test_dynamo_gates.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fail-loud gates of the dynamo trie parse. + +Two gates, both raised from :func:`from_dynamo_trace` itself (the trie IR has +no separate build pass): + +* Block alignment: a recorded ``input_length`` not spanned by its replay + hashes at the recorded ``trace_block_size`` raises + :class:`DynamoISLMismatchError` -- no reconstruction can honor both fields. +* Mixed block sizes: two replay turns recording different ``trace_block_size`` + values raise :class:`DynamoTraceAdapterError`. +""" + +from __future__ import annotations + +from pathlib import Path + +import orjson +import pytest + +from aiperf.dataset.graph.adapters.dynamo.trace import ( + DynamoTraceAdapterError, + from_dynamo_trace, +) +from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + DynamoISLMismatchError, +) + + +def _rec( + *, + ts: int, + sid: str, + input_tokens: int, + output_tokens: int, + hashes: list[int] | None = None, + block_size: int = 16, + input_length: int | None = None, +) -> dict: + """One current-schema ``dynamo.request.trace.v1`` ``request_end`` record.""" + req: dict = { + "request_id": f"r{sid}{ts}", + "model": "m", + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "cached_tokens": 0, + } + if hashes is not None: + req["replay"] = { + "trace_block_size": block_size, + "input_length": input_length if input_length is not None else input_tokens, + "input_sequence_hashes": hashes, + } + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": {"session_id": sid}, + "request": req, + } + + +def _write(tmp_path: Path, name: str, records: list[dict]) -> Path: + p = tmp_path / name + p.write_bytes(b"\n".join(orjson.dumps(r) for r in records)) + return p + + +def test_isl_mismatch_aborts_parse(tmp_path): + """input_length=100 cannot be spanned by 2 hashes at bs=16 (16 < 100 > 32).""" + p = _write( + tmp_path, + "bad_isl.jsonl", + [ + _rec( + ts=1000, + sid="s1", + input_tokens=100, + output_tokens=8, + hashes=[111, 222], + block_size=16, + input_length=100, + ) + ], + ) + with pytest.raises(DynamoISLMismatchError): + from_dynamo_trace(p) + + +def test_mixed_block_size_aborts_parse(tmp_path): + p = _write( + tmp_path, + "mixed_bs.jsonl", + [ + _rec(ts=1000, sid="s1", input_tokens=32, output_tokens=8, hashes=[1, 2]), + _rec( + ts=2000, + sid="s1", + input_tokens=64, + output_tokens=8, + hashes=[1, 2, 3, 4], + block_size=32, + input_length=128, + ), + ], + ) + with pytest.raises(DynamoTraceAdapterError, match="trace_block_size"): + from_dynamo_trace(p) diff --git a/tests/unit/dataset/graph/adapters/test_dynamo_multigraph.py b/tests/unit/dataset/graph/adapters/test_dynamo_multigraph.py new file mode 100644 index 0000000000..b1d5519bc0 --- /dev/null +++ b/tests/unit/dataset/graph/adapters/test_dynamo_multigraph.py @@ -0,0 +1,277 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dynamo multi-graph restructure: per-tree ``TraceRecord`` + ``graph_ref``. + +The dynamo adapter emits ONE ``GraphRecord`` per session-tree (mirroring the +weka multi-item path through ``merge_parsed_graphs``): each trace carries its +OWN top-level graph under ``ParsedGraph.graphs[root_session_id]`` and its +``graph_ref`` selects it. The whole-capture union is NO LONGER the emitted +shape -- ``ParsedGraph.graph`` is now the FIRST tree (back-compat), not the +union. + +Two guards, both pinned at this tree: + +1. **Content oracle** (byte-equivalence): the per-tree lowering CONTENT is + identical to the pre-change single union -- only the ``ParsedGraph`` SHAPE + changed. The pinned digests were captured from the pre-change + ``from_dynamo_trace(fixture).graph`` (the single union). After the + restructure, re-unioning the per-tree graphs via the retained test-only + ``_union_graphs`` helper must reproduce those EXACT digests (nodes, edges, + content-addressed pool keys) -- because the per-tree graphs ARE the same + graphs the old path unioned. + +2. **Scoping**: ``parsed.graphs`` is non-empty (multi-graph), each + ``TraceRecord.graph_ref`` equals its root session id, ``resolve_trace_graph`` + returns ONLY that tree's nodes, and two distinct traces resolve to DISJOINT + node sets. + +RED on the pre-change single-union code: ``parsed.graphs`` is empty and +``graph_ref`` is ``None``, so the scoping asserts fail and the re-union oracle +sees an empty graph. +""" + +from __future__ import annotations + +import hashlib +from pathlib import Path + +import orjson +import pytest + +from aiperf.common.environment import Environment +from aiperf.dataset.graph.adapters.dynamo.trace import ( + _union_graphs, + from_dynamo_trace, +) +from aiperf.dataset.graph.adapters.shared.content import CorpusContentSynthesizer +from aiperf.dataset.graph.models import LlmNode, resolve_trace_graph + +_SEED = 20260706 +_TOKENIZER = "builtin" + +# Pinned oracle digests, captured from the PRE-CHANGE ``from_dynamo_trace().graph`` +# (the single union) over the fixture below. The per-tree graphs are byte- +# identical to the graphs the old path unioned, so re-unioning them reproduces +# these EXACTLY. A drift here is a real content change, not a shape change -- +# investigate, do not re-pin blindly. +# EDGES re-pinned 2026-07-07: recorded edge delays now always replay (weka +# idle-warp parity) instead of the old zeroed-delay dependency-only default; +# nodes and pool digests were unaffected, confirming content synthesis did not +# drift. +# All three re-pinned 2026-07-07 for the data-inherent node-id scheme: node +# ids became {session_id}:{k} (recorded session id verbatim, 0-based turn), +# which re-keys nodes/edges, and the trace-scoped response/tiny synthesis +# seeds shift the node-id-seeded pool entries. Hash-block prompt content is +# unchanged; the determinism guard (two independent builds agree) still holds. +_ORACLE_NODES_DIGEST = "2e825a321ea8cd7e4cb77ea5f1c86eb8" +_ORACLE_EDGES_DIGEST = "97782ee32139b01c801d0485c6976ffe" +_ORACLE_POOL_DIGEST = "f963d9998288d2d9fd1c66cd52844ad9" + + +@pytest.fixture(autouse=True) +def _fresh_synth_cache(): + """Isolate the process-level synthesizer cache from other tests.""" + CorpusContentSynthesizer.reset_worker_cache() + yield + CorpusContentSynthesizer.reset_worker_cache() + + +def _re( + *, + ts: int, + sid: str, + hashes: list[int], + ilen: int, + parent: str | None = None, + otok: int = 8, + bs: int = 16, +) -> dict: + ctx: dict = {"session_id": sid, "trajectory_id": sid} + if parent is not None: + ctx["parent_trajectory_id"] = parent + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": ctx, + "request": { + "request_id": f"r-{sid}-{ts}", + "model": "m", + "input_tokens": ilen, + "output_tokens": otok, + "cached_tokens": 0, + "ttft_ms": 10.0, + "replay": { + "trace_block_size": bs, + "input_length": ilen, + "input_sequence_hashes": hashes, + }, + }, + } + + +def _multi_tree_records() -> list[dict]: + """Three session-trees with DISJOINT content/sessions. + + * ``aaa``: root (2 turns) + subagent ``bbb`` (within-tree parent<->child edge) + * ``ccc``: root (2 turns) + * ``ddd``: root (single turn) + + Roots sort ``aaa < ccc < ddd``, so the first tree (``ParsedGraph.graph`` + back-compat) is the ``aaa`` tree. + """ + return [ + _re(ts=1000, sid="aaa", hashes=[1, 2], ilen=32), + _re(ts=1500, sid="bbb", parent="aaa", hashes=[7, 8], ilen=32), + _re(ts=2000, sid="aaa", hashes=[1, 2, 3], ilen=48), + _re(ts=5000, sid="ccc", hashes=[20, 21], ilen=32), + _re(ts=5100, sid="ccc", hashes=[20, 21, 22], ilen=48), + _re(ts=9000, sid="ddd", hashes=[30, 31], ilen=32), + ] + + +def _write_jsonl(path: Path, records: list[dict]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("wb") as f: + for r in records: + f.write(orjson.dumps(r)) + f.write(b"\n") + + +def _fixture(tmp_path: Path) -> Path: + p = tmp_path / "trace.jsonl" + _write_jsonl(p, _multi_tree_records()) + return p + + +def _digest(text: str) -> str: + return hashlib.blake2b(text.encode(), digest_size=16).hexdigest() + + +def _edge_key(edge) -> str: + return ( + f"{edge.source}->{edge.target}" + f"|dap={edge.delay_after_predecessor_us}" + f"|min={edge.min_start_delay_us}" + f"|daps={edge.delay_after_predecessor_start_us}" + f"|dapft={edge.delay_after_predecessor_first_token_us}" + ) + + +def _node_session(node: LlmNode) -> str: + return node.metadata["dynamo"]["session_id"] + + +def _parse(path: Path): + return from_dynamo_trace( + path, content_root_seed=_SEED, content_tokenizer=_TOKENIZER + ) + + +# --- 1. CONTENT ORACLE: union(new per-tree graphs) == pre-change union ------ + + +def test_reunion_of_per_tree_graphs_matches_pinned_single_union( + tmp_path: Path, +) -> None: + """union(list(parsed.graphs.values())) reproduces the pinned pre-change union. + + The retained test-only ``_union_graphs`` re-flattens the per-tree graphs; + since those graphs ARE the graphs the old path unioned, the nodes, edges, + and content-addressed pool keys are byte-identical to the single union + captured before the restructure. + """ + parsed = _parse(_fixture(tmp_path)) + + assert parsed.graphs, "multi-graph shape required (graphs must be populated)" + union = _union_graphs(list(parsed.graphs.values())) + + assert _digest("\n".join(sorted(union.nodes))) == _ORACLE_NODES_DIGEST + assert ( + _digest("\n".join(sorted(_edge_key(e) for e in union.edges))) + == _ORACLE_EDGES_DIGEST + ) + assert parsed.segment_pool is not None + assert _digest("\n".join(sorted(parsed.segment_pool.by_id))) == _ORACLE_POOL_DIGEST + + +def test_union_node_and_edge_counts_unchanged(tmp_path: Path) -> None: + """The re-union has the same 6 nodes / 6 edges the pre-change union had.""" + parsed = _parse(_fixture(tmp_path)) + union = _union_graphs(list(parsed.graphs.values())) + assert len(union.nodes) == 6 + assert len(union.edges) == 6 + + +# --- 2. SCOPING: per-trace graph_ref + tree-scoped resolution --------------- + + +def test_each_trace_is_its_own_single_root_tree(tmp_path: Path) -> None: + parsed = _parse(_fixture(tmp_path)) + + # One trace per tree, id-sorted, each keyed into graphs by its root id. + assert [t.id for t in parsed.traces] == ["aaa", "ccc", "ddd"] + for trace in parsed.traces: + assert trace.graph_ref == trace.id, ( + f"trace {trace.id!r} graph_ref must equal its root session id" + ) + assert trace.graph_ref in parsed.graphs + # multi-root tag is dropped: each tree is its own single-root trace. + assert "multi-root" not in trace.tags + assert "from-dynamo-trace" in trace.tags + + +def test_resolve_trace_graph_returns_only_that_trees_nodes(tmp_path: Path) -> None: + parsed = _parse(_fixture(tmp_path)) + + expected_sessions = {"aaa": {"aaa", "bbb"}, "ccc": {"ccc"}, "ddd": {"ddd"}} + for trace in parsed.traces: + graph = resolve_trace_graph(parsed, trace) + sessions = {_node_session(n) for n in graph.nodes.values()} + assert sessions == expected_sessions[trace.id], ( + f"trace {trace.id!r} resolved to sessions {sessions}" + ) + + +def test_two_distinct_traces_resolve_to_disjoint_node_sets(tmp_path: Path) -> None: + parsed = _parse(_fixture(tmp_path)) + by_id = {t.id: t for t in parsed.traces} + + aaa_nodes = set(resolve_trace_graph(parsed, by_id["aaa"]).nodes) + ccc_nodes = set(resolve_trace_graph(parsed, by_id["ccc"]).nodes) + assert aaa_nodes and ccc_nodes + assert aaa_nodes.isdisjoint(ccc_nodes) + + +def test_back_compat_graph_is_first_tree_not_union(tmp_path: Path) -> None: + """``ParsedGraph.graph`` is the FIRST tree (lex-min root ``aaa``), not the union.""" + parsed = _parse(_fixture(tmp_path)) + first_sessions = {_node_session(n) for n in parsed.graph.nodes.values()} + assert first_sessions == {"aaa", "bbb"} + assert len(parsed.graph.nodes) < 6 # strictly fewer than the union's 6 + + +# --- 3. fused-parallel path emits the SAME multi-graph shape ---------------- + + +def test_fused_parallel_emits_same_multigraph( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Forcing the pool path (per-tree ParsedGraph blobs from workers) yields the + identical multi-graph shape + content as the serial per-tree build.""" + path = _fixture(tmp_path) + serial = _parse(path) + + monkeypatch.setattr(Environment.DATASET, "DYNAMO_GRAPH_PARALLEL_THRESHOLD", 0) + monkeypatch.setattr(Environment.DATASET, "DYNAMO_GRAPH_PARALLEL_WORKERS", 2) + parallel = _parse(path) + + assert [t.id for t in parallel.traces] == [t.id for t in serial.traces] + assert {t.graph_ref for t in parallel.traces} == {"aaa", "ccc", "ddd"} + for tid in ("aaa", "ccc", "ddd"): + assert set(parallel.graphs[tid].nodes) == set(serial.graphs[tid].nodes) + p_union = _union_graphs(list(parallel.graphs.values())) + assert _digest("\n".join(sorted(p_union.nodes))) == _ORACLE_NODES_DIGEST + assert parallel.segment_pool is not None + assert parallel.segment_pool.by_id == serial.segment_pool.by_id diff --git a/tests/unit/dataset/graph/adapters/test_dynamo_parallel_build.py b/tests/unit/dataset/graph/adapters/test_dynamo_parallel_build.py new file mode 100644 index 0000000000..ac23a45574 --- /dev/null +++ b/tests/unit/dataset/graph/adapters/test_dynamo_parallel_build.py @@ -0,0 +1,470 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused read+build parallel dynamo session-tree build: scan, shuffle, equivalence. + +The fused path replaces the earlier ship-collected-trees pool: a cheap grouping +scan (session ids + parent links only, NO ``input_sequence_hashes`` parse) +decides tree membership in the parent, raw record lines are shuffled to per-batch +temp files, and each batch is READ+BUILT inside a worker so the recorded hash +arrays never cross a process boundary. The acceptance bar is byte-equivalence to +the serial tree-scoped build: the pinned content seed + block size and CONTIGUOUS +weight-balanced batching make the parent's in-order union of worker results +identical to the serial per-tree loop -- same node keys, same edge set, same +content-addressed pool. + +The pool is real (forkserver on Linux); fixtures are kept tiny and workers low so +the multiprocess tests stay fast and non-flaky. +""" + +from __future__ import annotations + +import gzip +import tempfile +from pathlib import Path + +import orjson +import pytest + +from aiperf.common.environment import Environment +from aiperf.dataset.graph.adapters.dynamo import trace_parallel +from aiperf.dataset.graph.adapters.dynamo.trace import ( + DynamoTraceAdapterError, + from_dynamo_trace, + root_of_sessions, +) +from aiperf.dataset.graph.adapters.shared.content import CorpusContentSynthesizer + +_SEED = 12345 + + +@pytest.fixture(autouse=True) +def _fresh_synth_cache(): + """Isolate the process-level synthesizer cache from other tests.""" + CorpusContentSynthesizer.reset_worker_cache() + yield + CorpusContentSynthesizer.reset_worker_cache() + + +# --- fixture builders (real capture agent_context shape) ------------------ + + +def _re( + *, + ts: int, + sid: str, + hashes: list[int], + ilen: int, + parent: str | None = None, + otok: int = 8, + bs: int = 16, +) -> dict: + ctx: dict = {"session_id": sid, "trajectory_id": sid} + if parent is not None: + ctx["parent_trajectory_id"] = parent + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": ctx, + "request": { + "request_id": f"r-{sid}-{ts}", + "model": "m", + "input_tokens": ilen, + "output_tokens": otok, + "cached_tokens": 0, + "ttft_ms": 10.0, + "replay": { + "trace_block_size": bs, + "input_length": ilen, + "input_sequence_hashes": hashes, + }, + }, + } + + +def _write_jsonl(path: Path, records: list[dict]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("wb") as f: + for r in records: + f.write(orjson.dumps(r)) + f.write(b"\n") + + +def _multi_tree_records() -> list[dict]: + """One parent+subagent tree ('P'<-'S') plus three independent root trees.""" + return [ + _re(ts=1000, sid="P", hashes=[1, 2], ilen=32), + _re(ts=2000, sid="S", parent="P", hashes=[7, 8], ilen=32), + _re(ts=3000, sid="P", hashes=[1, 2, 3], ilen=48), + _re(ts=1000, sid="a", hashes=[10, 11], ilen=32), + _re(ts=1100, sid="a", hashes=[10, 11, 12], ilen=48), + _re(ts=5000, sid="b", hashes=[20, 21], ilen=32), + _re(ts=6000, sid="c", hashes=[30, 31], ilen=32), + _re(ts=6100, sid="c", hashes=[30, 31, 32], ilen=48), + ] + + +def _force_parallel(monkeypatch: pytest.MonkeyPatch, *, workers: int) -> None: + monkeypatch.setattr(Environment.DATASET, "DYNAMO_GRAPH_PARALLEL_THRESHOLD", 0) + monkeypatch.setattr(Environment.DATASET, "DYNAMO_GRAPH_PARALLEL_WORKERS", workers) + + +def _edge_set(pg) -> set[tuple[str, str]]: + return {(e.source, e.target) for e in pg.graph.edges} + + +# --- 1. contiguous weight batching (pure) --------------------------------- + + +def test_contiguous_batches_are_order_preserving_and_cover_all() -> None: + items = list(range(10)) + weights = [1] * 10 + batches = trace_parallel._contiguous_weight_batches(items, weights, num_batches=3) + flat = [x for batch in batches for x in batch] + assert flat == items # contiguous, every item once + assert 1 <= len(batches) <= 3 + + +def test_contiguous_batches_heavy_item_closes_its_own_batch() -> None: + items = list(range(7)) + weights = [100] + [1] * 6 # one heavy leader + batches = trace_parallel._contiguous_weight_batches(items, weights, num_batches=3) + assert batches[0] == [0] + assert [x for batch in batches for x in batch] == items + + +def test_contiguous_batches_single_batch_returns_all() -> None: + items = [0, 1, 2, 3] + assert trace_parallel._contiguous_weight_batches( + items, [1, 1, 1, 1], num_batches=1 + ) == [items] + + +# --- 2. grouping scan: NO hash parse, cross-file links, block size -------- + + +def test_scan_groups_sessions_and_parent_links(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl(p, _multi_tree_records()) + + scan = trace_parallel._scan_grouping(p, threads=2) + + assert scan.request_end_sessions == {"P", "S", "a", "b", "c"} + assert scan.parent_link == {"S": "P"} + assert scan.block_size == 16 + # weight is the summed record-line byte length per session (a hash-free proxy) + assert set(scan.session_weight) == {"P", "S", "a", "b", "c"} + assert all(w > 0 for w in scan.session_weight.values()) + + +def test_scan_ignores_no_context_and_marker_lines(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl(p, [_re(ts=1000, sid="P", hashes=[1, 2], ilen=32)]) + # Append a schema-less S3 marker + a replay-only (no agent_context) record. + with p.open("ab") as f: + f.write(orjson.dumps({"verification": "trace-s3-uploader"})) + f.write(b"\n") + f.write( + orjson.dumps( + { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": 9, + "request": { + "request_id": "noctx", + "replay": { + "trace_block_size": 16, + "input_length": 32, + "input_sequence_hashes": [99, 98], + }, + }, + } + ) + ) + f.write(b"\n") + + scan = trace_parallel._scan_grouping(p, threads=1) + # Only the session-bearing record is grouped; the marker / no-context record + # carry no session id and are dropped (exactly as the serial reader drops them). + assert scan.request_end_sessions == {"P"} + assert scan.parent_link == {} + + +def test_scan_handles_sink_envelope_and_gzip(tmp_path: Path) -> None: + """Real captures wrap each line in a ``{"timestamp","event"}`` envelope + gzip.""" + p = tmp_path / "trace.jsonl.gz" + with gzip.open(p, "wb") as f: + for rec in _multi_tree_records(): + f.write(orjson.dumps({"timestamp": 1, "event": rec})) + f.write(b"\n") + + scan = trace_parallel._scan_grouping(p, threads=2) + assert scan.request_end_sessions == {"P", "S", "a", "b", "c"} + assert scan.parent_link == {"S": "P"} + + +def test_scan_rejects_mixed_block_size(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _re(ts=1000, sid="P", hashes=[1, 2], ilen=32, bs=16), + _re(ts=2000, sid="Q", hashes=[1, 2], ilen=64, bs=32), + ], + ) + with pytest.raises(DynamoTraceAdapterError, match="mixed replay trace_block_size"): + trace_parallel._scan_grouping(p, threads=2) + + +def test_scan_cross_file_parent_link(tmp_path: Path) -> None: + """A subagent in one file linking to a parent in another is one tree.""" + d = tmp_path / "capture" + _write_jsonl(d / "a.jsonl", [_re(ts=1000, sid="P", hashes=[1, 2], ilen=32)]) + _write_jsonl( + d / "b.jsonl", [_re(ts=2000, sid="S", parent="P", hashes=[7, 8], ilen=32)] + ) + + scan = trace_parallel._scan_grouping(d, threads=2) + root_of = root_of_sessions(scan.request_end_sessions, scan.parent_link) + assert root_of["S"] == "P" # child rooted at cross-file parent + assert sorted(set(root_of.values())) == ["P"] # one tree + + +# --- 3. byte-equivalence: parallel == serial ------------------------------ + + +def test_parallel_build_byte_identical_to_sequential( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl(p, _multi_tree_records()) + + # Serial reference: default threshold (8) with 4 trees stays in-process. + serial = from_dynamo_trace(p, content_root_seed=_SEED) + + _force_parallel(monkeypatch, workers=2) + parallel = from_dynamo_trace(p, content_root_seed=_SEED) + + # Assert equality, not counts. + assert set(parallel.graph.nodes) == set(serial.graph.nodes) + assert list(parallel.graph.nodes) == list(serial.graph.nodes) # insertion order + assert _edge_set(parallel) == _edge_set(serial) + assert parallel.segment_pool is not None and serial.segment_pool is not None + assert parallel.segment_pool.by_id == serial.segment_pool.by_id + assert parallel.traces[0].tags == serial.traces[0].tags + assert parallel.traces[0].id == serial.traces[0].id + + +def test_parallel_preserves_within_tree_parent_subagent_edge( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl(p, _multi_tree_records()) + + _force_parallel(monkeypatch, workers=2) + parallel = from_dynamo_trace(p, content_root_seed=_SEED) + + # P's first turn (1000) finished before S started (2000): within-tree edge. + assert ("P:0", "S:0") in _edge_set(parallel) + + +def test_parallel_cross_file_subagent_tree_preserved( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A parent in one segment file + its subagent in another still form ONE tree. + + Grouping happens ONCE in the parent over the FULL scan (all files), so a + cross-file parent<->subagent link is unioned into a single tree BEFORE + batching; the whole tree's records are shuffled to one batch and its + within-tree edge survives. Byte-identical to the serial directory build. + """ + d = tmp_path / "capture" + _write_jsonl( + d / "a.jsonl", + [ + _re(ts=1000, sid="P", hashes=[1, 2], ilen=32), + _re(ts=3000, sid="P", hashes=[1, 2, 3], ilen=48), + _re(ts=1000, sid="x", hashes=[40, 41], ilen=32), + ], + ) + _write_jsonl( + d / "b.jsonl", + [ + _re(ts=2000, sid="S", parent="P", hashes=[7, 8], ilen=32), + _re(ts=5000, sid="y", hashes=[50, 51], ilen=32), + _re(ts=6000, sid="z", hashes=[60, 61], ilen=32), + ], + ) + + serial = from_dynamo_trace(d, content_root_seed=_SEED) + + _force_parallel(monkeypatch, workers=2) + parallel = from_dynamo_trace(d, content_root_seed=_SEED) + + assert ("P:0", "S:0") in _edge_set(parallel) + assert set(parallel.graph.nodes) == set(serial.graph.nodes) + assert _edge_set(parallel) == _edge_set(serial) + assert parallel.segment_pool.by_id == serial.segment_pool.by_id + + +# --- 4. threshold fallback: no pool below threshold ----------------------- + + +def test_below_threshold_returns_none_no_pool( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """At/below threshold, ``maybe_build_fused_parallel`` returns None (serial).""" + p = tmp_path / "trace.jsonl" + _write_jsonl(p, _multi_tree_records()) + monkeypatch.setattr(Environment.DATASET, "DYNAMO_GRAPH_PARALLEL_THRESHOLD", 8) + + def _no_pool(*_a, **_k): + raise AssertionError("pool must not spawn below threshold") + + monkeypatch.setattr(trace_parallel, "_build_fused_parallel", _no_pool) + result = trace_parallel.maybe_build_fused_parallel( + p, + content_root_seed=_SEED, + idle_gap_cap_seconds=60.0, + content_tokenizer=None, + prompt_corpus="coding", + release_replay=False, + max_depth=8, + ) + assert result is None + + +def test_below_threshold_end_to_end_equals_and_skips_pool( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl(p, _multi_tree_records()) + + reference = from_dynamo_trace(p, content_root_seed=_SEED) + + # Threshold above the 4-tree count keeps the build serial; the pool entry + # raising proves it is never reached. + monkeypatch.setattr(Environment.DATASET, "DYNAMO_GRAPH_PARALLEL_THRESHOLD", 8) + + def _no_pool(*_a, **_k): + raise AssertionError("pool must not spawn below threshold") + + monkeypatch.setattr(trace_parallel, "_build_fused_parallel", _no_pool) + serial = from_dynamo_trace(p, content_root_seed=_SEED) + + assert set(serial.graph.nodes) == set(reference.graph.nodes) + assert _edge_set(serial) == _edge_set(reference) + assert serial.segment_pool.by_id == reference.segment_pool.by_id + + +def test_single_tree_stays_serial( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Threshold 0 but one tree -> worker count collapses to 1 -> serial.""" + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _re(ts=1000, sid="only", hashes=[1, 2], ilen=32), + _re(ts=2000, sid="only", hashes=[1, 2, 3], ilen=48), + ], + ) + monkeypatch.setattr(Environment.DATASET, "DYNAMO_GRAPH_PARALLEL_THRESHOLD", 0) + monkeypatch.setattr(Environment.DATASET, "DYNAMO_GRAPH_PARALLEL_WORKERS", 4) + + def _no_pool(*_a, **_k): + raise AssertionError("single tree must not spawn a pool") + + monkeypatch.setattr(trace_parallel, "_build_fused_parallel", _no_pool) + result = trace_parallel.maybe_build_fused_parallel( + p, + content_root_seed=_SEED, + idle_gap_cap_seconds=60.0, + content_tokenizer=None, + prompt_corpus="coding", + release_replay=False, + max_depth=8, + ) + assert result is None + + +# --- 5. temp-dir cleanup on success AND exception ------------------------- + + +def test_temp_dir_cleaned_up_on_success( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + created: list[str] = [] + real_mkdtemp = tempfile.mkdtemp + + def _tracking_mkdtemp(*a, **k): + d = real_mkdtemp(*a, **k) + # Only the fused build's own temp dir; other machinery (pool, shm) may + # also call mkdtemp and clean up on its own schedule. + if Path(d).name.startswith("aiperf-dynamo-fused-"): + created.append(d) + return d + + monkeypatch.setattr(tempfile, "mkdtemp", _tracking_mkdtemp) + + p = tmp_path / "trace.jsonl" + _write_jsonl(p, _multi_tree_records()) + _force_parallel(monkeypatch, workers=2) + from_dynamo_trace(p, content_root_seed=_SEED) + + assert created, "fused path must have created a temp dir" + assert not any(Path(d).exists() for d in created), "temp dir left behind" + + +def test_temp_dir_cleaned_up_on_exception( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + created: list[str] = [] + real_mkdtemp = tempfile.mkdtemp + + def _tracking_mkdtemp(*a, **k): + d = real_mkdtemp(*a, **k) + # Only the fused build's own temp dir; other machinery (pool, shm) may + # also call mkdtemp and clean up on its own schedule. + if Path(d).name.startswith("aiperf-dynamo-fused-"): + created.append(d) + return d + + monkeypatch.setattr(tempfile, "mkdtemp", _tracking_mkdtemp) + + # Make the pool round blow up AFTER the temp dir + shuffle exist. + from aiperf.dataset.graph.adapters.weka import trace_parallel as weka_tp + + def _boom(*_a, **_k): + raise RuntimeError("boom in pool round") + + monkeypatch.setattr(weka_tp, "_run_pool_streaming", _boom) + + p = tmp_path / "trace.jsonl" + _write_jsonl(p, _multi_tree_records()) + _force_parallel(monkeypatch, workers=2) + + with pytest.raises(RuntimeError, match="boom in pool round"): + from_dynamo_trace(p, content_root_seed=_SEED) + + assert created, "fused path must have created a temp dir" + assert not any(Path(d).exists() for d in created), "temp dir left behind on error" + + +# --- 6. determinism -------------------------------------------------------- + + +def test_two_parallel_builds_identical( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl(p, _multi_tree_records()) + + _force_parallel(monkeypatch, workers=3) + first = from_dynamo_trace(p, content_root_seed=999) + second = from_dynamo_trace(p, content_root_seed=999) + + assert set(first.graph.nodes) == set(second.graph.nodes) + assert _edge_set(first) == _edge_set(second) + assert first.segment_pool.by_id == second.segment_pool.by_id diff --git a/tests/unit/dataset/graph/adapters/test_dynamo_reader_schema.py b/tests/unit/dataset/graph/adapters/test_dynamo_reader_schema.py new file mode 100644 index 0000000000..75df6a69bd --- /dev/null +++ b/tests/unit/dataset/graph/adapters/test_dynamo_reader_schema.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import gzip +from pathlib import Path + +import orjson + +from aiperf.dataset.graph.adapters.dynamo.trace_reader import ( + AgentTraceRecord, + iter_trace_records, +) + + +def _write(tmp_path, records) -> Path: + p = tmp_path / "t.jsonl.gz" + with gzip.open(p, "wb") as f: + for r in records: + f.write(orjson.dumps(r) + b"\n") + return p + + +def _req(session, ts, hashes=None): + rec = { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "agent_context": {"session_id": session}, + "request": { + "request_id": f"{session}-{ts}", + "input_tokens": 32, + "output_tokens": 16, + }, + } + if hashes is not None: + rec["request"]["replay"] = { + "trace_block_size": 16, + "input_length": 32, + "input_sequence_hashes": hashes, + } + return rec + + +def test_parse_current_schema(tmp_path): + p = _write(tmp_path, [_req("s1", 1000, [11, 22])]) + recs = list(iter_trace_records(p)) + assert len(recs) == 1 + assert recs[0].agent_context.session_id == "s1" + assert recs[0].request.replay.input_sequence_hashes == [11, 22] + + +def test_replay_only_record_without_agent_context_or_source(tmp_path): + # New schema: event_source + agent_context may be absent; model optional. + p = _write( + tmp_path, + [ + { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": 1000, + "request": { + "request_id": "r1", + "output_tokens": 4, + "replay": { + "trace_block_size": 2, + "input_length": 4, + "input_sequence_hashes": [1, 2], + }, + }, + } + ], + ) + recs = list(iter_trace_records(p)) + assert recs[0].agent_context is None + assert recs[0].event_source is None + assert recs[0].request.model is None + + +def test_session_id_filter(tmp_path): + p = _write(tmp_path, [_req("s1", 1000), _req("s2", 1001)]) + recs = list(iter_trace_records(p, session_id="s2")) + assert len(recs) == 1 and recs[0].agent_context.session_id == "s2" + + +def test_record_type_is_agent_trace_record(tmp_path): + p = _write(tmp_path, [_req("s1", 1000)]) + recs = list(iter_trace_records(p)) + assert isinstance(recs[0], AgentTraceRecord) + + +def test_collect_records_groups_by_session_and_parent_link(tmp_path): + from aiperf.dataset.graph.adapters.dynamo.trace import ( + _collect_records, + ) + + p = _write( + tmp_path, + [ + _req("root", 1000), + { + **_req("child", 1100), + "agent_context": { + "session_id": "child", + "parent_session_id": "root", + }, + }, + ], + ) + by_session, parent_link, skipped_no_context = _collect_records(p, None) + assert set(by_session) == {"root", "child"} + assert parent_link == {"child": "root"} + assert skipped_no_context == 0 + + +def test_collect_records_interns_duplicate_replay_hash_objects(tmp_path): + """Read-time interning: equal u64 hash values re-listed across turns AND + across sessions collapse to ONE int object after ``_collect_records`` -- + values unchanged, ``id()``-set cardinality equal to the unique-value count. + + Uses values above CPython's small-int cache so ``orjson.loads`` allocates a + fresh object per occurrence (small ints like 11/22 are cached singletons and + would make the identity assertion vacuous). + """ + from aiperf.dataset.graph.adapters.dynamo.trace import ( + _collect_records, + ) + + a, b, c, d = 2**63 + 1, 2**63 + 2, 2**63 + 3, 2**63 + 4 + p = _write( + tmp_path, + [ + _req("s1", 1000, [a, b]), + _req("s1", 1001, [a, b, c]), + _req("s2", 1002, [a, d]), + ], + ) + by_session, _parent_link, _skipped = _collect_records(p, None) + + all_hashes = [ + h + for recs in by_session.values() + for r in recs + for h in r.request.replay.input_sequence_hashes + ] + # Values are preserved exactly (interning shares objects, never mutates). + assert sorted(all_hashes) == sorted([a, b, a, b, c, a, d]) + # Every occurrence of an equal value is now the SAME object; four unique + # values -> four distinct ids across all seven slots. + assert len({id(h) for h in all_hashes}) == len(set(all_hashes)) == 4 diff --git a/tests/unit/dataset/graph/adapters/test_dynamo_recorded_streaming.py b/tests/unit/dataset/graph/adapters/test_dynamo_recorded_streaming.py new file mode 100644 index 0000000000..bef2f5d622 --- /dev/null +++ b/tests/unit/dataset/graph/adapters/test_dynamo_recorded_streaming.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dynamo recorded-mode streaming derivation through ``from_dynamo_trace``. + +A recorded ``ttft_ms`` proves the original request streamed, so the lowered +``LlmNode.streaming`` must be ``ttft is not None`` (mirroring weka's recorded +``"n"``/``"s"`` discriminator); the build-plane envelope carries the mode to +the worker's per-request stream override from this native field. +""" + +from __future__ import annotations + +from pathlib import Path + +import orjson + +from aiperf.dataset.graph.adapters.dynamo.trace import from_dynamo_trace + + +def _rec(*, ts: int, sid: str, ttft_ms: float | None = None) -> dict: + req: dict = { + "request_id": f"{sid}-{ts}", + "model": "m", + "input_tokens": 32, + "output_tokens": 16, + "cached_tokens": 0, + "request_received_ms": ts - 100, + "total_time_ms": 100.0, + } + if ttft_ms is not None: + req["ttft_ms"] = ttft_ms + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": {"session_id": sid}, + "request": req, + } + + +def _write(tmp_path: Path, record: dict) -> Path: + p = tmp_path / "trace.jsonl" + p.write_bytes(orjson.dumps(record)) + return p + + +def test_ttft_bearing_turn_is_streaming(tmp_path: Path) -> None: + """A recorded ``ttft_ms`` lowers to a streaming node + stream override.""" + p = _write(tmp_path, _rec(ts=1000, sid="s1", ttft_ms=250.0)) + pb = from_dynamo_trace(p) + node = pb.graph.nodes["s1:0"] + assert node.streaming is True + + +def test_ttft_absent_turn_is_not_streaming(tmp_path: Path) -> None: + """No recorded ``ttft_ms`` lowers to a non-streaming node + stream override.""" + p = _write(tmp_path, _rec(ts=1000, sid="s1")) + pb = from_dynamo_trace(p) + node = pb.graph.nodes["s1:0"] + assert node.streaming is False + + +def test_ttft_zero_is_streaming(tmp_path: Path) -> None: + """``ttft_ms=0.0`` still streamed -- derivation is is-not-None, not truthiness.""" + p = _write(tmp_path, _rec(ts=1000, sid="s1", ttft_ms=0.0)) + pb = from_dynamo_trace(p) + node = pb.graph.nodes["s1:0"] + assert node.streaming is True diff --git a/tests/unit/dataset/graph/adapters/test_dynamo_seed_alignment.py b/tests/unit/dataset/graph/adapters/test_dynamo_seed_alignment.py new file mode 100644 index 0000000000..1680b9d05a --- /dev/null +++ b/tests/unit/dataset/graph/adapters/test_dynamo_seed_alignment.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import orjson + +from aiperf.common import random_generator as rng +from aiperf.dataset.graph.adapters.dynamo import trie_lowering +from aiperf.dataset.graph.adapters.dynamo.trace import from_dynamo_trace + + +def _dynamo_record(ts: int, sid: str, input_tokens: int, hashes: list[int]) -> dict: + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": {"session_id": sid}, + "request": { + "request_id": f"r{ts}", + "model": "m", + "input_tokens": input_tokens, + "output_tokens": 8, + "cached_tokens": 0, + "replay": { + "trace_block_size": 16, + "input_length": input_tokens, + "input_sequence_hashes": hashes, + }, + }, + } + + +def _write_trace(tmp_path: Path) -> Path: + p = tmp_path / "dyn_seed.jsonl" + records = [ + _dynamo_record(1000, "s1", 32, [111, 222]), + _dynamo_record(2000, "s1", 64, [111, 222, 333, 444]), + ] + p.write_bytes(b"\n".join(orjson.dumps(r) for r in records)) + return p + + +def _parse_capturing_seed(monkeypatch, path: Path, seed: int | None) -> int | None: + captured: list[int | None] = [] + real = trie_lowering.dynamo_recon_callbacks + + def spy(tokenizer: str, corpus: str, root_seed: int | None, **kwargs: Any): + captured.append(root_seed) + return real(tokenizer, corpus, root_seed, **kwargs) + + monkeypatch.setattr(trie_lowering, "dynamo_recon_callbacks", spy) + from_dynamo_trace(path, content_root_seed=seed) + assert len(captured) == 1 + return captured[0] + + +def test_from_dynamo_trace_explicit_seed_passes_through(monkeypatch, tmp_path) -> None: + assert _parse_capturing_seed(monkeypatch, _write_trace(tmp_path), 1234) == 1234 + + +def test_from_dynamo_trace_none_seed_uses_ambient_root_seed( + monkeypatch, tmp_path +) -> None: + rng.reset() + rng.init(777) + assert _parse_capturing_seed(monkeypatch, _write_trace(tmp_path), None) == 777 + + +def test_from_dynamo_trace_unseeded_generates_per_run_seed( + monkeypatch, tmp_path +) -> None: + # No ambient root seed: each parse resolves fresh OS entropy — concrete + # int threaded to the synthesizer, distinct across parses. + rng.reset() + path = _write_trace(tmp_path) + first = _parse_capturing_seed(monkeypatch, path, None) + rng.reset() + second = _parse_capturing_seed(monkeypatch, path, None) + assert isinstance(first, int) + assert isinstance(second, int) + assert first != second diff --git a/tests/unit/dataset/graph/adapters/test_dynamo_trace.py b/tests/unit/dataset/graph/adapters/test_dynamo_trace.py new file mode 100644 index 0000000000..e974f499d4 --- /dev/null +++ b/tests/unit/dataset/graph/adapters/test_dynamo_trace.py @@ -0,0 +1,809 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the Dynamo agent-trace -> flat segment-trie ParsedGraph adapter.""" + +from __future__ import annotations + +from pathlib import Path + +import orjson +import pytest + +from aiperf.dataset.graph.adapters.dynamo.trace import ( + DynamoTraceAdapter, + DynamoTraceAdapterError, + EmptyDynamoTraceError, + from_dynamo_trace, +) +from aiperf.dataset.graph.models import LlmNode, ParsedGraph, StaticEdge + +# --- fixture builders ----------------------------------------------------- + + +def _ctx( + *, + session_id: str, + parent_session_id: str | None = None, +) -> dict: + out: dict = { + "session_id": session_id, + } + if parent_session_id is not None: + out["parent_session_id"] = parent_session_id + return out + + +def _request_end( + *, + ts: int, + session_id: str, + parent_session_id: str | None = None, + request_id: str | None = None, + model: str = "m", + input_tokens: int | None = 10, + output_tokens: int | None = 20, + cached_tokens: int | None = 5, + kv_hit_rate: float | None = 0.5, + ttft_ms: float | None = 100.0, + replay: dict | None = None, +) -> dict: + req: dict = { + "request_id": request_id or f"r{ts}", + "model": model, + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "cached_tokens": cached_tokens, + "kv_hit_rate": kv_hit_rate, + "ttft_ms": ttft_ms, + } + if replay is not None: + req["replay"] = replay + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": _ctx( + session_id=session_id, + parent_session_id=parent_session_id, + ), + "request": req, + } + + +def _tool_end( + *, + ts: int, + session_id: str, + parent_session_id: str | None = None, + name: str = "search", + duration_ms: float | None = 50.0, + status: str = "succeeded", + tool_call_id: str | None = None, +) -> dict: + return { + "schema": "dynamo.request.trace.v1", + "event_type": "tool_end", + "event_time_unix_ms": ts, + "event_source": "harness", + "agent_context": _ctx( + session_id=session_id, + parent_session_id=parent_session_id, + ), + "tool": { + "tool_call_id": tool_call_id or f"tc{ts}", + "tool_class": name, + "duration_ms": duration_ms, + "status": status, + }, + } + + +def _write_jsonl(path: Path, records: list[dict]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("wb") as f: + for r in records: + f.write(orjson.dumps(r)) + f.write(b"\n") + + +def _static_edges(pb: ParsedGraph) -> set[tuple[str, str]]: + return {(e.source, e.target) for e in pb.graph.edges if isinstance(e, StaticEdge)} + + +# --- 1. flat lowering: one LlmNode per request_end ------------------------ + + +def test_single_session_three_turns_flat_nodes(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="p1"), + _request_end(ts=1100, session_id="p1"), + _request_end(ts=1200, session_id="p1"), + ], + ) + + pb = from_dynamo_trace(p) + + assert isinstance(pb, ParsedGraph) + aid = "p1" + assert sorted(pb.graph.nodes) == [f"{aid}:0", f"{aid}:1", f"{aid}:2"] + assert all(isinstance(n, LlmNode) for n in pb.graph.nodes.values()) + assert pb.segment_pool is not None + # Sequential recorded turns chain :0 -> :1 -> :2 (finished-before edges). + edges = _static_edges(pb) + assert ("START", f"{aid}:0") in edges + assert (f"{aid}:0", f"{aid}:1") in edges + assert (f"{aid}:1", f"{aid}:2") in edges + + +def test_single_trace_record_with_base_tag(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="p1"), + _request_end(ts=1100, session_id="p1"), + ], + ) + + pb = from_dynamo_trace(p) + + assert len(pb.traces) == 1 + assert pb.traces[0].id == "p1" + assert "from-dynamo-trace" in pb.traces[0].tags + assert "multi-root" not in pb.traces[0].tags + + +def test_multi_root_file_parses_one_trace_per_root_tree(tmp_path: Path) -> None: + """Two parentless root sessions become TWO per-tree traces (multi-graph).""" + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="p1"), + _request_end(ts=1010, session_id="p2"), + _request_end(ts=1100, session_id="p1"), + _request_end(ts=1110, session_id="p2"), + ], + ) + + pb = from_dynamo_trace(p) + + # One id-sorted trace per tree; each selects its own graph via graph_ref, and + # the multi-root tag is dropped (each tree is its own single-root trace). + assert [t.id for t in pb.traces] == ["p1", "p2"] + for trace in pb.traces: + assert trace.graph_ref == trace.id + assert "multi-root" not in trace.tags + # graph is the FIRST tree (p1); p2's node lives in its own graph, not here. + assert "p1:0" in pb.graph.nodes + assert "p2:0" not in pb.graph.nodes + assert "p2:0" in pb.graphs["p2"].nodes + + +def test_every_node_output_channel_declared_in_state(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="p1"), + _request_end(ts=1100, session_id="child", parent_session_id="p1"), + _request_end(ts=1200, session_id="p1"), + ], + ) + + pb = from_dynamo_trace(p) + + for nid, node in pb.graph.nodes.items(): + assert node.output == f"{nid}_out" + assert f"{nid}_out" in pb.graph.state + + +# --- 2. edge timing: recorded delays replay by default (idle-warped) -------- + + +def _two_turn_edge(pb: ParsedGraph) -> StaticEdge: + aid = "p1" + return next( + e + for e in pb.graph.edges + if isinstance(e, StaticEdge) + and (e.source, e.target) == (f"{aid}:0", f"{aid}:1") + ) + + +def test_edges_keep_recorded_delay_by_default(tmp_path: Path) -> None: + """Recorded end-to-start gaps replay on the edges (weka idle-warp parity): + a 1s gap sits below the shared 60s cap, so it survives unwarped.""" + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="p1"), + _request_end(ts=2000, session_id="p1"), + ], + ) + + pb = from_dynamo_trace(p) + + edge = _two_turn_edge(pb) + # 1000 ms recorded end-to-start gap (zero api_time) -> 1_000_000 us. + assert edge.delay_after_predecessor_us == pytest.approx(1_000_000.0) + # Arrival stamping rides the same warped clock. + aid = "p1" + assert pb.graph.nodes[f"{aid}:1"].arrival_offset_us == 1_000_000 + + +def test_idle_gap_cap_compresses_edge_delay(tmp_path: Path) -> None: + """A gap above the cap is compressed to the cap on the binding edge.""" + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="p1"), + _request_end(ts=2000, session_id="p1"), + ], + ) + + pb = from_dynamo_trace(p, idle_gap_cap_seconds=0.25) + + edge = _two_turn_edge(pb) + # The 1s recorded idle gap is warped down to the 0.25s cap. + assert edge.delay_after_predecessor_us == pytest.approx(250_000.0) + + +# --- 3. session identity via x-dynamo-* headers ----------------------------- + + +def test_session_headers_every_turn_final_on_last(tmp_path: Path) -> None: + """Session identity is HEADER-borne (x-dynamo-*), never body nvext: dynamo's + NvExt is deny_unknown_fields and rejects body-level agent_context.""" + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="sess-X"), + _request_end(ts=1100, session_id="sess-X"), + _request_end(ts=1200, session_id="sess-X"), + ], + ) + + pb = from_dynamo_trace(p) + aid = "sess-X" + + def headers(nid: str) -> dict: + node = pb.graph.nodes[nid] + assert "nvext" not in (node.extra_body or {}), ( + "body nvext is rejected by dynamo (deny_unknown_fields)" + ) + return node.extra_headers + + h1, h2, h3 = (headers(f"{aid}:{k}") for k in (0, 1, 2)) + for h in (h1, h2, h3): + assert h["x-dynamo-session-id"] == "sess-X" + assert "x-dynamo-session-final" not in h1 + assert "x-dynamo-session-final" not in h2 + assert h3["x-dynamo-session-final"] == "true" + + +def test_child_session_headers_carry_parent_session_id(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="parent"), + _request_end(ts=1100, session_id="child", parent_session_id="parent"), + _request_end(ts=1200, session_id="parent"), + ], + ) + + pb = from_dynamo_trace(p) + child = pb.graph.nodes["child:0"] + h = child.extra_headers + assert h["x-dynamo-session-id"] == "child" + assert h["x-dynamo-parent-session-id"] == "parent" + # Single-turn child session: its one turn is also the final one. + assert h["x-dynamo-session-final"] == "true" + parent_a1 = pb.graph.nodes["parent:0"] + assert "x-dynamo-parent-session-id" not in parent_a1.extra_headers + # Only the parent emits a trace; the child flattens into the same graph. + assert [t.id for t in pb.traces] == ["parent"] + assert "multi-root" not in pb.traces[0].tags + + +def test_session_headers_reach_store_envelope(tmp_path: Path) -> None: + """The envelope carries extra_headers so the worker can attach them to the + request headers (Turn.extra_headers -> transport merge).""" + from aiperf.dataset.graph.segment_ir.store_builder import ( + _prompt_segment_ids, + _trie_envelope, + ) + + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="s1"), + _request_end(ts=1100, session_id="s1"), + ], + ) + pb = from_dynamo_trace(p) + aid = "s1" + for k, final in ((0, False), (1, True)): + node = pb.graph.nodes[f"{aid}:{k}"] + env = _trie_envelope(node, _prompt_segment_ids(node) or []) + h = env["extra_headers"] + assert h["x-dynamo-session-id"] == "s1" + assert ("x-dynamo-session-final" in h) is final + assert "nvext" not in env["dispatch_overrides"] + + +# --- 4. recorded output pinning ---------------------------------------------- + + +def test_recorded_output_always_pins_native_max_tokens(tmp_path: Path) -> None: + """Recorded output_tokens always pin the native ``LlmNode.max_tokens`` + (weka parity); a recorded 0 upgrades to 1 with a warning (wire_output_cap).""" + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="p1", output_tokens=256), + _request_end(ts=1100, session_id="p1", output_tokens=300), + _request_end(ts=1200, session_id="p1", output_tokens=0), + ], + ) + + pb = from_dynamo_trace(p) + aid = "p1" + assert pb.graph.nodes[f"{aid}:0"].max_tokens == 256 + assert pb.graph.nodes[f"{aid}:1"].max_tokens == 300 + assert pb.graph.nodes[f"{aid}:2"].max_tokens == 1 + + +# --- 5. empty trace file ---------------------------------------------------- + + +def test_empty_trace_file_raises(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + p.write_text("") + + with pytest.raises(EmptyDynamoTraceError): + from_dynamo_trace(p) + + +def test_replay_only_trace_error_names_missing_agent_context(tmp_path: Path) -> None: + """Every record skipped for lacking agent_context must produce an error + that says so, not a misleading "no trace records found".""" + p = tmp_path / "trace.jsonl" + replay_only = _request_end(ts=1000, session_id="ignored") + del replay_only["agent_context"] + replay_only_2 = _request_end(ts=1100, session_id="ignored", request_id="r2") + del replay_only_2["agent_context"] + _write_jsonl(p, [replay_only, replay_only_2]) + + with pytest.raises(EmptyDynamoTraceError) as ei: + from_dynamo_trace(p) + msg = str(ei.value) + assert "2 records had no agent_context" in msg + assert "replay-only" in msg + assert "session identity" in msg + + +# --- 6. virtual-hash fallback tag -------------------------------------------- + + +def test_no_replay_metrics_tags_virtual_fallback(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl(p, [_request_end(ts=1000, session_id="p1")]) + + pb = from_dynamo_trace(p) + + assert "virtual-hash-fallback" in pb.traces[0].tags + + +def test_recorded_replay_metrics_have_no_fallback_tag( + tmp_path: Path, +) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end( + ts=1000, + session_id="p1", + input_tokens=32, + replay={ + "trace_block_size": 16, + "input_length": 32, + "input_sequence_hashes": [11, 22], + }, + ), + ], + ) + + pb = from_dynamo_trace(p) + + assert "virtual-hash-fallback" not in pb.traces[0].tags + + +# --- 7. session_id_filter restricts sessions -------------------------------- + + +def test_session_id_filter_restricts_sessions(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="p1"), + _request_end(ts=1010, session_id="p2"), + _request_end(ts=1100, session_id="p1"), + ], + ) + + pb = from_dynamo_trace(p, session_id_filter="p1") + + assert [t.id for t in pb.traces] == ["p1"] + assert set(pb.graph.nodes) == {"p1:0", "p1:1"} + + +def test_session_id_filter_no_matches_raises(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl(p, [_request_end(ts=1000, session_id="p1")]) + + with pytest.raises(EmptyDynamoTraceError): + from_dynamo_trace(p, session_id_filter="p-OTHER") + + +# --- 8. expected + observed metadata ---------------------------------------- + + +def test_expected_tokens_populated_from_request(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end( + ts=1000, + session_id="p1", + input_tokens=128, + output_tokens=256, + cached_tokens=64, + ), + ], + ) + + pb = from_dynamo_trace(p) + a1 = pb.graph.nodes["p1:0"] + assert isinstance(a1, LlmNode) + assert a1.expected is not None + assert a1.expected.input_tokens == 128 + assert a1.expected.output_tokens == 256 + assert a1.expected.cache_read_tokens == 64 + assert a1.expected.cache_creation_tokens is None + + +def test_node_metadata_carries_only_identity_breadcrumbs(tmp_path: Path) -> None: + """No recorded-scalar round-trip is stamped: node metadata carries only the + dynamo identity breadcrumbs (session/turn, small_prompt) -- everything else + lives on native fields or in the capture file, and metadata survives the + content-free sidecar broadcast.""" + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="p1", kv_hit_rate=0.42), + _request_end(ts=1100, session_id="p1", kv_hit_rate=0.99), + ], + ) + + pb = from_dynamo_trace(p) + a1 = pb.graph.nodes["p1:0"] + assert set(a1.metadata) == {"dynamo", "trie"} + assert set(a1.metadata["dynamo"]) == { + "session_id", + "parent_session_id", + "turn_index", + "small_prompt", + } + + +def test_tool_events_are_recognized_but_not_lowered(tmp_path: Path) -> None: + """tool_start/tool_end/tool_error records parse cleanly and produce no + per-node tool metadata: tool time is implicit in the recorded end-to-start + gaps the replay honors, and no consumer reads a breakdown.""" + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="p1"), + _tool_end(ts=1100, session_id="p1", name="search", duration_ms=80.0), + _tool_end(ts=1150, session_id="p1", name="fetch", duration_ms=20.0), + _request_end(ts=1200, session_id="p1"), + ], + ) + + pb = from_dynamo_trace(p) + aid = "p1" + assert sorted(pb.graph.nodes) == [f"{aid}:0", f"{aid}:1"] + for node in pb.graph.nodes.values(): + assert "tool_breakdown" not in node.metadata["dynamo"] + + +# --- 9. cycle and depth guards ---------------------------------------------- + + +def test_cycle_in_parent_link_raises(tmp_path: Path) -> None: + """Mutually-pointing parent_link records raise DynamoTraceAdapterError.""" + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="A", parent_session_id="B"), + _request_end(ts=1100, session_id="B", parent_session_id="A"), + ], + ) + + with pytest.raises(DynamoTraceAdapterError) as ei: + from_dynamo_trace(p) + assert "cycle" in str(ei.value).lower() + + +def test_depth_overflow_respects_env_setting(tmp_path: Path, monkeypatch) -> None: + """A linear chain longer than AIPERF_DYNAMO_MAX_SUBAGENT_DEPTH raises.""" + # Build a 5-level chain p1 -> p2 -> p3 -> p4 -> p5 + pids = [f"p{i}" for i in range(1, 6)] + records = [] + ts = 1000 + for i, pid in enumerate(pids): + parent = pids[i - 1] if i > 0 else None + records.append(_request_end(ts=ts, session_id=pid, parent_session_id=parent)) + ts += 10 + p = tmp_path / "trace.jsonl" + _write_jsonl(p, records) + + # Patch the adapter module's *local* `Environment` reference instead of + # mutating the shared `Environment.DYNAMO` Pydantic singleton. The latter + # was xdist-flaky: concurrent workers could observe a torn singleton + # mid-monkeypatch. Replacing the module-level name binding is fully + # isolated to this test and still exercises the env-read code path in + # `from_dynamo_trace`. + from types import SimpleNamespace + + from aiperf.dataset.graph.adapters.dynamo import trace as _dt_mod + + # Cap depth at 3 -> chain depth 5 should raise. + monkeypatch.setattr( + _dt_mod, + "Environment", + SimpleNamespace(DYNAMO=SimpleNamespace(MAX_SUBAGENT_DEPTH=3)), + ) + with pytest.raises(DynamoTraceAdapterError) as ei: + from_dynamo_trace(p) + assert "depth" in str(ei.value).lower() + + # Cap depth at 10 -> chain depth 5 should pass. + monkeypatch.setattr( + _dt_mod, + "Environment", + SimpleNamespace(DYNAMO=SimpleNamespace(MAX_SUBAGENT_DEPTH=10)), + ) + pb = from_dynamo_trace(p) + assert isinstance(pb, ParsedGraph) + + +# --- deliberate divergences from dynamo-source behavior ---------------------- + + +def test_self_parent_session_is_root_not_cycle(tmp_path: Path) -> None: + """Dynamo's generic header mapping passes parent==session through verbatim; + it means "no parent", never a cycle.""" + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [_request_end(ts=1000, session_id="s1", parent_session_id="s1")], + ) + pb = from_dynamo_trace(p) + assert pb.traces, "self-parent session must parse as a root" + + +def test_parent_from_first_non_self_record_not_recs0(tmp_path: Path) -> None: + """The parent header may appear only on later calls of a session; the + chain parent must still be picked up (first non-self parent wins).""" + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="root"), + _request_end(ts=1100, session_id="child"), + _request_end(ts=1200, session_id="child", parent_session_id="root"), + ], + ) + pb = from_dynamo_trace(p) + # One root trace: child is linked under root, not emitted as a second root. + assert len(pb.traces) == 1 + assert "multi-root" not in pb.traces[0].tags + + +def test_parent_found_when_earliest_record_is_parentless_tool_event( + tmp_path: Path, +) -> None: + """A harness tool event without the parent header can be the session's + EARLIEST record; the parent from the later request_end must still bind + (chain parent comes from the parent_link map, not any single record).""" + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="root"), + _tool_end(ts=1050, session_id="child"), # no parent header + _request_end(ts=1200, session_id="child", parent_session_id="root"), + ], + ) + pb = from_dynamo_trace(p) + assert len(pb.traces) == 1 + assert pb.traces[0].id == "root" + assert "multi-root" not in pb.traces[0].tags + child = pb.graph.nodes["child:0"] + h = child.extra_headers + assert h["x-dynamo-parent-session-id"] == "root" + + +def test_duplicate_records_across_dir_files_dedup(tmp_path: Path) -> None: + """Dual sinks share one output path: aggregated dirs can duplicate records.""" + rec = _request_end(ts=1000, session_id="s1", request_id="r-dup") + d = tmp_path / "capture" + d.mkdir() + _write_jsonl(d / "a.jsonl", [rec]) + _write_jsonl(d / "b.jsonl", [rec]) + pb = from_dynamo_trace(d) + graph = pb.graph + llm_nodes = [n for n in graph.nodes.values() if isinstance(n, LlmNode)] + assert len(llm_nodes) == 1, "duplicated request_end must lower to ONE node" + + +def test_can_load_sniffs_sink_envelope(tmp_path: Path) -> None: + """Real dynamo captures wrap every line in {"timestamp","event"}.""" + import gzip as _gzip + + p = tmp_path / "trace.000000.jsonl.gz" + wrapped = { + "timestamp": 5, + "event": _request_end(ts=1000, session_id="s1"), + } + with _gzip.open(p, "wb") as f: + f.write(orjson.dumps(wrapped) + b"\n") + assert DynamoTraceAdapter.can_load(tmp_path) is True + assert from_dynamo_trace(tmp_path / "trace").traces + + +@pytest.mark.asyncio +async def test_session_headers_survive_interned_store_round_trip( + tmp_path: Path, +) -> None: + """The INTERNED unified store (what the worker actually reads via + read_node_envelope) must carry extra_headers end to end.""" + from aiperf.dataset.graph.segment_ir.store_builder import ( + build_unified_trie_store_interned, + flat_trie_ordinals, + ) + from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, + ) + from aiperf.graph.worker_materialize import read_node_envelope + + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="s1"), + _request_end(ts=1100, session_id="s1"), + ], + ) + pb = from_dynamo_trace(p) + store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id="bench-hdr" + ) + await build_unified_trie_store_interned(pb, store) + + client = GraphSegmentUnifiedClient( + base_path=tmp_path, benchmark_id="bench-hdr" + ).open() + trace = pb.traces[0] + ordinals = flat_trie_ordinals(pb, trace) + aid = "s1" + env1 = read_node_envelope(client, trace.id, ordinals[f"{aid}:0"], "profiling") + env2 = read_node_envelope(client, trace.id, ordinals[f"{aid}:1"], "profiling") + assert env1 is not None and env2 is not None + assert env1["extra_headers"]["x-dynamo-session-id"] == "s1" + assert "x-dynamo-session-final" not in env1["extra_headers"] + assert env2["extra_headers"]["x-dynamo-session-final"] == "true" + assert "nvext" not in env1["dispatch_overrides"] + + +# --- read-time hash-id interning propagation -------------------------------- + + +def _replay(hashes: list[int], *, bs: int = 16) -> dict: + """Block-aligned replay dict (input_length = len(hashes) * bs).""" + return { + "trace_block_size": bs, + "input_length": len(hashes) * bs, + "input_sequence_hashes": hashes, + } + + +def test_lowered_hash_ids_share_interned_objects(tmp_path: Path) -> None: + """The interned canonical objects survive lowering: after ``_collect_records`` + canonicalizes the record window, ``dynamo_trie_nodes`` (via the ``list()`` + companion) copies the SAME objects into every ``TrieRequest.hash_ids``, so an + equal value shared across turns AND sessions is ONE object across all nodes -- + while ``release_replay=True`` still frees every recorded replay. + """ + from aiperf.common.environment import Environment + from aiperf.dataset.graph.adapters.dynamo.trace import _collect_chains + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + dynamo_trie_nodes, + ) + + a, b, c, d = 2**63 + 1, 2**63 + 2, 2**63 + 3, 2**63 + 4 + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="s1", replay=_replay([a, b])), + _request_end(ts=2000, session_id="s1", replay=_replay([a, b, c])), + _request_end(ts=1500, session_id="s2", replay=_replay([a, d])), + ], + ) + chains = _collect_chains(p, None, max_depth=Environment.DYNAMO.MAX_SUBAGENT_DEPTH) + nodes, _bs, _tags = dynamo_trie_nodes(chains, release_replay=True) + + ids_by_value: dict[int, set[int]] = {} + for n in nodes: + for h in n.request.hash_ids: + ids_by_value.setdefault(h, set()).add(id(h)) + + assert set(ids_by_value) == {a, b, c, d} + # Every distinct value maps to exactly ONE object across all nodes. + assert all(len(ids) == 1 for ids in ids_by_value.values()), ids_by_value + # release_replay still nulled every recorded replay after the copy. + assert all( + t.record.request.replay is None for ch in chains.values() for t in ch.turns + ) + + +def test_intern_does_not_disturb_virtual_hash_fallback(tmp_path: Path) -> None: + """Non-interference: a mixed recorded/virtual session still lowers correctly + through the interning path -- the recorded turn keeps its exact values and the + non-replay turn extends the per-session prefix with fresh negative virtual ids + (values-only assertions; the ``None``-guard skips records with no replay). + """ + from aiperf.common.environment import Environment + from aiperf.dataset.graph.adapters.dynamo.trace import _collect_chains + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + dynamo_trie_nodes, + ) + + a, b = 2**63 + 1, 2**63 + 2 + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(ts=1000, session_id="s1", replay=_replay([a, b])), + # No replay metadata -> virtual-hash fallback at lowering. + _request_end(ts=2000, session_id="s1", input_tokens=64), + ], + ) + chains = _collect_chains(p, None, max_depth=Environment.DYNAMO.MAX_SUBAGENT_DEPTH) + nodes, bs, tags = dynamo_trie_nodes(chains) + + assert "virtual-hash-fallback" in tags + h1, h2 = nodes[0].request.hash_ids, nodes[1].request.hash_ids + assert h1 == [a, b] + assert h2[:2] == [a, b] and len(h2) == 64 // bs + assert all(h < 0 for h in h2[2:]) diff --git a/tests/unit/dataset/graph/adapters/test_dynamo_trace_fidelity.py b/tests/unit/dataset/graph/adapters/test_dynamo_trace_fidelity.py new file mode 100644 index 0000000000..ead155ae46 --- /dev/null +++ b/tests/unit/dataset/graph/adapters/test_dynamo_trace_fidelity.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fidelity-parity coverage: output-token enforcement + idle-warped edge delays. + +The trie IR carries recorded timing on the EDGES (interval-order +``delay_after_predecessor_us`` / START ``min_start_delay_us``), not on +per-node ``pre_wait`` WaitSpecs. Recorded end-to-start gaps always replay, +warped through the SAME ``ActiveIdleWarp`` the weka adapter uses: +``idle_gap_cap_seconds`` defaults to the shared 60s cap, an explicit float +compresses over-long idle gaps to that cap, and ``None`` disables the warp +(raw recorded gaps). +""" + +from __future__ import annotations + +from pathlib import Path + +from aiperf.dataset.graph.adapters.dynamo.trace import from_dynamo_trace +from aiperf.dataset.graph.models import GraphRecord, LlmNode, StaticEdge + +FIXTURES = Path(__file__).resolve().parent / "fixtures" / "dynamo_nested" + +CADENCE_FIXTURE = "nested_2_level.jsonl.gz" + + +def _llm_nodes(graph: GraphRecord) -> list[LlmNode]: + return [n for n in graph.nodes.values() if isinstance(n, LlmNode)] + + +def _static_edges(pb) -> list[StaticEdge]: + return [e for e in pb.graph.edges if isinstance(e, StaticEdge)] + + +# --- recorded output pinning (weka parity, always on) ------------------------- + + +def test_output_cap_always_pinned_to_recorded() -> None: + """Every node pins max_output_tokens to the recorded output_tokens by + default (weka parity); a recorded 0 upgrades to 1 (wire_output_cap).""" + pb = from_dynamo_trace(FIXTURES / "nested_2_level.jsonl.gz") + pinned = 0 + for node in _llm_nodes(pb.graph): + cap = node.max_tokens + recorded = node.expected.output_tokens if node.expected else None + assert cap == (recorded if recorded else 1) + if recorded: + pinned += 1 + assert pinned >= 1 + + +# --- idle-warped recorded timing on edges ------------------------------------ + + +def test_default_replays_recorded_edge_delays() -> None: + """Delays are never silently dropped: the default keeps the recorded gaps.""" + pb = from_dynamo_trace(FIXTURES / CADENCE_FIXTURE) + delays = [e.delay_after_predecessor_us or 0.0 for e in _static_edges(pb)] + assert any(d > 0.0 for d in delays), ( + "the default parse must keep the recorded end-to-start gaps on the " + f"binding edges; got {delays}" + ) + # Arrival stamping rides the same warped clock. + assert any((n.arrival_offset_us or 0) > 0 for n in _llm_nodes(pb.graph)) + + +def test_idle_gap_cap_none_keeps_raw_recorded_gaps() -> None: + """Explicit None disables the warp: identical to a cap above every gap.""" + raw = from_dynamo_trace(FIXTURES / CADENCE_FIXTURE, idle_gap_cap_seconds=None) + uncompressed = from_dynamo_trace( + FIXTURES / CADENCE_FIXTURE, idle_gap_cap_seconds=1e9 + ) + raw_delays = [e.delay_after_predecessor_us or 0.0 for e in _static_edges(raw)] + assert raw_delays == [ + e.delay_after_predecessor_us or 0.0 for e in _static_edges(uncompressed) + ] + assert any(d > 0.0 for d in raw_delays) + + +def test_idle_gap_cap_zero_collapses_recorded_gaps() -> None: + """cap=0.0 collapses every idle gap, so all warped delays/offsets hit 0.""" + pb = from_dynamo_trace( + FIXTURES / CADENCE_FIXTURE, + idle_gap_cap_seconds=0.0, + ) + for e in _static_edges(pb): + assert (e.delay_after_predecessor_us or 0.0) == 0.0 + assert (e.min_start_delay_us or 0.0) == 0.0 + for n in _llm_nodes(pb.graph): + assert (n.arrival_offset_us or 0) == 0 + + +# --- DynamoTraceAdapter.parse knob sources ----------------------------------- + + +def test_parse_forwards_ctx_idle_gap_cap() -> None: + """The run's synthesis.idle_gap_cap_seconds reaches the dynamo build.""" + from aiperf.dataset.graph.adapters.dynamo import trace as dt + from aiperf.dataset.graph.parse_context import GraphParseContext + + pb = dt.DynamoTraceAdapter.parse( + FIXTURES / CADENCE_FIXTURE, + ctx=GraphParseContext(idle_gap_cap_seconds=0.0), + ) + for e in [e for e in pb.graph.edges if isinstance(e, StaticEdge)]: + assert (e.delay_after_predecessor_us or 0.0) == 0.0 + assert (e.min_start_delay_us or 0.0) == 0.0 diff --git a/tests/unit/dataset/graph/adapters/test_dynamo_trace_nested_fixtures.py b/tests/unit/dataset/graph/adapters/test_dynamo_trace_nested_fixtures.py new file mode 100644 index 0000000000..fce69fadba --- /dev/null +++ b/tests/unit/dataset/graph/adapters/test_dynamo_trace_nested_fixtures.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Flat-IR structural coverage for the nested-subagent Dynamo fixtures. + +Every parent/child session flattens into ONE graph of ``{session_id}:{k}`` +``LlmNode``s. Concurrency is EMERGENT from the recorded intervals: a child whose +first request overlaps a still-running parent turn START-ANCHORS to that +parent (a single ``delay_after_predecessor_start_us`` edge, so the child +tracks the parent's dispatch causally), while disjoint intervals yield a +finished-before edge. +""" + +from __future__ import annotations + +from pathlib import Path + +import orjson +import pytest + +from aiperf.dataset.graph.adapters.dynamo.trace import ( + DynamoTraceAdapterError, + from_dynamo_trace, +) +from aiperf.dataset.graph.models import LlmNode, StaticEdge + +FIXTURES = Path(__file__).resolve().parent / "fixtures" / "dynamo_nested" + + +def _node_ids(pb) -> set[str]: + return set(pb.graph.nodes) + + +def _incoming(pb, target: str) -> list[StaticEdge]: + return [ + e for e in pb.graph.edges if isinstance(e, StaticEdge) and e.target == target + ] + + +def test_nested_2_level_flattens_parent_and_child() -> None: + pb = from_dynamo_trace(FIXTURES / "nested_2_level.jsonl.gz") + assert _node_ids(pb) == { + "sess_A:0", + "sess_A:1", + "sess_A:2", + "sess_B:0", + "sess_B:1", + } + assert all(isinstance(n, LlmNode) for n in pb.graph.nodes.values()) + assert [t.id for t in pb.traces] == ["sess_A"] + + +def test_nested_3_level_flattens_all_three_sessions() -> None: + pb = from_dynamo_trace(FIXTURES / "nested_3_level.jsonl.gz") + for sid in ("sess_A", "sess_B", "sess_C"): + assert f"{sid}:0" in _node_ids(pb) + assert [t.id for t in pb.traces] == ["sess_A"] + + +def test_parallel_subagents_flatten_side_by_side() -> None: + pb = from_dynamo_trace(FIXTURES / "parallel_subagents.jsonl.gz") + for sid in ("sess_B", "sess_C"): + assert f"{sid}:0" in _node_ids(pb) + assert f"{sid}:1" in _node_ids(pb) + + +def test_parallel_two_root_parses_one_trace_per_root() -> None: + pb = from_dynamo_trace(FIXTURES / "parallel_two_root.jsonl.gz") + # Two independent roots -> two per-tree traces (multi-graph), no multi-root tag. + assert [t.id for t in pb.traces] == ["sess_A", "sess_B"] + for trace in pb.traces: + assert trace.graph_ref == trace.id + assert "multi-root" not in trace.tags + assert set(pb.graphs) == {"sess_A", "sess_B"} + + +def test_tool_call_id_linkage_fixture_flattens_without_subgraph_nodes() -> None: + pb = from_dynamo_trace(FIXTURES / "tool_call_id_linkage.jsonl.gz") + assert all(isinstance(n, LlmNode) for n in pb.graph.nodes.values()) + assert "sess_B:0" in _node_ids(pb) + + +def test_mixed_turn_tool_events_do_not_lower() -> None: + """Recorded tool events parse cleanly but produce no per-node metadata: + tool time is implicit in the recorded end-to-start gaps the replay honors.""" + pb = from_dynamo_trace(FIXTURES / "mixed_turn.jsonl.gz") + a2 = pb.graph.nodes["sess_A:1"] + assert "tool_breakdown" not in a2.metadata["dynamo"] + + +def test_cycle_AB_A_raises_with_cycle_message() -> None: + with pytest.raises(DynamoTraceAdapterError) as excinfo: + from_dynamo_trace(FIXTURES / "cycle_AB_A.jsonl.gz") + assert "cycle" in str(excinfo.value).lower(), ( + f"DynamoTraceAdapterError message should mention 'cycle'; got: {excinfo.value!s}" + ) + + +# --- emergent concurrency: overlap vs disjoint intervals -------------------- + + +def _rec( + *, + ts: int, + sid: str, + parent: str | None = None, + received: int | None = None, + total: float | None = None, +) -> dict: + ctx: dict = {"session_id": sid} + if parent is not None: + ctx["parent_session_id"] = parent + req: dict = { + "request_id": f"{sid}-{ts}", + "model": "m", + "input_tokens": 32, + "output_tokens": 16, + "cached_tokens": 0, + } + if received is not None: + req["request_received_ms"] = received + if total is not None: + req["total_time_ms"] = total + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": ctx, + "request": req, + } + + +def test_child_overlapping_parent_turn_start_anchors_to_parent( + tmp_path: Path, +) -> None: + """The child starts while the parent's turn is still running -> one + start-anchored edge from the parent (``delay_after_predecessor_start_us`` + set), so the child tracks the parent's dispatch causally instead of + freezing to the recorded wall clock.""" + p = tmp_path / "overlap.jsonl" + records = [ + # Parent turn 1 active [1000, 1400]; child fires at 1150, inside it. + _rec(ts=1400, sid="parent", received=1000, total=400.0), + _rec(ts=1150, sid="child", parent="parent"), + ] + p.write_bytes(b"\n".join(orjson.dumps(r) for r in records)) + + pb = from_dynamo_trace(p) + parent_a1 = "parent:0" + child_a1 = "child:0" + incoming = _incoming(pb, child_a1) + assert [e.source for e in incoming] == [parent_a1], ( + f"overlapping child must start-anchor to the parent; got {incoming}" + ) + (edge,) = incoming + # The start-anchor edge keeps the recorded parent-start-to-child-start gap + # (150ms, below the shared 60s idle-gap cap) while the end-delay stays None. + assert edge.delay_after_predecessor_start_us == pytest.approx(150_000) + assert edge.delay_after_predecessor_us is None + + +def test_child_overlapping_parent_turn_start_anchor_warps_under_cap( + tmp_path: Path, +) -> None: + """Same overlap fixture with the idle-gap cap collapsed to zero: the + start-anchor delay rides the warped clock (every idle gap compressed), while + the anchor kind survives (end-delay stays None).""" + p = tmp_path / "overlap_warp.jsonl" + records = [ + # Parent turn 1 active [1000, 1400]; child fires at 1150, inside it. + _rec(ts=1400, sid="parent", received=1000, total=400.0), + _rec(ts=1150, sid="child", parent="parent"), + ] + p.write_bytes(b"\n".join(orjson.dumps(r) for r in records)) + + pb = from_dynamo_trace(p, idle_gap_cap_seconds=0.0) + parent_a1 = "parent:0" + child_a1 = "child:0" + incoming = _incoming(pb, child_a1) + assert [e.source for e in incoming] == [parent_a1], ( + f"overlapping child must start-anchor to the parent; got {incoming}" + ) + (edge,) = incoming + # Both starts sit inside ONE active interval (no idle gap between them), so + # the start-to-start gap survives even a cap of zero. + assert edge.delay_after_predecessor_start_us == pytest.approx(150_000) + assert edge.delay_after_predecessor_us is None + + +def test_child_after_parent_turn_gets_finished_before_edge(tmp_path: Path) -> None: + """Disjoint intervals: the parent turn finished before the child started, + so the child carries a finished-before edge from it.""" + p = tmp_path / "disjoint.jsonl" + records = [ + # Parent turn 1 active [1000, 1100]; child fires at 1150, after it. + _rec(ts=1100, sid="parent", received=1000, total=100.0), + _rec(ts=1150, sid="child", parent="parent"), + ] + p.write_bytes(b"\n".join(orjson.dumps(r) for r in records)) + + pb = from_dynamo_trace(p) + parent_a1 = "parent:0" + child_a1 = "child:0" + incoming = _incoming(pb, child_a1) + assert [e.source for e in incoming] == [parent_a1] diff --git a/tests/unit/dataset/graph/adapters/test_dynamo_trace_reader.py b/tests/unit/dataset/graph/adapters/test_dynamo_trace_reader.py new file mode 100644 index 0000000000..b404b4c4e7 --- /dev/null +++ b/tests/unit/dataset/graph/adapters/test_dynamo_trace_reader.py @@ -0,0 +1,478 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the Dynamo agent-trace reader.""" + +from __future__ import annotations + +import gzip +from pathlib import Path + +import orjson +import pytest +from pytest import param + +from aiperf.dataset.graph.adapters.dynamo.trace_reader import ( + AgentContext, + AgentReplayMetrics, + AgentTraceRecord, + DynamoTraceReadError, + discover_segments, + iter_raw_records, + iter_trace_records, +) + + +def _request_end( + *, + event_time_unix_ms: int = 1_000, + session_id: str = "sess-1", + request_id: str = "req-1", + model: str = "my-model", + replay: dict | None = None, +) -> dict: + rec: dict = { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": event_time_unix_ms, + "event_source": "dynamo", + "agent_context": { + "session_id": session_id, + }, + "request": {"request_id": request_id, "model": model}, + } + if replay is not None: + rec["request"]["replay"] = replay + return rec + + +def _tool_event( + *, + event_type: str, + event_time_unix_ms: int = 1_500, + session_id: str = "sess-1", + tool_call_id: str = "call-1", + tool_class: str = "web_search", + status: str | None = None, +) -> dict: + rec: dict = { + "schema": "dynamo.request.trace.v1", + "event_type": event_type, + "event_time_unix_ms": event_time_unix_ms, + "event_source": "harness", + "agent_context": { + "session_id": session_id, + }, + "tool": {"tool_call_id": tool_call_id, "tool_class": tool_class}, + } + if status is not None: + rec["tool"]["status"] = status + return rec + + +def _write_jsonl(path: Path, records: list[dict]) -> None: + with path.open("wb") as f: + for r in records: + f.write(orjson.dumps(r)) + f.write(b"\n") + + +def _write_jsonl_gz(path: Path, records: list[dict]) -> None: + with gzip.open(path, "wb") as f: + for r in records: + f.write(orjson.dumps(r)) + f.write(b"\n") + + +def test_iter_plain_jsonl_round_trips_three_records(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(), + _tool_event(event_type="tool_start"), + _tool_event(event_type="tool_end", status="succeeded"), + ], + ) + out = list(iter_trace_records(p)) + assert len(out) == 3 + assert out[0].event_type == "request_end" + assert out[0].request is not None + assert out[0].request.model == "my-model" + assert out[1].event_type == "tool_start" + assert out[1].tool is not None and out[1].tool.tool_call_id == "call-1" + assert out[2].event_type == "tool_end" + assert out[2].tool is not None and out[2].tool.status == "succeeded" + + +def test_iter_gzipped_jsonl_round_trips(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl.gz" + _write_jsonl_gz( + p, + [ + _request_end(), + _tool_event(event_type="tool_start"), + _tool_event(event_type="tool_end", status="succeeded"), + ], + ) + out = list(iter_trace_records(p)) + assert [r.event_type for r in out] == ["request_end", "tool_start", "tool_end"] + + +def test_segmented_directory_iterates_in_segment_order(tmp_path: Path) -> None: + seg_dir = tmp_path / "segdir" + seg_dir.mkdir() + _write_jsonl_gz(seg_dir / "trace.000000.jsonl.gz", [_request_end(request_id="r0")]) + _write_jsonl_gz(seg_dir / "trace.000001.jsonl.gz", [_request_end(request_id="r1")]) + _write_jsonl_gz(seg_dir / "trace.000002.jsonl.gz", [_request_end(request_id="r2")]) + + out_dir = list(iter_trace_records(seg_dir)) + assert [r.request.request_id for r in out_dir if r.request] == ["r0", "r1", "r2"] + + +def test_segmented_prefix_iterates_in_segment_order(tmp_path: Path) -> None: + seg_dir = tmp_path / "segdir" + seg_dir.mkdir() + _write_jsonl_gz(seg_dir / "trace.000000.jsonl.gz", [_request_end(request_id="r0")]) + _write_jsonl_gz(seg_dir / "trace.000001.jsonl.gz", [_request_end(request_id="r1")]) + _write_jsonl_gz(seg_dir / "trace.000002.jsonl.gz", [_request_end(request_id="r2")]) + + prefix = seg_dir / "trace" + segments = discover_segments(prefix) + assert [s.name for s in segments] == [ + "trace.000000.jsonl.gz", + "trace.000001.jsonl.gz", + "trace.000002.jsonl.gz", + ] + out_pref = list(iter_trace_records(prefix)) + assert [r.request.request_id for r in out_pref if r.request] == ["r0", "r1", "r2"] + + +@pytest.mark.parametrize( + "event_filter,expected_count", + [ + param({"request_end"}, 2, id="request_end_only"), + param({"tool_start", "tool_end"}, 2, id="tool_pair"), + param({"tool_error"}, 1, id="tool_error_only"), + ], +) # fmt: skip +def test_filter_by_event_type( + tmp_path: Path, event_filter: set[str], expected_count: int +) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(request_id="a"), + _tool_event(event_type="tool_start"), + _request_end(request_id="b"), + _tool_event(event_type="tool_end", status="succeeded"), + _tool_event(event_type="tool_error", status="error"), + ], + ) + out = list(iter_trace_records(p, event_types=event_filter)) + assert len(out) == expected_count + assert {r.event_type for r in out} <= event_filter + + +def test_filter_by_session_id_keeps_subset(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(session_id="sess-A", request_id="a1"), + _request_end(session_id="sess-B", request_id="b1"), + _request_end(session_id="sess-A", request_id="a2"), + _request_end(session_id="sess-B", request_id="b2"), + ], + ) + out = list(iter_trace_records(p, session_id="sess-A")) + assert [r.request.request_id for r in out if r.request] == ["a1", "a2"] + + +def test_filter_by_subagent_session_id_keeps_subset(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(session_id="planner", request_id="p1"), + _request_end(session_id="researcher", request_id="r1"), + _request_end(session_id="planner", request_id="p2"), + ], + ) + out = list(iter_trace_records(p, session_id="planner")) + assert [r.request.request_id for r in out if r.request] == ["p1", "p2"] + + +def test_filter_by_time_range_inclusive(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end(event_time_unix_ms=500, request_id="t500"), + _request_end(event_time_unix_ms=1000, request_id="t1000"), + _request_end(event_time_unix_ms=2000, request_id="t2000"), + _request_end(event_time_unix_ms=3000, request_id="t3000"), + _request_end(event_time_unix_ms=4000, request_id="t4000"), + ], + ) + out = list(iter_trace_records(p, time_range_ms=(1000, 3000))) + assert [r.request.request_id for r in out if r.request] == [ + "t1000", + "t2000", + "t3000", + ] + + +def test_schema_mismatch_raises(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + bad = _request_end() + bad["schema"] = "other.schema.v1" + _write_jsonl(p, [bad]) + with pytest.raises(DynamoTraceReadError): + list(iter_trace_records(p)) + + +def test_replay_metrics_round_trip(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + big_hash = 14_879_255_164_371_896_291 # > 2^63, valid u64 from Dynamo + _write_jsonl( + p, + [ + _request_end( + replay={ + "trace_block_size": 64, + "input_length": 128, + "input_sequence_hashes": [1, 2, 3, big_hash], + } + ) + ], + ) + out = list(iter_trace_records(p)) + assert len(out) == 1 + assert out[0].request is not None + replay = out[0].request.replay + assert replay is not None + assert replay.trace_block_size == 64 + assert replay.input_length == 128 + assert replay.input_sequence_hashes == [1, 2, 3, big_hash] + + +def test_negative_replay_hash_rejected(tmp_path: Path) -> None: + """A negative recorded hash collides with the virtual negative-id namespace.""" + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _request_end( + replay={ + "trace_block_size": 16, + "input_length": 32, + "input_sequence_hashes": [1, -2, 3], + } + ) + ], + ) + with pytest.raises(DynamoTraceReadError): + list(iter_trace_records(p)) + + # Model-level rejection is direct (pydantic ValidationError before the reader + # wraps it) so callers building AgentReplayMetrics get the same guard. + with pytest.raises(ValueError, match="non-negative"): + AgentReplayMetrics( + trace_block_size=16, + input_length=32, + input_sequence_hashes=[1, -2, 3], + ) + + +def test_schema_alias_round_trip_via_model_dump() -> None: + rec = _request_end() + parsed = AgentTraceRecord.model_validate(rec) + dumped = parsed.model_dump(by_alias=True, exclude_none=True) + assert dumped["schema"] == "dynamo.request.trace.v1" + assert "schema_" not in dumped + + +def test_directory_with_no_trace_files_raises(tmp_path: Path) -> None: + empty = tmp_path / "empty" + empty.mkdir() + with pytest.raises(DynamoTraceReadError): + discover_segments(empty) + + +def test_segmented_prefix_no_match_raises(tmp_path: Path) -> None: + seg_dir = tmp_path / "segdir" + seg_dir.mkdir() + with pytest.raises(DynamoTraceReadError): + discover_segments(seg_dir / "missing-prefix") + + +def test_invalid_json_line_raises(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + p.write_bytes(b'{"schema": "dynamo.request.trace.v1"}\nnot-json\n') + with pytest.raises(DynamoTraceReadError): + list(iter_raw_records(p)) + + +def test_blank_lines_are_skipped(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + payload = ( + orjson.dumps(_request_end()) + + b"\n\n \n" + + orjson.dumps(_request_end(request_id="r2")) + + b"\n" + ) + p.write_bytes(payload) + out = list(iter_trace_records(p)) + assert len(out) == 2 + + +def test_agent_context_required_fields_missing_raises(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + rec = _request_end() + del rec["agent_context"]["session_id"] + _write_jsonl(p, [rec]) + with pytest.raises(DynamoTraceReadError): + list(iter_trace_records(p)) + + +def test_agent_context_empty_string_passes_consumer_side(tmp_path: Path) -> None: + """Consumer-side parity note: Dynamo's deserialize_non_empty_string is server-side + only. Empty strings pass here — this test pins that behavior.""" + p = tmp_path / "trace.jsonl" + rec = _request_end() + rec["agent_context"]["session_id"] = "" + _write_jsonl(p, [rec]) + out = list(iter_trace_records(p)) + assert len(out) == 1 + assert out[0].agent_context.session_id == "" + + +def test_extra_fields_ignored() -> None: + rec = _request_end() + rec["future_field"] = {"new": "thing"} + rec["request"]["another_future"] = 42 + parsed = AgentTraceRecord.model_validate(rec) + assert not hasattr(parsed, "future_field") + assert parsed.request is not None + + +def test_agent_context_parent_session_id_optional() -> None: + ac = AgentContext(session_id="z") + assert ac.parent_session_id is None + ac2 = AgentContext( + session_id="z", + parent_session_id="parent", + ) + assert ac2.parent_session_id == "parent" + + +# --- dynamo file-sink envelope + discovery parity ----------------------------- + + +def _wrap(rec: dict, ts: int = 12) -> dict: + """Wrap a record the way both dynamo file sinks do (telemetry/jsonl_gz.rs).""" + return {"timestamp": ts, "event": rec} + + +def test_sink_envelope_unwrapped_plain_jsonl(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl(p, [_wrap(_request_end(request_id="req-w1"))]) + out = list(iter_trace_records(p)) + assert len(out) == 1 + assert out[0].request is not None and out[0].request.request_id == "req-w1" + + +def test_sink_envelope_unwrapped_gz_segments(tmp_path: Path) -> None: + _write_jsonl_gz( + tmp_path / "trace.000000.jsonl.gz", + [_wrap(_request_end(event_time_unix_ms=1_000, request_id="a"))], + ) + _write_jsonl_gz( + tmp_path / "trace.000001.jsonl.gz", + [_wrap(_request_end(event_time_unix_ms=2_000, request_id="b"), ts=99)], + ) + out = list(iter_trace_records(tmp_path / "trace")) + assert [r.request.request_id for r in out] == ["a", "b"] + + +def test_bare_and_wrapped_records_mix(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, [_request_end(request_id="bare"), _wrap(_request_end(request_id="wrapped"))] + ) + assert [r.request.request_id for r in iter_trace_records(p)] == ["bare", "wrapped"] + + +def test_prefix_with_output_path_suffix_is_stripped(tmp_path: Path) -> None: + # DYN_REQUEST_TRACE_OUTPUT_PATH=/x/trace.jsonl.gz produces /x/trace.000000.jsonl.gz; + # passing the configured path verbatim must resolve the segments. + _write_jsonl_gz(tmp_path / "trace.000000.jsonl.gz", [_request_end()]) + for given in ("trace.jsonl.gz", "trace.jsonl", "trace"): + segs = discover_segments(tmp_path / given) + assert [s.name for s in segs] == ["trace.000000.jsonl.gz"] + + +def test_directory_segment_order_is_numeric_past_six_digits(tmp_path: Path) -> None: + _write_jsonl_gz( + tmp_path / "t.1000000.jsonl.gz", [_request_end(event_time_unix_ms=3_000)] + ) + _write_jsonl_gz( + tmp_path / "t.999999.jsonl.gz", [_request_end(event_time_unix_ms=2_000)] + ) + _write_jsonl_gz( + tmp_path / "t.000001.jsonl.gz", [_request_end(event_time_unix_ms=1_000)] + ) + segs = discover_segments(tmp_path) + assert [s.name for s in segs] == [ + "t.000001.jsonl.gz", + "t.999999.jsonl.gz", + "t.1000000.jsonl.gz", + ] + + +def test_truncated_final_gzip_member_raises_read_error(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl.gz" + _write_jsonl_gz(p, [_request_end()]) + data = p.read_bytes() + p.write_bytes(data[: len(data) - 8]) # chop the gzip trailer mid-member + with pytest.raises(DynamoTraceReadError, match="truncated or corrupt gzip"): + list(iter_raw_records(p)) + + +def test_trace_block_size_zero_rejected() -> None: + with pytest.raises(ValueError, match="trace_block_size"): + AgentReplayMetrics( + trace_block_size=0, input_length=4, input_sequence_hashes=[1] + ) + + +# valid 10-byte gzip header followed by bytes that are not a deflate stream: +# reading raises zlib.error (not EOFError / BadGzipFile). +_GZ_HEADER_PLUS_GARBAGE = ( + b"\x1f\x8b\x08\x00\x00\x00\x00\x00\x00\x03" + b"this-is-not-deflate-data" +) + + +def test_corrupt_gzip_deflate_raises_read_error(tmp_path: Path) -> None: + """zlib.error from corrupt deflate data must wrap into DynamoTraceReadError.""" + p = tmp_path / "trace.jsonl.gz" + p.write_bytes(_GZ_HEADER_PLUS_GARBAGE) + with pytest.raises(DynamoTraceReadError, match="truncated or corrupt gzip"): + list(iter_raw_records(p)) + + +def test_non_utf8_jsonl_raises_read_error(tmp_path: Path) -> None: + """Binary bytes behind a .jsonl name must wrap into DynamoTraceReadError.""" + p = tmp_path / "trace.jsonl" + p.write_bytes(b"\xff\xfe\x00garbage-bytes") + with pytest.raises(DynamoTraceReadError, match="not valid UTF-8"): + list(iter_raw_records(p)) + + +def test_uppercase_gz_suffix_opens_as_gzip(tmp_path: Path) -> None: + """Detection lowercases suffixes; the segment opener must match.""" + p = tmp_path / "trace.JSONL.GZ" + _write_jsonl_gz(p, [_request_end(request_id="upper")]) + out = list(iter_trace_records(p)) + assert [r.request.request_id for r in out if r.request] == ["upper"] diff --git a/tests/unit/dataset/graph/adapters/test_dynamo_tree_build.py b/tests/unit/dataset/graph/adapters/test_dynamo_tree_build.py new file mode 100644 index 0000000000..6195fad1d5 --- /dev/null +++ b/tests/unit/dataset/graph/adapters/test_dynamo_tree_build.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-session-tree dynamo build: grouping, cross-parent edge drop, byte-identity. + +Dynamo graphs are built per session-tree (a root session plus every descendant +linked by ``agent_context.parent_trajectory_id``; ``trajectory_id == +session_id`` in real captures). Building each tree independently drops +cross-parent interval-order edges by construction while preserving within-tree +edges (parent<->subagent + intra-session). The differential oracle is the OLD +single global build, reached by calling the per-tree seam +(``_build_graph_from_chains``) with the WHOLE chain set. +""" + +from __future__ import annotations + +from pathlib import Path + +import orjson + +from aiperf.common.environment import Environment +from aiperf.dataset.graph.adapters.dynamo.trace import ( + _build_graph_from_chains, + _Chain, + _collect_chains, + from_dynamo_trace, + group_chains_into_trees, +) +from aiperf.dataset.graph.adapters.shared.content import resolve_effective_root_seed +from aiperf.dataset.graph.adapters.shared.idle_gap import DEFAULT_IDLE_GAP_CAP_SECONDS +from aiperf.dataset.graph.segment_ir.pool import SegmentPool + +_MAX_DEPTH = Environment.DYNAMO.MAX_SUBAGENT_DEPTH + +# --- fixture builders (REAL agent_context field shape) -------------------- + + +def _ctx(*, sid: str, parent: str | None, use_trajectory: bool) -> dict: + """Agent context in the real capture shape (``trajectory_id == session_id``). + + ``use_trajectory=True`` mirrors real captures: ``trajectory_id`` is set to + the session id and a subagent's ``parent_trajectory_id`` names the parent's + trajectory id (== parent session id). ``use_trajectory=False`` exercises the + legacy ``parent_session_id`` fallback (older / hand-authored traces). + """ + ctx: dict = {"session_id": sid} + if use_trajectory: + ctx["trajectory_id"] = sid + if parent is not None: + ctx["parent_trajectory_id"] = parent + elif parent is not None: + ctx["parent_session_id"] = parent + return ctx + + +def _re( + *, + ts: int, + sid: str, + hashes: list[int], + ilen: int, + parent: str | None = None, + otok: int = 8, + bs: int = 16, + use_trajectory: bool = True, +) -> dict: + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": _ctx(sid=sid, parent=parent, use_trajectory=use_trajectory), + "request": { + "request_id": f"r-{sid}-{ts}", + "model": "m", + "input_tokens": ilen, + "output_tokens": otok, + "cached_tokens": 0, + "ttft_ms": 10.0, + "replay": { + "trace_block_size": bs, + "input_length": ilen, + "input_sequence_hashes": hashes, + }, + }, + } + + +def _write_jsonl(path: Path, records: list[dict]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("wb") as f: + for r in records: + f.write(orjson.dumps(r)) + f.write(b"\n") + + +def _two_independent_trees() -> list[dict]: + """Two parentless roots with DISJOINT hashes and DISJOINT intervals. + + A single global build finished-before-links p1's last turn to p2's first + (p1 all ends by 1100, p2 starts at 5000); a per-tree build must not. + """ + return [ + _re(ts=1000, sid="p1", hashes=[1, 2], ilen=32), + _re(ts=1100, sid="p1", hashes=[1, 2, 3], ilen=48), + _re(ts=5000, sid="p2", hashes=[4, 5], ilen=32), + _re(ts=5100, sid="p2", hashes=[4, 5, 6], ilen=48), + ] + + +def _oracle_global_build(chains: dict[str, _Chain], *, seed: int) -> tuple: + """The OLD single global build: the per-tree seam over the WHOLE chain set.""" + pool = SegmentPool() + graph, tags = _build_graph_from_chains( + chains, + pool=pool, + content_root_seed=resolve_effective_root_seed(seed), + block_size=16, + idle_gap_cap_seconds=DEFAULT_IDLE_GAP_CAP_SECONDS, + content_tokenizer=None, + prompt_corpus="coding", + release_replay=False, + ) + return graph, pool, tags + + +def _agent_prefix(node_id: str) -> str: + """The session id (drops the ``:{turn}`` suffix from ``{session_id}:{k}``).""" + return node_id.rsplit(":", 1)[0] + + +def _nonterminal_edges(pb) -> list[tuple[str, str]]: + return [ + (e.source, e.target) + for e in pb.graph.edges + if e.source not in ("START", "END") and e.target not in ("START", "END") + ] + + +# --- 1. group_chains_into_trees (pure) ------------------------------------ + + +def _mk_chains(links: dict[str, str | None]) -> dict[str, _Chain]: + return { + sid: _Chain(sid, parent_session_id=parent, turns=[]) + for sid, parent in links.items() + } + + +def _plink(chains: dict[str, _Chain]) -> dict[str, str]: + return { + sid: c.parent_session_id + for sid, c in chains.items() + if c.parent_session_id is not None + } + + +def test_group_parent_with_two_subagents_forms_one_tree() -> None: + chains = _mk_chains({"p": None, "s1": "p", "s2": "p"}) + trees = group_chains_into_trees(chains, _plink(chains)) + assert len(trees) == 1 + assert set(trees[0]) == {"p", "s1", "s2"} + + +def test_group_nested_subagent_chain_forms_one_tree() -> None: + # p -> s1 -> s2 (grandchild): the fixpoint walk unions all three. + chains = _mk_chains({"p": None, "s1": "p", "s2": "s1"}) + trees = group_chains_into_trees(chains, _plink(chains)) + assert len(trees) == 1 + assert set(trees[0]) == {"p", "s1", "s2"} + + +def test_group_two_unrelated_roots_form_two_trees() -> None: + chains = _mk_chains({"a": None, "b": None}) + trees = group_chains_into_trees(chains, _plink(chains)) + assert [sorted(t) for t in trees] == [["a"], ["b"]] + + +def test_group_external_parent_is_its_own_root() -> None: + # 'a' points at 'x' which is NOT in the chain set -> 'a' is a forest root. + chains = _mk_chains({"a": "x", "b": None}) + trees = group_chains_into_trees(chains, {"a": "x"}) + assert [sorted(t) for t in trees] == [["a"], ["b"]] + + +def test_group_cycle_guard_terminates_and_covers_all() -> None: + # A -> B -> A: _guard_chain_forest rejects this upstream; the grouping guard + # must still terminate and partition every session exactly once. + chains = _mk_chains({"a": "b", "b": "a"}) + trees = group_chains_into_trees(chains, {"a": "b", "b": "a"}) + covered = sorted(sid for t in trees for sid in t) + assert covered == ["a", "b"] + + +# --- 2. end-to-end linkage from the REAL field shape ---------------------- + + +def test_parent_trajectory_id_links_subagent_into_one_tree(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _re(ts=1000, sid="P", hashes=[1, 2], ilen=32), + _re(ts=2000, sid="S", parent="P", hashes=[7, 8], ilen=32), + _re(ts=3000, sid="P", hashes=[1, 2, 3], ilen=48), + ], + ) + chains = _collect_chains(p, None, max_depth=_MAX_DEPTH) + # Linkage came from parent_trajectory_id (parent_session_id was never set). + assert chains["S"].parent_session_id == "P" + trees = group_chains_into_trees(chains, _plink(chains)) + assert len(trees) == 1 + assert set(trees[0]) == {"P", "S"} + + +def test_two_unrelated_trajectories_form_two_trees(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _re(ts=1000, sid="a", hashes=[1, 2], ilen=32), + _re(ts=2000, sid="b", hashes=[4, 5], ilen=32), + ], + ) + chains = _collect_chains(p, None, max_depth=_MAX_DEPTH) + trees = group_chains_into_trees(chains, _plink(chains)) + assert [sorted(t) for t in trees] == [["a"], ["b"]] + + +def test_parent_session_id_fallback_when_no_trajectory(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _re(ts=1000, sid="P", hashes=[1, 2], ilen=32, use_trajectory=False), + _re( + ts=2000, + sid="S", + parent="P", + hashes=[7, 8], + ilen=32, + use_trajectory=False, + ), + ], + ) + chains = _collect_chains(p, None, max_depth=_MAX_DEPTH) + assert chains["S"].parent_session_id == "P" + trees = group_chains_into_trees(chains, _plink(chains)) + assert len(trees) == 1 + assert set(trees[0]) == {"P", "S"} + + +# --- 3. cross-parent edge drop + within-tree edge preservation ------------ + + +def test_independent_trees_have_zero_cross_session_edges(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl(p, _two_independent_trees()) + pb = from_dynamo_trace(p, content_root_seed=7) + + for s, t in _nonterminal_edges(pb): + assert _agent_prefix(s) == _agent_prefix(t), f"cross-session edge {s} -> {t}" + + # The specific p1-finished-before-p2 edge a global build emits is absent. + edges = {(e.source, e.target) for e in pb.graph.edges} + assert ("p1:1", "p2:0") not in edges + + +def test_global_build_carries_the_cross_edge_tree_build_drops(tmp_path: Path) -> None: + # Differential oracle: the OLD single global build DOES emit the cross-parent + # edge, proving the tree build actively drops it (not that it never existed). + p = tmp_path / "trace.jsonl" + _write_jsonl(p, _two_independent_trees()) + chains = _collect_chains(p, None, max_depth=_MAX_DEPTH) + graph, _pool, _tags = _oracle_global_build(chains, seed=7) + edges = {(e.source, e.target) for e in graph.edges} + assert ("p1:1", "p2:0") in edges + + +def test_parent_subagent_within_tree_cross_session_edge_present( + tmp_path: Path, +) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + _re(ts=1000, sid="P", hashes=[1, 2], ilen=32), + _re(ts=2000, sid="S", parent="P", hashes=[7, 8], ilen=32), + _re(ts=3000, sid="P", hashes=[1, 2, 3], ilen=48), + ], + ) + pb = from_dynamo_trace(p, content_root_seed=7) + edges = {(e.source, e.target) for e in pb.graph.edges} + + cross = [ + (s, t) + for (s, t) in _nonterminal_edges(pb) + if _agent_prefix(s) != _agent_prefix(t) + ] + assert cross, "within-tree parent<->subagent edge expected" + # P's first turn finished (1000) before S started (2000): interval edge. + assert ("P:0", "S:0") in edges + + +# --- 4. pool byte-identity vs the global build + determinism -------------- + + +def test_tree_scoped_pool_identical_to_per_tree_seam_union(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl(p, _two_independent_trees()) + seed = 12345 + + pb = from_dynamo_trace(p, content_root_seed=seed) + + # Reference: build each independent tree ALONE through the same per-tree seam, + # then union. Trace-scoped response/tail seeds make each tree's content scope + # on its OWN root, so the reference is this per-tree-scoped union -- NOT the + # single global build (which scopes every tree by one root, giving different + # response bytes for every non-root tree). Hashes are disjoint, so the union + # is a plain by_id merge with no cross-tree collision. + chains = _collect_chains(p, None, max_depth=_MAX_DEPTH) + trees = group_chains_into_trees(chains, _plink(chains)) + union: dict = {} + for tree in trees: + tree_chains = {sid: chains[sid] for sid in tree} + _graph, tree_pool, _tags = _oracle_global_build(tree_chains, seed=seed) + union.update(tree_pool.by_id) + + assert pb.segment_pool is not None + assert len(union) > 0 + assert pb.segment_pool.by_id == union + + +def test_tree_scoped_build_is_deterministic(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + _two_independent_trees() + + [_re(ts=2000, sid="S", parent="p1", hashes=[1, 2, 9], ilen=48)], + ) + + a = from_dynamo_trace(p, content_root_seed=999) + b = from_dynamo_trace(p, content_root_seed=999) + + assert a.segment_pool is not None and b.segment_pool is not None + assert set(a.graph.nodes) == set(b.graph.nodes) + assert set(a.segment_pool.by_id) == set(b.segment_pool.by_id) diff --git a/tests/unit/dataset/graph/adapters/test_dynamo_trie_ir.py b/tests/unit/dataset/graph/adapters/test_dynamo_trie_ir.py new file mode 100644 index 0000000000..3d50c8e9b9 --- /dev/null +++ b/tests/unit/dataset/graph/adapters/test_dynamo_trie_ir.py @@ -0,0 +1,257 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the dynamo -> shared-LCP-segment-trie content lowering. + +THE regression this file guards: a multi-turn session's later turns must +materialize at the block-aligned COVERED token count of their recorded hashes +(turn 2 of a 2-turn extending kv trace = 64 tokens at bs=16), never the old +channel-replay inflation (history + full per-turn reconstruction = 144). +Shared-prefix identity (same leading segment ids across extending turns) and +seed-pinned determinism are asserted alongside. +""" + +from __future__ import annotations + +from pathlib import Path + +import orjson + +from aiperf.dataset.graph.adapters.dynamo.trace import from_dynamo_trace +from aiperf.dataset.graph.models import LlmNode +from aiperf.dataset.graph.segment_ir.store_builder import ( + _prompt_segment_ids, + _trie_llm_nodes, +) + + +def _rec( + *, + ts: int, + sid: str, + input_tokens: int, + output_tokens: int, + hashes: list[int] | None = None, + block_size: int = 16, + input_length: int | None = None, +) -> dict: + """One current-schema ``dynamo.request.trace.v1`` ``request_end`` record.""" + req: dict = { + "request_id": f"r{ts}", + "model": "m", + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "cached_tokens": 0, + } + if hashes is not None: + req["replay"] = { + "trace_block_size": block_size, + "input_length": input_length if input_length is not None else input_tokens, + "input_sequence_hashes": hashes, + } + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": {"session_id": sid}, + "request": req, + } + + +def _dyn_fixture(tmp_path: Path) -> Path: + """Write a current-schema linear 2-turn single-session dynamo trace. + + Turn 2's replay hashes EXTEND turn 1's (``[111,222]`` -> ``[111,222,333,444]``) + so the accumulated turn-2 prompt shares turn 1's leading messages. + """ + p = tmp_path / "dyn_recorded.jsonl" + records = [ + _rec( + ts=1000, + sid="s1", + input_tokens=32, + output_tokens=8, + hashes=[111, 222], + input_length=32, + ), + _rec( + ts=2000, + sid="s1", + input_tokens=64, + output_tokens=12, + hashes=[111, 222, 333, 444], + input_length=64, + ), + ] + p.write_bytes(b"\n".join(orjson.dumps(r) for r in records)) + return p + + +def _dyn_fixture_no_hashes(tmp_path: Path) -> Path: + """Same 2 extending turns, but with NO replay metadata (virtual-hash path).""" + p = tmp_path / "dyn_virtual.jsonl" + records = [ + _rec(ts=1000, sid="s1", input_tokens=32, output_tokens=8), + _rec(ts=2000, sid="s1", input_tokens=64, output_tokens=12), + ] + p.write_bytes(b"\n".join(orjson.dumps(r) for r in records)) + return p + + +def _builtin_encode(root_seed: int): + """The builtin-tokenizer ``encode`` of the SAME cached synthesizer the parse used.""" + from aiperf.dataset.graph.adapters.shared.content import ( + get_or_build_synthesizer, + ) + + synth = get_or_build_synthesizer( + "builtin", prompt_corpus="coding", root_seed=root_seed + ) + return synth._pg.tokenizer.encode + + +def _materialized_token_count(pool, path: list[str], encode) -> int: + return sum(len(encode(m["content"])) for m in pool.materialize(path)) + + +def _arrival_ordered_llm_nodes(parsed) -> list[tuple[str, LlmNode]]: + return sorted( + ((nid, n) for nid, n in parsed.graph.nodes.items() if isinstance(n, LlmNode)), + key=lambda kv: (kv[1].arrival_offset_us or 0, kv[0]), + ) + + +def test_turn2_isl_is_covered_count_not_inflated(tmp_path): + """THE regression: 2-turn extending chain must materialize turn 2 at 64 + tokens (covered count == input_length, bs=16), not ~144 (history + full + reconstruction double-count).""" + p = _dyn_fixture(tmp_path) + parsed = from_dynamo_trace(p, content_root_seed=1234, content_tokenizer="builtin") + pool = parsed.segment_pool + assert pool is not None + nodes = _arrival_ordered_llm_nodes(parsed) + assert len(nodes) == 2 + + encode = _builtin_encode(1234) + p1 = _prompt_segment_ids(nodes[0][1]) + p2 = _prompt_segment_ids(nodes[1][1]) + assert p1 and p2 + assert _materialized_token_count(pool, p1, encode) == 32 + assert _materialized_token_count(pool, p2, encode) == 64 + assert p2[: len(p1)] == p1 # prefix identity survives + + +def test_virtual_mode_isl_is_covered_count(tmp_path): + """Virtual (no recorded hashes) path: the same covered-count contract holds. + + Seed 1234 is pinned to a corpus window where the builtin decode->encode + round trip is block-exact (a handful of windows merge one token across a + message boundary, e.g. seed 7 counts 31/63); the covered-count sizing + under test is seed-independent. + """ + p = _dyn_fixture_no_hashes(tmp_path) + parsed = from_dynamo_trace( + p, + content_root_seed=1234, + content_tokenizer="builtin", + ) + pool = parsed.segment_pool + assert pool is not None + nodes = _arrival_ordered_llm_nodes(parsed) + assert len(nodes) == 2 + + encode = _builtin_encode(1234) + p1 = _prompt_segment_ids(nodes[0][1]) + p2 = _prompt_segment_ids(nodes[1][1]) + assert p1 and p2 + assert _materialized_token_count(pool, p1, encode) == 32 + assert _materialized_token_count(pool, p2, encode) == 64 + assert p2[: len(p1)] == p1 + + +def test_trie_parse_is_deterministic_under_pinned_seed(tmp_path): + """Parsing twice with the same seed yields identical paths + pool bytes.""" + p = _dyn_fixture(tmp_path) + parsed_a = from_dynamo_trace(p, content_root_seed=7, content_tokenizer="builtin") + parsed_b = from_dynamo_trace(p, content_root_seed=7, content_tokenizer="builtin") + node_a = list(_trie_llm_nodes(parsed_a, parsed_a.traces[0]).values())[-1] + node_b = list(_trie_llm_nodes(parsed_b, parsed_b.traces[0]).values())[-1] + path_a = _prompt_segment_ids(node_a) + path_b = _prompt_segment_ids(node_b) + assert path_a and path_a == path_b + assert parsed_a.segment_pool.materialize( + path_a + ) == parsed_b.segment_pool.materialize(path_b) + + +def test_every_node_materializes_from_the_pool(tmp_path): + """Every trie node stamps a prompt_segment_ids path the pool can materialize.""" + p = _dyn_fixture(tmp_path) + parsed = from_dynamo_trace(p, content_root_seed=1234) + pool = parsed.segment_pool + assert pool is not None + for trace in parsed.traces: + nodes = _trie_llm_nodes(parsed, trace) + assert nodes + for nid, node in nodes.items(): + path = _prompt_segment_ids(node) + assert path, f"{nid} has no prompt_segment_ids" + msgs = pool.materialize(path) + assert msgs and all(set(m) == {"role", "content"} for m in msgs) + + +def test_post_strip_sidecar_carries_no_recorded_replay_hashes(tmp_path): + """The content-free graph_meta sidecar must not embed the recorded + input_sequence_hashes: strip_replay_text clears only prompt and + metadata["trie"], so the replay payload must never ride node metadata in + the first place (hash consumption happens at lowering time via + TrieRequest.hash_ids; no recorded-scalar round-trip is stamped).""" + from aiperf.dataset.graph.codecs import encode_graph_meta_sidecar + from aiperf.dataset.graph.graph_meta_sidecar import strip_replay_text + + p = _dyn_fixture(tmp_path) + parsed = from_dynamo_trace(p, content_root_seed=1234, content_tokenizer="builtin") + stripped = strip_replay_text(parsed) + for nid, node in stripped.graph.nodes.items(): + assert set(node.metadata) == {"dynamo", "trie"}, nid + blob = encode_graph_meta_sidecar(stripped, source_fingerprint={"kind": "test"}) + assert b"input_sequence_hashes" not in blob + + +def test_recorded_hashes_are_content_global_across_files(tmp_path): + """Dynamo hash scope is deliberately GLOBAL, unlike weka's ``local`` scope. + + Recorded ``input_sequence_hashes`` are chained sequence hashes over the + actual prompt tokens -- equal hash means equal upstream content by + construction -- so the same hashes in two DIFFERENT trace files (e.g. two + captures of the same session) must synthesize identical bytes. Contrast + ``test_weka_trie_hash_scope.py``, where weka's ``hash_id_scope: "local"`` + requires the opposite (weka's ``"global"`` scope matches dynamo). + """ + pools = [] + for name in ("capture-a", "capture-b"): + d = tmp_path / name + d.mkdir() + p = d / "trace.jsonl" + records = [ + _rec( + ts=1000, + sid=f"session-{name}", + input_tokens=32, + output_tokens=8, + hashes=[111, 222], + input_length=32, + ), + ] + p.write_bytes(b"\n".join(orjson.dumps(r) for r in records)) + parsed = from_dynamo_trace( + p, + content_root_seed=1234, + content_tokenizer="builtin", + ) + node = list(_trie_llm_nodes(parsed, parsed.traces[0]).values())[0] + pools.append(parsed.segment_pool.materialize(_prompt_segment_ids(node))) + + assert pools[0] == [ + {"role": m["role"], "content": m["content"]} for m in pools[1] + ], "equal recorded hashes must synthesize identical bytes across files" diff --git a/tests/unit/dataset/graph/adapters/test_dynamo_trie_lowering.py b/tests/unit/dataset/graph/adapters/test_dynamo_trie_lowering.py new file mode 100644 index 0000000000..481239c6e6 --- /dev/null +++ b/tests/unit/dataset/graph/adapters/test_dynamo_trie_lowering.py @@ -0,0 +1,570 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the dynamo -> shared-trie normalization layer.""" + +from __future__ import annotations + +import pytest + +from aiperf.dataset.graph.adapters.dynamo.trace import _Chain, _Turn +from aiperf.dataset.graph.adapters.dynamo.trace_reader import AgentTraceRecord + + +def _rec_obj( + ts, sid, itok, otok, hashes=None, bs=16, ilen=None, received=None, total=None +): + req = { + "request_id": f"r{ts}", + "model": "m", + "input_tokens": itok, + "output_tokens": otok, + "cached_tokens": 0, + } + if received is not None: + req["request_received_ms"] = received + if total is not None: + req["total_time_ms"] = total + if hashes is not None: + req["replay"] = { + "trace_block_size": bs, + "input_length": ilen if ilen is not None else itok, + "input_sequence_hashes": hashes, + } + return AgentTraceRecord.model_validate( + { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": {"session_id": sid}, + "request": req, + } + ) + + +def _chain(sid, turns, parent=None): + return _Chain( + sid, + parent_session_id=parent, + turns=[_Turn(record=r) for r in turns], + ) + + +def test_dynamo_recon_callbacks_offset_cache_smoke_and_shape(): + """``dynamo_recon_callbacks`` decodes through the offset cache and returns + tokens byte-identical to a hand-built list-cache ``_decode_block_tokens`` at + the same block size; the offset cache stores plain int offsets keyed by hash + id (one int per unique hash, not the decoded block list).""" + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + dynamo_recon_callbacks, + ) + from aiperf.dataset.graph.adapters.shared.content import ( + CorpusContentSynthesizer, + get_or_build_synthesizer, + ) + + CorpusContentSynthesizer.reset_worker_cache() + bs = 16 + hash_ids = [111, 222, 333, 111, 222] # repeats exercise the hit path + + cb = dynamo_recon_callbacks( + "builtin", "coding", 1234, block_size=bs, trace_scope="t" + ) + got = cb.decode_block_tokens(hash_ids) + + # Hand-built oracle: the SAME shared synth via the list-cache path. + synth = get_or_build_synthesizer("builtin", prompt_corpus="coding", root_seed=1234) + expected = synth._decode_block_tokens(hash_ids, block_size=bs, cache={}) + assert got == expected + assert len(got) == len(hash_ids) * bs + + offsets: dict[int, int] = {} + synth._decode_block_tokens_offset_cached( + hash_ids, block_size=bs, offset_cache=offsets + ) + assert offsets and all(type(v) is int for v in offsets.values()) + assert set(offsets) == set(hash_ids) + + +def test_recorded_hashes_used_when_present_and_relative_seconds(): + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + dynamo_trie_nodes, + ) + + chains = { + "s1": _chain( + "s1", + [ + _rec_obj( + 2000, "s1", 32, 8, hashes=[111, 222], received=1000, total=1000 + ), + _rec_obj( + 4000, + "s1", + 64, + 12, + hashes=[111, 222, 333, 444], + received=3000, + total=1000, + ), + ], + ) + } + nodes, bs, _tags = dynamo_trie_nodes(chains) + assert bs == 16 + assert [n.request.hash_ids for n in nodes] == [[111, 222], [111, 222, 333, 444]] + assert nodes[0].request.t == 0.0 and nodes[1].request.t == 2.0 + assert nodes[0].request.api_time == 1.0 + + +def test_virtual_hashes_extend_per_session_and_never_collide_with_recorded(): + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + dynamo_trie_nodes, + ) + + chains = { + "s1": _chain( + "s1", + [ + _rec_obj(1000, "s1", 32, 8), + _rec_obj(2000, "s1", 64, 8), + ], + ) + } + nodes, bs, _ = dynamo_trie_nodes(chains) + h1, h2 = nodes[0].request.hash_ids, nodes[1].request.hash_ids + assert len(h1) == 32 // bs and len(h2) == 64 // bs + assert h2[: len(h1)] == h1 + assert all(h < 0 for h in h1 + h2) + + +def test_mixed_block_size_raises(): + from aiperf.dataset.graph.adapters.dynamo.trace import ( + DynamoTraceAdapterError, + ) + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + dynamo_trie_nodes, + ) + + chains = { + "s1": _chain( + "s1", + [ + _rec_obj(1000, "s1", 32, 8, hashes=[1, 2], bs=16), + _rec_obj(2000, "s1", 64, 8, hashes=[1, 2, 3, 4], bs=32), + ], + ) + } + with pytest.raises(DynamoTraceAdapterError, match="trace_block_size"): + dynamo_trie_nodes(chains) + + +def test_misaligned_input_length_raises(): + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + DynamoISLMismatchError, + dynamo_trie_nodes, + ) + + chains = { + "s1": _chain("s1", [_rec_obj(1000, "s1", 100, 8, hashes=[1, 2], ilen=100)]) + } + with pytest.raises(DynamoISLMismatchError): + dynamo_trie_nodes(chains) + + +def test_kv_fallback_turn_gets_virtual_hashes_and_tag(): + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + dynamo_trie_nodes, + ) + + chains = { + "s1": _chain( + "s1", + [ + _rec_obj(1000, "s1", 32, 8, hashes=[111, 222]), + _rec_obj(2000, "s1", 64, 8), # no replay metadata + ], + ) + } + nodes, _bs, tags = dynamo_trie_nodes(chains) + assert "virtual-hash-fallback" in tags + h2 = nodes[1].request.hash_ids + assert h2[:2] == [111, 222] and all(h < 0 for h in h2[2:]) + + +def test_node_ids_and_session_header_metadata(): + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + dynamo_trie_nodes, + ) + + chains = { + "root": _chain( + "root", [_rec_obj(1000, "root", 32, 8), _rec_obj(2000, "root", 64, 8)] + ), + "child": _chain("child", [_rec_obj(1500, "child", 16, 4)], parent="root"), + } + nodes, _bs, _t = dynamo_trie_nodes(chains) + ids = {n.node_id for n in nodes} + assert "root:0" in ids and "child:0" in ids + root_a1 = next(n for n in nodes if n.node_id == "root:0") + h1 = root_a1.dynamo_meta["extra_headers"] + assert h1["x-dynamo-session-id"] == "root" + assert "x-dynamo-session-final" not in h1 + root_a2 = next(n for n in nodes if n.node_id == "root:1") + assert root_a2.dynamo_meta["extra_headers"]["x-dynamo-session-final"] == "true" + child = next(n for n in nodes if n.node_id == "child:0") + ch = child.dynamo_meta["extra_headers"] + assert ch["x-dynamo-parent-session-id"] == "root" + assert ch["x-dynamo-session-final"] == "true" # single-turn session + + +def test_session_id_with_colon_raises(): + """Node ids are ``{session_id}:{k}`` with the session id VERBATIM, so a + recorded session id containing the reserved ``:`` separator (or empty) + cannot form a node id and must fail loud at lowering.""" + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + dynamo_trie_nodes, + ) + + chains = {"a:b": _chain("a:b", [_rec_obj(1000, "a:b", 32, 8)])} + with pytest.raises(ValueError, match="cannot form node ids"): + dynamo_trie_nodes(chains) + + +def test_same_ms_turn_tiebreak_is_numeric_past_index_nine(): + """12 turns at the same start_ms must order :0.. :11, not the lexicographic + :0, :1, :10, :11, :2, ... (order is the trie's content-parent ground truth).""" + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + dynamo_trie_nodes, + ) + + n_turns = 12 + chains = {"s1": _chain("s1", [_rec_obj(1000, "s1", 32, 8) for _ in range(n_turns)])} + nodes, _bs, _tags = dynamo_trie_nodes(chains) + assert [n.node_id for n in nodes] == [f"s1:{k}" for k in range(n_turns)] + assert [n.order for n in nodes] == list(range(n_turns)) + + +def test_virtual_hashes_shrinking_input_reuses_prefix_without_new_ids(): + """prev[:m] truncation branch: a turn with FEWER covered blocks than the + session's prior turn reuses the leading virtual ids and allocates none.""" + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + dynamo_trie_nodes, + ) + + chains = { + "s1": _chain( + "s1", + [ + _rec_obj(1000, "s1", 64, 8), # 4 covered blocks at bs=16 + _rec_obj(2000, "s1", 32, 8), # shrinks to 2 covered blocks + _rec_obj(3000, "s1", 64, 8), # grows again past the truncation + ], + ) + } + nodes, bs, _ = dynamo_trie_nodes(chains) + h1, h2, h3 = (n.request.hash_ids for n in nodes) + assert len(h1) == 64 // bs and len(h2) == 32 // bs + assert h2 == h1[: len(h2)] # pure truncation, no fresh ids + # Regrowth extends the TRUNCATED prefix with fresh ids (the counter never + # rewinds), so the shared prefix survives but the tail is new. + assert h3[: len(h2)] == h2 + assert set(h3[len(h2) :]).isdisjoint(set(h1)) + assert all(h < 0 for h in h1 + h2 + h3) + + +def _built_nodes(chains): + """Run the shared trie build with cheap stub callbacks; return (nodes, pool, result).""" + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + dynamo_trie_nodes, + ) + from aiperf.dataset.graph.segment_ir.pool import SegmentPool + from aiperf.dataset.graph.segment_ir.trie_content import ( + ReconCallbacks, + build_trie_ir, + ) + + nodes, bs, _tags = dynamo_trie_nodes(chains) + pool = SegmentPool() + result = build_trie_ir( + nodes, + block_size=bs, + callbacks=ReconCallbacks( + # One block decodes to EXACTLY bs tokens (the ISL gate checks the + # actual assembled token count, so a mis-sized stub aborts the build). + decode_block_tokens=lambda hash_ids: [ + t for h in hash_ids for t in [abs(h) % 9973] * bs + ], + sample_partial_tail_tokens=lambda n, seed: list(range(n)), + decode_tokens_to_text=lambda toks: " ".join(map(str, toks)), + ), + pool=pool, + idle_gap_cap_seconds=None, + small_prompt_fallback=True, + ) + return nodes, pool, result + + +def test_build_dynamo_llm_node_overrides_metadata_and_stamp(): + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + build_dynamo_llm_node, + ) + + chains = { + "root": _chain( + "root", + [ + _rec_obj(1000, "root", 32, 8, received=1000, total=500), + _rec_obj(2000, "root", 64, 12, received=2000, total=500), + ], + ) + } + nodes, pool, result = _built_nodes(chains) + node = nodes[1] + llm = build_dynamo_llm_node( + node, + build=result.builds[node.node_id], + ) + assert llm.output == f"{node.node_id}_out" + assert llm.streaming is False + # arrival stamps the warped start assigned by the trie build. + assert llm.arrival_offset_us == int(round(node.start * 1_000_000.0)) + # Body params ride the native node fields (no extra_body); + # session identity is header-borne, never body nvext. + assert llm.extra_body is None + headers = llm.extra_headers + assert headers["x-dynamo-session-id"] == "root" + assert headers["x-dynamo-session-final"] == "true" + assert llm.model == "m" + assert llm.max_tokens == 12 + # expected tokens mirror the recorded request. + assert llm.expected is not None and llm.expected.output_tokens == 12 + # metadata shape: observed round-trip + dynamo identity + trie stamp. + assert "observed" not in llm.metadata + dyn = llm.metadata["dynamo"] + assert dyn["session_id"] == "root" + assert dyn["parent_session_id"] is None + assert dyn["turn_index"] == 1 # 0-based; nodes[1] is the second turn + assert "tool_breakdown" not in dyn + assert dyn["small_prompt"] is False + trie = llm.metadata["trie"] + assert trie["prompt_segment_ids"] == result.builds[node.node_id].prompt_path + # Dynamo stamps ONLY prompt_segment_ids -- the build-synthesis response_id / + # hash_ids extras are dropped (nothing on the build, sidecar, store, or + # dispatch plane reads them; weka is the sole remaining extras author). + assert set(trie) == {"prompt_segment_ids"} + assert pool.materialize(trie["prompt_segment_ids"]) + + +def test_build_dynamo_llm_node_omits_inline_prompt(): + """The trie-route LlmNode carries NO inline prompt: content lives only in the + segment pool, addressed via ``metadata["trie"]["prompt_segment_ids"]``. + + ``node.prompt`` is dead weight on this route (the store, sidecar, and worker + all go through the segment pool), so the adapter stamps ``prompt=[]`` and the + pool path still materializes the real conversation content.""" + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + build_dynamo_llm_node, + ) + + chains = { + "root": _chain( + "root", + [ + _rec_obj(1000, "root", 32, 8, received=1000, total=500), + _rec_obj(2000, "root", 64, 12, received=2000, total=500), + ], + ) + } + nodes, pool, result = _built_nodes(chains) + node = nodes[1] + llm = build_dynamo_llm_node( + node, + build=result.builds[node.node_id], + ) + assert llm.prompt == [] + # Content is still reachable through the segment pool via the trie envelope. + seg_ids = llm.metadata["trie"]["prompt_segment_ids"] + messages = pool.materialize(seg_ids) + assert messages, "the pool path must still materialize the real prompt" + for msg in messages: + assert msg["role"] in {"system", "user", "assistant", "tool"} + assert isinstance(msg["content"], str) + + +def test_node_meta_carries_no_recorded_scalar_round_trip(): + """No recorded-scalar round-trip rides the node meta: node metadata + survives the graph_meta sidecar strip (which clears only prompt and + metadata["trie"]), so every extra key would bloat the content-free + structural plane. Hashes are consumed via TrieRequest.hash_ids; scalars + live on the native fields (model / max_tokens / expected).""" + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + dynamo_trie_nodes, + ) + + chains = {"s1": _chain("s1", [_rec_obj(1000, "s1", 32, 8, hashes=[111, 222])])} + nodes, _bs, _tags = dynamo_trie_nodes(chains) + assert set(nodes[0].dynamo_meta) == { + "extra_headers", + "session_id", + "parent_session_id", + "turn_index", + "expected", + } + assert nodes[0].request.hash_ids == [111, 222] + + +def test_build_dynamo_llm_node_zero_output_upgrades_cap_to_one(): + """A recorded output_tokens == 0 must NOT reach the wire as + max_output_tokens=0 (a meaningless cap); it upgrades to 1 with a warning + (wire_output_cap) so the turn stays pinned.""" + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + build_dynamo_llm_node, + ) + + chains = {"s1": _chain("s1", [_rec_obj(1000, "s1", 32, 0)])} + nodes, _pool, result = _built_nodes(chains) + llm = build_dynamo_llm_node( + nodes[0], + build=result.builds[nodes[0].node_id], + ) + assert llm.max_tokens == 1 + + +def _replay_release_chains(): + """Two recorded-replay sessions, freshly built per call so a release + mutation in one lowering never leaks into another test's chains.""" + return { + "s1": _chain( + "s1", + [ + _rec_obj(1000, "s1", 32, 8, hashes=[111, 222]), + _rec_obj(2000, "s1", 64, 12, hashes=[111, 222, 333, 444]), + ], + ), + "s2": _chain("s2", [_rec_obj(3000, "s2", 48, 8, hashes=[555, 666, 777])]), + } + + +def _all_records(chains): + return [t.record for c in chains.values() for t in c.turns] + + +def test_dynamo_trie_nodes_release_replay_off_keeps_recorded_replay(): + """Default (release_replay=False): every record's replay stays intact -- the + re-lowering safety pin. A second lowering of the same chains must still read + recorded hashes, never degrade to the virtual-hash fallback.""" + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + dynamo_trie_nodes, + ) + + chains = _replay_release_chains() + dynamo_trie_nodes(chains) + assert all(r.request.replay is not None for r in _all_records(chains)) + + +def test_dynamo_trie_nodes_release_replay_frees_records_without_changing_nodes(): + """release_replay=True nulls each recorded req.replay AFTER its hashes are + copied into the TrieRequest, so the emitted nodes are identical (same + hashes, ISL, order, causal parents, block size, tags) to the default.""" + from aiperf.dataset.graph.adapters.dynamo.trie_lowering import ( + dynamo_trie_nodes, + ) + + keep = _replay_release_chains() + release = _replay_release_chains() + nodes_keep, bs_keep, tags_keep = dynamo_trie_nodes(keep) + nodes_rel, bs_rel, tags_rel = dynamo_trie_nodes(release, release_replay=True) + + assert all(r.request.replay is None for r in _all_records(release)) + assert all(r.request.replay is not None for r in _all_records(keep)) + + def _ident(nodes): + return [ + ( + n.node_id, + n.request.hash_ids, + n.request.input_length, + n.order, + n.causal_parent_id, + ) + for n in nodes + ] + + assert _ident(nodes_rel) == _ident(nodes_keep) + assert bs_rel == bs_keep and tags_rel == tags_keep + + +def _parse_route_sids(path, *, direct_store): + """Parse a capture and return ``(parsed, [every prompt_segment_id str, in order])``. + + Reads each ``LlmNode``'s ``read_prompt_segment_ids`` path across the whole + graph; the returned strings preserve their object identity so a caller can + compare ``id()``-set cardinality against distinct VALUE cardinality. + """ + from aiperf.dataset.graph.adapters.dynamo.trace import from_dynamo_trace + from aiperf.dataset.graph.models import LlmNode + from aiperf.dataset.graph.segment_ir.envelope import read_prompt_segment_ids + + parsed = from_dynamo_trace( + path, + content_root_seed=1234, + content_tokenizer="builtin", + direct_store=direct_store, + ) + sids: list[str] = [] + for node in parsed.graph.nodes.values(): + if isinstance(node, LlmNode): + sids.extend(read_prompt_segment_ids(node) or []) + return parsed, sids + + +@pytest.mark.parametrize("route", ["eager", "direct"]) # fmt: skip +def test_prompt_segment_ids_are_canonical_objects_on_both_routes(route, tmp_path): + """Every equal ``prompt_segment_ids`` string across all nodes is ONE object. + + Each turn re-lists its whole message chain, so absent interning the same + segment value appears as ~``H/N`` distinct fresh hexdigest strings per node. + Pool interning collapses those to one canonical object per + unique value on BOTH dynamo routes -- so the count of distinct ``id()``s + equals the count of distinct VALUES. On the eager route each stamped sid is + additionally the very object the ``SegmentPool`` first-born for that value. + """ + from tests.harness.dynamo_synth_corpus import write_synthetic_dynamo_capture + + capture = tmp_path / "route.jsonl" + write_synthetic_dynamo_capture( + capture, + sessions=3, + turns_per_session=6, + new_blocks_per_turn=2, + block_size=16, + seed=7, + ) + + if route == "eager": + parsed, sids = _parse_route_sids(capture, direct_store=None) + else: + from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + ) + + store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id="route-direct" + ) + try: + parsed, sids = _parse_route_sids(capture, direct_store=store) + finally: + store.abort() + + assert sids, "the capture produced no prompt_segment_ids" + # Duplicates across turns/nodes exist (each turn re-lists its chain)... + assert len(sids) > len(set(sids)), "test corpus too small to exercise dedup" + # ...but every equal-valued sid is the SAME object: distinct objects == values. + assert len({id(s) for s in sids}) == len(set(sids)) + + if route == "eager": + by_id = parsed.segment_pool.by_id + assert all(s is by_id[s].id for s in sids), ( + "eager route must stamp the first-born canonical Segment.id object" + ) diff --git a/tests/unit/dataset/graph/adapters/test_dynamo_validate_gate.py b/tests/unit/dataset/graph/adapters/test_dynamo_validate_gate.py new file mode 100644 index 0000000000..b9647f991d --- /dev/null +++ b/tests/unit/dataset/graph/adapters/test_dynamo_validate_gate.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dynamo adapter output locked through the full graph validator. + +Defeats the adapter-tests-skip-validator trap (the weka adapter locks its +output the same way, ``test_start_anchor_weka_stamping``): structural +regressions in the trie lowering -- dangling rule-56 edge endpoints, rule-1 +cycles, rule-55 anchor-shape violations -- would otherwise pass a suite that +only asserts node-level fields. Every ``from_dynamo_trace`` shape (linear +recorded, virtual-hash fallback, nested/subagent, parallel, multi-root) must +produce ZERO ERROR-severity issues from :func:`validate`. WARNING-severity +issues are allowed: dynamo stamps ``expected.cache_read_tokens`` from the +recorded ``cached_tokens``, which rule-15 flags as informational. +""" + +from __future__ import annotations + +from pathlib import Path + +import orjson +import pytest +from pytest import param + +from aiperf.dataset.graph.adapters.dynamo.trace import from_dynamo_trace +from aiperf.dataset.graph.models import ParsedGraph +from aiperf.dataset.graph.validator import ( + ValidationIssue, + ValidationSeverity, + validate, +) + +FIXTURES = Path(__file__).resolve().parent / "fixtures" / "dynamo_nested" + + +def _blocking(parsed: ParsedGraph) -> list[ValidationIssue]: + return [i for i in validate(parsed) if i.severity is ValidationSeverity.ERROR] + + +def _rec( + *, + ts: int, + sid: str, + input_tokens: int, + output_tokens: int, + hashes: list[int] | None = None, +) -> dict: + req: dict = { + "request_id": f"r{ts}", + "model": "m", + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "cached_tokens": 0, + } + if hashes is not None: + req["replay"] = { + "trace_block_size": 16, + "input_length": input_tokens, + "input_sequence_hashes": hashes, + } + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": {"session_id": sid}, + "request": req, + } + + +def _write(tmp_path: Path, name: str, records: list[dict]) -> Path: + p = tmp_path / name + p.write_bytes(b"\n".join(orjson.dumps(r) for r in records)) + return p + + +@pytest.mark.parametrize( + "fixture_name", + [ + param("nested_2_level.jsonl.gz", id="nested_2_level"), + param("nested_3_level.jsonl.gz", id="nested_3_level"), + param("parallel_subagents.jsonl.gz", id="parallel_subagents"), + param("parallel_two_root.jsonl.gz", id="parallel_two_root_multi_root"), + param("mixed_turn.jsonl.gz", id="mixed_turn_tool_events"), + param("tool_call_id_linkage.jsonl.gz", id="tool_call_id_linkage"), + ], +) # fmt: skip +def test_nested_fixture_output_has_no_validator_errors(fixture_name: str) -> None: + parsed = from_dynamo_trace(FIXTURES / fixture_name) + blocking = _blocking(parsed) + assert blocking == [], blocking + + +def test_linear_recorded_trace_output_has_no_validator_errors(tmp_path: Path) -> None: + """Trie lowering on recorded replay hashes (turn 2 extends turn 1).""" + p = _write( + tmp_path, + "linear_recorded.jsonl", + [ + _rec( + ts=1000, sid="s1", input_tokens=32, output_tokens=8, hashes=[111, 222] + ), + _rec( + ts=2000, + sid="s1", + input_tokens=64, + output_tokens=12, + hashes=[111, 222, 333, 444], + ), + ], + ) + parsed = from_dynamo_trace(p) + blocking = _blocking(parsed) + assert blocking == [], blocking + + +def test_linear_virtual_fallback_output_has_no_validator_errors( + tmp_path: Path, +) -> None: + """Trie lowering on the virtual-hash fallback (no replay metadata).""" + p = _write( + tmp_path, + "linear_virtual.jsonl", + [ + _rec(ts=1000, sid="s1", input_tokens=32, output_tokens=8), + _rec(ts=2000, sid="s1", input_tokens=64, output_tokens=12), + ], + ) + parsed = from_dynamo_trace(p) + assert "virtual-hash-fallback" in parsed.traces[0].tags + blocking = _blocking(parsed) + assert blocking == [], blocking + + +def test_replayed_delays_output_has_no_validator_errors(tmp_path: Path) -> None: + """The default replay keeps recorded edge delays; rule-54/57 must still pass.""" + p = _write( + tmp_path, + "linear_cadence.jsonl", + [ + _rec( + ts=1000, sid="s1", input_tokens=32, output_tokens=8, hashes=[111, 222] + ), + _rec( + ts=2000, + sid="s1", + input_tokens=64, + output_tokens=12, + hashes=[111, 222, 333, 444], + ), + ], + ) + parsed = from_dynamo_trace(p) + blocking = _blocking(parsed) + assert blocking == [], blocking diff --git a/tests/unit/dataset/graph/adapters/test_peak_context.py b/tests/unit/dataset/graph/adapters/test_peak_context.py new file mode 100644 index 0000000000..69d7f835df --- /dev/null +++ b/tests/unit/dataset/graph/adapters/test_peak_context.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the per-trace peak-context helpers (weka + dynamo).""" + +from __future__ import annotations + +from aiperf.dataset.graph.adapters.dynamo.trace_reader import ( + AgentReplayMetrics, + AgentRequestMetrics, + AgentTraceRecord, +) +from aiperf.dataset.graph.adapters.shared.peak_context import ( + dynamo_tree_peak_context, + weka_trace_peak_context, +) +from aiperf.dataset.graph.adapters.weka.trace_models import ( + WekaNormalRequest, + WekaSubagentEntry, + WekaTrace, +) + + +def _weka_trace_parent_and_subagent() -> WekaTrace: + """Top-level parent (in=100/out=50) plus a subagent child (in=200/out=400).""" + return WekaTrace( + id="trace-1", + models=["m"], + block_size=16, + hash_id_scope="local", + requests=[ + WekaNormalRequest( + t=0.0, type="n", model="m", input_length=100, output_length=50 + ), + WekaSubagentEntry( + t=1.0, + type="subagent", + agent_id="agent_001", + subagent_type="Explore", + status="completed", + models=["m"], + requests=[ + WekaNormalRequest( + t=1.5, + type="n", + model="m", + input_length=200, + output_length=400, + ), + ], + ), + ], + ) + + +def test_weka_peak_context_uncapped_uses_raw_output_lengths() -> None: + trace = _weka_trace_parent_and_subagent() + # max(100+50, 200+400) == 600. + assert weka_trace_peak_context(trace, max_osl=None) == 600 + + +def test_weka_peak_context_max_osl_caps_top_level_only() -> None: + trace = _weka_trace_parent_and_subagent() + # Parent leg capped to 100+min(50,10)=110, subagent child stays 200+400=600. + # If the child were also capped it would drop to 200+10=210, so a peak of + # 600 proves the subagent body is left uncapped. + assert weka_trace_peak_context(trace, max_osl=10) == 600 + + +def test_weka_peak_context_max_osl_reduces_dominant_top_level_leg() -> None: + trace = WekaTrace( + id="trace-2", + models=["m"], + block_size=16, + hash_id_scope="local", + requests=[ + WekaNormalRequest( + t=0.0, type="n", model="m", input_length=100, output_length=50 + ), + ], + ) + assert weka_trace_peak_context(trace, max_osl=None) == 150 + # Only leg is top-level, so max_osl=10 caps it to 100+10. + assert weka_trace_peak_context(trace, max_osl=10) == 110 + + +def _dynamo_record(input_length: int, output_tokens: int) -> AgentTraceRecord: + return AgentTraceRecord( + schema="dynamo.request.trace.v1", + event_type="request_end", + event_time_unix_ms=0, + request=AgentRequestMetrics( + request_id="req-1", + output_tokens=output_tokens, + replay=AgentReplayMetrics( + trace_block_size=16, + input_length=input_length, + input_sequence_hashes=[1, 2, 3], + ), + ), + ) + + +def test_dynamo_peak_context_uses_replay_input_length_plus_output_tokens() -> None: + records = [_dynamo_record(input_length=28832, output_tokens=174)] + # 28832 + 174 == 29006. + assert dynamo_tree_peak_context(records) == 29006 + + +def test_dynamo_peak_context_peaks_over_all_records() -> None: + records = [ + _dynamo_record(input_length=28832, output_tokens=174), + _dynamo_record(input_length=100, output_tokens=10), + ] + assert dynamo_tree_peak_context(records) == 29006 + + +def test_dynamo_peak_context_falls_back_to_input_tokens_then_one() -> None: + with_input_tokens = AgentTraceRecord( + schema="dynamo.request.trace.v1", + event_type="request_end", + event_time_unix_ms=0, + request=AgentRequestMetrics( + request_id="req-2", input_tokens=42, output_tokens=8 + ), + ) + assert dynamo_tree_peak_context([with_input_tokens]) == 50 + + request_free = AgentTraceRecord( + schema="dynamo.request.trace.v1", + event_type="tool_end", + event_time_unix_ms=0, + ) + # No request: input_length falls back to 1, output_tokens to 0. + assert dynamo_tree_peak_context([request_free]) == 1 diff --git a/tests/unit/dataset/graph/adapters/test_selection_filter_then_cap.py b/tests/unit/dataset/graph/adapters/test_selection_filter_then_cap.py new file mode 100644 index 0000000000..56fa29f121 --- /dev/null +++ b/tests/unit/dataset/graph/adapters/test_selection_filter_then_cap.py @@ -0,0 +1,327 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The ai-dynamo/aiperf#1106 regression: schema-only filter-then-cap selection. + +Both recorded-trace loaders (weka + dynamo) must honor ``--num-dataset-entries`` +and ``--max-context-length`` at LOAD time: reject traces whose peak context +exceeds the ceiling, then build only the first N ELIGIBLE traces. The historic +bug built every trace and cloned to fill lanes. + +Fixture shape (both adapters): 20 traces, 8 of which exceed the chosen +``max_context_length``. The 8 over-limit traces sort FIRST in each adapter's +deterministic scan order (weka: dir file name; dynamo: root session id), so a +``num_dataset_entries=10`` load that (correctly) filters-then-caps builds exactly +10 eligible traces after scanning all 8 rejects -- while the BUGGY cap-then-filter +would take the first 10 (8 rejects + 2 eligible) and build only 2. +""" + +from __future__ import annotations + +from pathlib import Path + +import orjson +import pytest + +from aiperf.common.environment import Environment +from aiperf.dataset.graph.adapters.dynamo.trace import from_dynamo_trace +from aiperf.dataset.graph.adapters.shared.content import CorpusContentSynthesizer +from aiperf.dataset.graph.adapters.shared.selection import SelectionStats +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.models import LlmNode + +_MAX_CTX = 1000 +_NUM_ENTRIES = 10 +_OVER_LIMIT_COUNT = 8 +_TOTAL = 20 +# Peak of an over-limit trace (2000 + 10) and an eligible one (100 + 10). +_OVER_INPUT = 2000 +_UNDER_INPUT = 100 +_OUTPUT = 10 + + +@pytest.fixture(autouse=True) +def _fresh_synth_cache(): + """Isolate the process-level synthesizer cache from other tests.""" + CorpusContentSynthesizer.reset_worker_cache() + yield + CorpusContentSynthesizer.reset_worker_cache() + + +@pytest.fixture(autouse=True) +def _force_serial(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep both adapters' builds serial + in-process for deterministic speed.""" + monkeypatch.setattr(Environment.DATASET, "WEKA_GRAPH_PARALLEL_THRESHOLD", 1000) + monkeypatch.setattr(Environment.DATASET, "DYNAMO_GRAPH_PARALLEL_THRESHOLD", 1000) + + +def _is_over_limit(index: int) -> bool: + """The first ``_OVER_LIMIT_COUNT`` traces (in scan order) are over-limit.""" + return index < _OVER_LIMIT_COUNT + + +# --- weka ----------------------------------------------------------------- + + +def _weka_trace_dict(index: int) -> dict: + """One valid weka trace whose sole top-level request is over/under the cap.""" + input_length = _OVER_INPUT if _is_over_limit(index) else _UNDER_INPUT + return { + "id": f"trace-{index:02d}", + "models": ["m"], + "block_size": 64, + "hash_id_scope": "local", + "requests": [ + { + "t": 0.0, + "type": "n", + "model": "m", + "in": input_length, + "out": _OUTPUT, + "hash_ids": [1, 2], + } + ], + } + + +def _write_weka_dir(root: Path) -> Path: + """20 weka trace files named ``t00..t19`` (over-limit ones sort first).""" + root.mkdir(parents=True, exist_ok=True) + for index in range(_TOTAL): + (root / f"t{index:02d}.json").write_bytes(orjson.dumps(_weka_trace_dict(index))) + return root + + +def test_weka_filter_then_cap_builds_exactly_n_eligible(tmp_path: Path) -> None: + root = _write_weka_dir(tmp_path / "weka") + stats_out: list[SelectionStats] = [] + + parsed = from_weka_trace( + root, + num_dataset_entries=_NUM_ENTRIES, + max_context_length=_MAX_CTX, + selection_out=stats_out, + ) + + # Exactly N eligible built -- NOT N minus the rejects that fell in the prefix. + assert len(parsed.traces) == _NUM_ENTRIES + assert len(stats_out) == 1 + stats = stats_out[0] + assert stats.loaded == _NUM_ENTRIES + assert stats.eligible == _NUM_ENTRIES + assert stats.rejected_by_maxctx == _OVER_LIMIT_COUNT + assert stats.largest_observed == _OVER_INPUT + _OUTPUT + # Every built trace is an eligible (under-limit) one. + built_ids = {t.id for t in parsed.traces} + assert all(int(tid.split("-")[1]) >= _OVER_LIMIT_COUNT for tid in built_ids) + + +def test_weka_both_none_builds_everything(tmp_path: Path) -> None: + root = _write_weka_dir(tmp_path / "weka") + stats_out: list[SelectionStats] = [] + + parsed = from_weka_trace(root, selection_out=stats_out) + + assert len(parsed.traces) == _TOTAL + assert stats_out == [] # no selection performed when both knobs are unset + + +# --- dynamo --------------------------------------------------------------- + + +def _dynamo_record(session_id: str, input_tokens: int) -> dict: + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": 1000, + "event_source": "dynamo", + "agent_context": {"session_id": session_id, "trajectory_id": session_id}, + "request": { + "request_id": f"r-{session_id}", + "model": "m", + "input_tokens": input_tokens, + "output_tokens": _OUTPUT, + "ttft_ms": 10.0, + }, + } + + +def _write_dynamo_trace(path: Path) -> Path: + """20 independent single-turn root sessions (= 20 trees), rejects sort first.""" + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("wb") as f: + for index in range(_TOTAL): + input_tokens = _OVER_INPUT if _is_over_limit(index) else _UNDER_INPUT + f.write(orjson.dumps(_dynamo_record(f"s{index:02d}", input_tokens))) + f.write(b"\n") + return path + + +_BLOCK_SIZE = 16 + + +def _aligned_hashes(input_length: int, *, base: int) -> list[int]: + """Block-aligned hash list for ``input_length`` (satisfies _assert_block_aligned). + + Dynamo records full-block hashes plus one partial tail, so a consistent + record has ``ceil(input_length / block_size)`` hashes; distinct ``base`` per + session keeps content-addressed segments from aliasing across trees. + """ + n = max(1, -(-input_length // _BLOCK_SIZE)) + return [base + i for i in range(n)] + + +def _dynamo_replay_record( + session_id: str, *, input_length: int, input_tokens: int, base: int +) -> dict: + """A request_end carrying a real ``replay`` block (input_length + hashes). + + ``input_length`` (the field the peak helpers must read) is set INDEPENDENTLY + of ``input_tokens`` so a peak computed off the wrong field would diverge. + """ + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": 1000, + "event_source": "dynamo", + "agent_context": {"session_id": session_id, "trajectory_id": session_id}, + "request": { + "request_id": f"r-{session_id}", + "model": "m", + "input_tokens": input_tokens, + "output_tokens": _OUTPUT, + "ttft_ms": 10.0, + "replay": { + "trace_block_size": _BLOCK_SIZE, + "input_length": input_length, + "input_sequence_hashes": _aligned_hashes(input_length, base=base), + }, + }, + } + + +def _write_dynamo_replay_trace(path: Path) -> Path: + """20 replay-bearing trees; peak is decided by ``replay.input_length``. + + ``input_tokens`` is INVERTED relative to ``input_length``: the 8 over-limit + trees (by input_length) carry a small input_tokens, and the 12 eligible ones + carry a large input_tokens. So a peak computed off ``input_tokens`` (the + wrong field) would REJECT the eligibles and KEEP the over-limits -- a total + inversion the fused-vs-serial parity assertion would catch. + """ + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("wb") as f: + for index in range(_TOTAL): + if _is_over_limit(index): + input_length, input_tokens = _OVER_INPUT, _UNDER_INPUT + else: + input_length, input_tokens = _UNDER_INPUT, _OVER_INPUT + record = _dynamo_replay_record( + f"s{index:02d}", + input_length=input_length, + input_tokens=input_tokens, + base=index * 1000 + 1, + ) + f.write(orjson.dumps(record)) + f.write(b"\n") + return path + + +def _session_ids(parsed) -> set[str]: + """Every dynamo session across ALL selected trees. + + Dynamo now emits one ``GraphRecord`` per session-tree under + ``parsed.graphs`` (multi-graph), so gather node sessions across every + per-tree graph (falling back to the single ``parsed.graph`` for a + single-graph parse); ``parsed.graph`` alone is only the FIRST tree. + """ + graphs = list(parsed.graphs.values()) or [parsed.graph] + return { + node.metadata["dynamo"]["session_id"] + for graph in graphs + for node in graph.nodes.values() + if isinstance(node, LlmNode) + } + + +def test_dynamo_filter_then_cap_builds_exactly_n_eligible(tmp_path: Path) -> None: + path = _write_dynamo_trace(tmp_path / "trace.jsonl") + stats_out: list[SelectionStats] = [] + + parsed = from_dynamo_trace( + path, + num_dataset_entries=_NUM_ENTRIES, + max_context_length=_MAX_CTX, + selection_out=stats_out, + ) + + # One node per selected single-turn tree; exactly N eligible built. + assert len(_session_ids(parsed)) == _NUM_ENTRIES + assert len(stats_out) == 1 + stats = stats_out[0] + assert stats.loaded == _NUM_ENTRIES + assert stats.eligible == _NUM_ENTRIES + assert stats.rejected_by_maxctx == _OVER_LIMIT_COUNT + assert stats.largest_observed == _OVER_INPUT + _OUTPUT + # The kept trees are the first N eligible (s08..s17), never any over-limit one. + expected = { + f"s{i:02d}" for i in range(_OVER_LIMIT_COUNT, _OVER_LIMIT_COUNT + _NUM_ENTRIES) + } + assert _session_ids(parsed) == expected + + +def test_dynamo_both_none_builds_everything(tmp_path: Path) -> None: + path = _write_dynamo_trace(tmp_path / "trace.jsonl") + stats_out: list[SelectionStats] = [] + + parsed = from_dynamo_trace(path, selection_out=stats_out) + + assert len(_session_ids(parsed)) == _TOTAL + assert stats_out == [] + + +def test_dynamo_fused_parallel_selection_matches_serial( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Selection applies identically on the fused-parallel build path. + + Uses a REPLAY-bearing corpus (populated ``input_sequence_hashes`` + + ``replay.input_length``), so the fused path's hash-free scan exercises the + ``replay.input_length`` branch of ``_line_peak_context`` AND the prefix-bound + hash-exclusion cut (the production path for real captures) -- not just the + ``input_tokens`` fallback. ``input_tokens`` is INVERTED relative to + ``input_length``, so a scan reading the wrong field would flip the selected + set; forcing the pool path (threshold 0, 2 workers) must still select the + SAME trees and stats as the serial ``dynamo_tree_peak_context`` path. + """ + path = _write_dynamo_replay_trace(tmp_path / "trace.jsonl") + + serial_stats: list[SelectionStats] = [] + serial = from_dynamo_trace( + path, + num_dataset_entries=_NUM_ENTRIES, + max_context_length=_MAX_CTX, + selection_out=serial_stats, + ) + + # Selection was decided by replay.input_length (not the inverted + # input_tokens): the kept trees are the low-input_length ones, s08..s17. + expected = { + f"s{i:02d}" for i in range(_OVER_LIMIT_COUNT, _OVER_LIMIT_COUNT + _NUM_ENTRIES) + } + assert _session_ids(serial) == expected + assert serial_stats[0].largest_observed == _OVER_INPUT + _OUTPUT + + monkeypatch.setattr(Environment.DATASET, "DYNAMO_GRAPH_PARALLEL_THRESHOLD", 0) + monkeypatch.setattr(Environment.DATASET, "DYNAMO_GRAPH_PARALLEL_WORKERS", 2) + parallel_stats: list[SelectionStats] = [] + parallel = from_dynamo_trace( + path, + num_dataset_entries=_NUM_ENTRIES, + max_context_length=_MAX_CTX, + selection_out=parallel_stats, + ) + + assert _session_ids(parallel) == _session_ids(serial) + assert len(_session_ids(parallel)) == _NUM_ENTRIES + assert parallel_stats[0] == serial_stats[0] + assert parallel_stats[0].rejected_by_maxctx == _OVER_LIMIT_COUNT diff --git a/tests/unit/dataset/graph/adapters/test_weka_dynamo_parity.py b/tests/unit/dataset/graph/adapters/test_weka_dynamo_parity.py new file mode 100644 index 0000000000..f0a9589411 --- /dev/null +++ b/tests/unit/dataset/graph/adapters/test_weka_dynamo_parity.py @@ -0,0 +1,747 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Cross-adapter parity: one logical trace, two formats, ONE lowered shape. + +Both recorded-trace adapters lower through the shared trie pipeline +(``build_trie_ir`` -> ``assemble_trie_graph``), so the SAME logical trace -- +identical request starts / durations / hash prefixes / output lengths / +streaming flags -- expressed once as a weka JSON trace and once as a dynamo +JSONL capture must produce identical ``LlmNode`` shapes, identical +interval-order edges, and (under weka ``hash_id_scope: "global"``, the bare +``hash_id`` reseed namespace dynamo always uses) byte-identical synthesized +prompt content with identical content-addressed ``prompt_segment_ids``. + +Each helper builds BOTH fixtures from one ``_TurnSpec`` list, so the two +files can never drift apart. Node ids are ``{scope}:{turn}`` built from +RECORDED identity on both sides (weka: trace id / agent_id; dynamo: session +ids), so with aligned fixture identifiers the two formats produce IDENTICAL +node ids -- and the whole compare keys on them RAW: every ``LlmNode`` field +(a new field enters the compare automatically), every edge verbatim, the +channel state, and the FULL segment pool byte-for-byte. Response/tiny bytes +are included: both adapters seed partial-tail synthesis with the same +trace-scoped prefix (weka: trace id; dynamo: root session id) over the same +node ids, so content-addressed segment ids -- and therefore bytes -- match. + +DOCUMENTED divergences -- the only values excluded from the compare, each +pinned by a dedicated test so any future unification tightens this suite +deliberately: + +* ``extra_headers``: dynamo replays recorded x-dynamo-* session-identity + headers; weka records none. +* ``expected`` CACHE fields: dynamo records engine ``cached_tokens``; weka + has no equivalent recording. The recorded ``input_tokens`` / + ``output_tokens`` expectations ARE compared (both formats record them). +* ``metadata["dynamo"]``: dynamo's adapter identity breadcrumb (session / + turn). The shared ``metadata["trie"]`` sub-dict IS compared raw. +* ``TraceRecord`` provenance: adapter ``tags`` and dynamo's multi-graph + ``graph_ref``. + +Partial tails are NOT a divergence: dynamo records one extra partial-tail +hash for a non-block-aligned input (weka records full blocks only), but its +lowering strips it -- engines cache FULL blocks only, and the sub-block tail +is seed-sampled identically on both sides -- so the same recording lowers +byte-identically wherever the tail sits, and covered-count-0 tiny prompts +synthesize the same sampled user message via ``small_prompt_fallback`` on +BOTH adapters. +""" + +from __future__ import annotations + +import random +import zlib +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import msgspec +import orjson +import pytest + +from aiperf.dataset.graph.adapters.dynamo.trace import from_dynamo_trace +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.models import LlmNode, ParsedGraph, StaticEdge + +_SEED = 20260707 +_BLOCK_SIZE = 16 +_MODEL = "recorded-model" +_EPOCH_MS = 1_000_000 # dynamo absolute base; both timelines start at t=0. + + +@dataclass(frozen=True) +class _TurnSpec: + """One logical recorded call, expressible in both formats.""" + + t: float + """Start, seconds from trace start (whole milliseconds only).""" + api: float + """Server processing time, seconds (whole milliseconds only).""" + hashes: tuple[int, ...] + """KV block hashes covering the block-aligned prompt prefix.""" + out: int + """Recorded output tokens (0 exercises the wire_output_cap upgrade).""" + ttft: float | None = None + """Recorded time-to-first-token, seconds. None = non-streaming.""" + session: str = "root" + """"root" or the child scope id verbatim (weka subagent ``agent_id`` == + dynamo child session id), e.g. ``"agent_001"``.""" + tail: int = 0 + """Recorded input tokens past the covered blocks (partial-tail sampling); + with ``hashes=()`` this is a tiny sub-block prompt.""" + + +# Aligned recorded identifiers: the weka trace id equals the dynamo root +# session id, and the weka subagent agent_id equals the dynamo child session +# id, so the {scope}:{turn} node ids come out IDENTICAL across formats. +_TRACE_ID = "parity" +_CHILD_ID = "agent_001" + + +def _in_len(spec: _TurnSpec) -> int: + return len(spec.hashes) * _BLOCK_SIZE + spec.tail + + +def _sessions_of(specs: list[_TurnSpec]) -> list[str]: + """Child scope ids in first-appearance order.""" + seen: list[str] = [] + for s in specs: + if s.session != "root" and s.session not in seen: + seen.append(s.session) + return seen + + +def _resolve_parents( + specs: list[_TurnSpec], parents: dict[str, str] | None +) -> dict[str, str]: + """Child scope -> parent scope ("root" unless the test says otherwise).""" + return {sid: (parents or {}).get(sid, "root") for sid in _sessions_of(specs)} + + +def _weka_trace_dict( + specs: list[_TurnSpec], + *, + trace_id: str = _TRACE_ID, + parents: dict[str, str] | None = None, + hash_scope: str = "global", +) -> dict[str, Any]: + """One weka JSON trace: each child session becomes a (possibly nested) + subagent entry placed in its PARENT's request stream at the child's first + start. ``hash_id_scope: "global"`` selects the bare-hash-id content + namespace dynamo uses, so equal hashes synthesize equal bytes.""" + tree = _resolve_parents(specs, parents) + + def req(spec: _TurnSpec) -> dict[str, Any]: + body: dict[str, Any] = { + "t": spec.t, + "type": "s" if spec.ttft is not None else "n", + "model": _MODEL, + "in": _in_len(spec), + "out": spec.out, + "hash_ids": list(spec.hashes), + "api_time": spec.api, + } + if spec.ttft is not None: + body["ttft"] = spec.ttft + return body + + def stream_for(scope: str) -> list[dict[str, Any]]: + own = [(s.t, req(s)) for s in specs if s.session == scope] + for child in (c for c, p in tree.items() if p == scope): + first_t = min(s.t for s in specs if s.session == child) + own.append( + ( + first_t, + { + "t": first_t, + "type": "subagent", + "agent_id": child, + "subagent_type": "X", + "status": "completed", + "requests": stream_for(child), + "models": [_MODEL], + }, + ) + ) + own.sort(key=lambda pair: pair[0]) + return [entry for _, entry in own] + + return { + "id": trace_id, + "models": [_MODEL], + "block_size": _BLOCK_SIZE, + "hash_id_scope": hash_scope, + "requests": stream_for("root"), + } + + +def _write_weka( + tmp_path: Path, + specs: list[_TurnSpec], + *, + parents: dict[str, str] | None = None, + hash_scope: str = "global", +) -> Path: + p = tmp_path / "parity.weka.json" + p.write_bytes( + orjson.dumps(_weka_trace_dict(specs, parents=parents, hash_scope=hash_scope)) + ) + return p + + +def _dynamo_records( + specs: list[_TurnSpec], + *, + trace_id: str = _TRACE_ID, + parents: dict[str, str] | None = None, +) -> list[dict[str, Any]]: + """The SAME logical trace as dynamo request_end records: root turns on the + root session, child turns on child sessions linked by parent_session_id + (one session-tree -> one graph, matching the weka single-trace shape).""" + tree = _resolve_parents(specs, parents) + records = [] + for i, spec in enumerate(specs): + sid = trace_id if spec.session == "root" else spec.session + ctx: dict[str, Any] = {"session_id": sid} + if spec.session != "root": + parent = tree[spec.session] + ctx["parent_session_id"] = trace_id if parent == "root" else parent + # Dynamo hashes span the WHOLE input incl. one partial-tail block for a + # non-block-aligned length ((n-1)*bs < input_length <= n*bs); weka + # records full blocks only. The partial hash is content over the actual + # tail tokens, so it never recurs: mint a unique id per tail turn + # (trace-salted so multi-trace captures cannot collide). + hashes = list(spec.hashes) + if spec.tail: + hashes.append(900_000 + (zlib.crc32(trace_id.encode()) % 50) * 1_000 + i) + req: dict[str, Any] = { + "request_id": f"r-{sid}-{spec.t}", + "model": _MODEL, + "input_tokens": _in_len(spec), + "output_tokens": spec.out, + "cached_tokens": 0, + "request_received_ms": _EPOCH_MS + round(spec.t * 1000), + # Whole-ms ints: n/1000*1000 is not always n in binary floats, and + # a sub-ms drift here would quantize differently downstream. + "total_time_ms": round(spec.api * 1000), + "replay": { + "trace_block_size": _BLOCK_SIZE, + "input_length": _in_len(spec), + "input_sequence_hashes": hashes, + }, + } + if spec.ttft is not None: + req["ttft_ms"] = round(spec.ttft * 1000) + records.append( + { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": _EPOCH_MS + round((spec.t + spec.api) * 1000), + "event_source": "dynamo", + "agent_context": ctx, + "request": req, + } + ) + return records + + +def _write_dynamo( + tmp_path: Path, + specs: list[_TurnSpec], + *, + parents: dict[str, str] | None = None, +) -> Path: + p = tmp_path / "parity.dynamo.jsonl" + p.write_bytes( + b"\n".join(orjson.dumps(r) for r in _dynamo_records(specs, parents=parents)) + ) + return p + + +def _parse_pair( + tmp_path: Path, + specs: list[_TurnSpec], + *, + idle_gap_cap_seconds: float | None = 60.0, + parents: dict[str, str] | None = None, + hash_scope: str = "global", +) -> tuple[ParsedGraph, ParsedGraph]: + weka = from_weka_trace( + _write_weka(tmp_path, specs, parents=parents, hash_scope=hash_scope), + content_root_seed=_SEED, + content_tokenizer="builtin", + idle_gap_cap_seconds=idle_gap_cap_seconds, + ) + dyn = from_dynamo_trace( + _write_dynamo(tmp_path, specs, parents=parents), + content_root_seed=_SEED, + content_tokenizer="builtin", + idle_gap_cap_seconds=idle_gap_cap_seconds, + ) + return weka, dyn + + +# --- the raw compare ---------------------------------------------------------- + +# LlmNode fields excluded from the raw compare; each is pinned as a REAL, +# documented recording-capability divergence by test_documented_divergences_hold +# below. metadata is popped whole but its shared "trie" sub-dict (incl. the +# content-addressed prompt_segment_ids) re-enters the compare verbatim, and +# expected re-enters with its recorded input/output expectations (only the +# engine cache fields are a dynamo-only recording). +_DOCUMENTED_DIVERGENT_FIELDS = ("metadata", "expected", "extra_headers") + + +def _node_fields(node: LlmNode) -> dict[str, Any]: + """Every LlmNode field raw -- a field added tomorrow is compared tomorrow.""" + d = msgspec.structs.asdict(node) + for field in _DOCUMENTED_DIVERGENT_FIELDS: + d.pop(field) + d["trie_metadata"] = node.metadata["trie"] + exp = node.expected + d["expected_recorded"] = ( + None if exp is None else (exp.input_tokens, exp.output_tokens) + ) + return d + + +def _assert_node_parity(weka_node: LlmNode, dyn_node: LlmNode, nid: str) -> None: + """Raw field compare -- block counts included. + + Dynamo RECORDS one extra partial-tail hash for a non-block-aligned input + (weka records full blocks only), but its lowering strips it (engines + cache full blocks only; the tail is seed-sampled), so the two formats + lower the same recording to identical block structure everywhere. + """ + assert _node_fields(weka_node) == _node_fields(dyn_node), nid + + +def _assert_parity(weka: ParsedGraph, dyn: ParsedGraph) -> None: + assert len(weka.traces) == len(dyn.traces) == 1 + assert weka.traces[0].id == dyn.traces[0].id + assert weka.traces[0].initial_state == dyn.traces[0].initial_state + assert weka.traces[0].replay_outputs == dyn.traces[0].replay_outputs + # Node ids are recorded data on both sides -- with aligned fixture + # identifiers they are IDENTICAL, so the compare keys on them raw and a + # divergence localizes to (node id, field). + assert sorted(weka.graph.nodes) == sorted(dyn.graph.nodes) + for nid in weka.graph.nodes: + _assert_node_parity(weka.graph.nodes[nid], dyn.graph.nodes[nid], nid) + # Edges VERBATIM: sources/targets are the shared node ids, so no index + # rewrite -- every edge type and every delay field, byte-for-byte. + assert sorted(map(repr, weka.graph.edges)) == sorted(map(repr, dyn.graph.edges)) + assert weka.graph.state == dyn.graph.state + # Segment pools byte-identical: ids are content-addressed, values compare + # raw -- prompt AND response/tiny bytes (equal trace-scoped tail seeds). + assert weka.segment_pool is not None and dyn.segment_pool is not None + assert weka.segment_pool.by_id == dyn.segment_pool.by_id + + +def _assert_shape_parity(weka: ParsedGraph, dyn: ParsedGraph) -> None: + """Everything except synthesized CONTENT: the compare for weka's LOCAL + hash-scope namespace (the one documented content divergence), where + segment ids and pool bytes legitimately differ while ids, timing, edges, + counts, and state must not.""" + assert sorted(weka.graph.nodes) == sorted(dyn.graph.nodes) + for nid in weka.graph.nodes: + wd, dd = ( + _node_fields(weka.graph.nodes[nid]), + _node_fields(dyn.graph.nodes[nid]), + ) + wd.pop("trie_metadata") + dd.pop("trie_metadata") + assert wd == dd, nid + assert sorted(map(repr, weka.graph.edges)) == sorted(map(repr, dyn.graph.edges)) + assert weka.graph.state == dyn.graph.state + + +# --- scenarios ---------------------------------------------------------------- + +_LINEAR = [ + # Streaming turn, then a zero-output non-streaming turn (wire_output_cap + # upgrade parity), then a streaming continuation -- each extending the + # recorded hash prefix so content parents chain. + _TurnSpec(t=0.0, api=1.0, hashes=(1, 2), out=8, ttft=0.5), + _TurnSpec(t=2.5, api=1.0, hashes=(1, 2, 3, 4), out=0), + _TurnSpec(t=5.0, api=0.5, hashes=(1, 2, 3, 4, 5, 6), out=16, ttft=0.25), +] + + +def test_linear_session_lowers_identically(tmp_path: Path) -> None: + """3-turn session: node fields, edges, and ALL segment BYTES match.""" + weka, dyn = _parse_pair(tmp_path, _LINEAR) + _assert_parity(weka, dyn) + # Anti-vacuous: the compare covered real values, not empty defaults -- + # keyed by the data-inherent {scope}:{turn} node ids. + nodes = [weka.graph.nodes[f"{_TRACE_ID}:{k}"] for k in range(3)] + assert [n.max_tokens for n in nodes] == [8, 1, 16] + assert [n.streaming for n in nodes] == [True, False, True] + assert all(n.model == _MODEL for n in nodes) + assert [n.theoretical_prefix_cache_total_blocks for n in nodes] == [2, 4, 6] + assert [n.theoretical_prefix_cache_hit_blocks for n in nodes] == [0, 2, 4] + # Recorded expectations, NOT the capped wire value: the zero-output turn + # keeps expected.output_tokens == 0 while max_tokens upgrades to 1. + assert [n.expected.input_tokens for n in nodes] == [32, 64, 96] + assert [n.expected.output_tokens for n in nodes] == [8, 0, 16] + assert all(n.metadata["trie"]["prompt_segment_ids"] for n in nodes) + + +def test_idle_gap_warp_engages_identically(tmp_path: Path) -> None: + """A 200s recorded idle gap compresses to the 60s cap on BOTH timelines.""" + specs = [ + _TurnSpec(t=0.0, api=1.0, hashes=(1, 2), out=8, ttft=0.5), + _TurnSpec(t=201.0, api=1.0, hashes=(1, 2, 3, 4), out=8, ttft=0.5), + ] + weka, dyn = _parse_pair(tmp_path, specs, idle_gap_cap_seconds=60.0) + _assert_parity(weka, dyn) + # Anti-vacuous: the warp actually engaged (raw gap 200s -> capped 60s). + for pb in (weka, dyn): + second = pb.graph.nodes[f"{_TRACE_ID}:1"] + assert second.arrival_offset_us == pytest.approx(61_000_000) + (edge,) = [ + e + for e in pb.graph.edges + if isinstance(e, StaticEdge) and e.delay_after_predecessor_us + ] + assert edge.delay_after_predecessor_us == pytest.approx(60_000_000) + + +def test_subagent_child_session_concurrency_parity(tmp_path: Path) -> None: + """Weka subagent inner requests and a dynamo child session lower to the + SAME concurrency shape: the overlapping child start-anchors to the + in-flight root turn, and the later root turn waits on both.""" + specs = [ + _TurnSpec(t=0.0, api=2.0, hashes=(1, 2), out=8, ttft=0.5), + _TurnSpec(t=0.5, api=0.5, hashes=(7, 8), out=4, session=_CHILD_ID), + _TurnSpec(t=3.0, api=1.0, hashes=(1, 2, 3, 4), out=8, ttft=0.5), + ] + weka, dyn = _parse_pair(tmp_path, specs) + _assert_parity(weka, dyn) + # Anti-vacuous: the child really start-anchored to the in-flight root + # turn -- the edge endpoints are the shared recorded ids themselves. + for pb in (weka, dyn): + anchored = [ + e + for e in pb.graph.edges + if isinstance(e, StaticEdge) + and e.delay_after_predecessor_start_us is not None + ] + assert len(anchored) == 1 + assert anchored[0].source == f"{_TRACE_ID}:0" + assert anchored[0].target == f"{_CHILD_ID}:0" + assert anchored[0].delay_after_predecessor_start_us == pytest.approx(500_000) + + +def test_response_segments_byte_identical(tmp_path: Path) -> None: + """Assistant response segments exist AND byte-match across formats: both + adapters seed partial-tail synthesis with the same trace-scoped prefix + (weka: trace id; dynamo: root session id) over the same node ids, so the + content-addressed segment ids -- hence bytes -- are equal. The pool-wide + equality lives in ``_assert_parity``; this pins that it is NOT vacuous + (response segments really got synthesized on both sides).""" + weka, dyn = _parse_pair(tmp_path, _LINEAR) + _assert_parity(weka, dyn) + for pb in (weka, dyn): + for nid, node in pb.graph.nodes.items(): + if node.max_tokens == 1: + continue # zero-output turn: empty response segment shape varies + tip = node.metadata["trie"]["prompt_segment_ids"][-1] + assert any( + s.role == "assistant" and s.parent_id == tip + for s in pb.segment_pool.by_id.values() + ), nid + + +def test_documented_divergences_hold(tmp_path: Path) -> None: + """Pin the exclusion list: every value skipped by the raw compare is a + REAL divergence today. When one is unified, move it INTO the compare.""" + weka, dyn = _parse_pair(tmp_path, _LINEAR) + for node in dyn.graph.nodes.values(): + assert node.extra_headers is not None # session identity headers + assert node.expected.cache_read_tokens is not None # engine recording + assert set(node.metadata) == {"dynamo", "trie"} + for node in weka.graph.nodes.values(): + assert node.extra_headers is None + assert node.expected.cache_read_tokens is None # no weka equivalent + assert node.expected.cache_creation_tokens is None + assert set(node.metadata) == {"trie"} + # TraceRecord provenance: adapter tags + dynamo's multi-graph graph_ref. + assert weka.traces[0].tags == ["from-weka-trace"] + assert dyn.traces[0].tags == ["from-dynamo-trace"] + assert weka.traces[0].graph_ref is None + assert dyn.traces[0].graph_ref == _TRACE_ID + + +def test_partial_tail_turns_byte_identical_anywhere(tmp_path: Path) -> None: + """Non-block-aligned turns, each format's NATIVE encoding: weka records + full-block hashes + an uncovered tail; dynamo records one extra + partial-tail hash spanning it, which its lowering STRIPS (engines cache + full blocks only; the tail is seed-sampled, never decoded). The same + recording therefore lowers byte-identically wherever the tail sits -- + leaf turns, MID-CHAIN turns whose prefix later turns extend, and child + sessions -- with identical block totals.""" + specs = [ + _TurnSpec(t=0.0, api=1.0, hashes=(1, 2), out=8, ttft=0.5), + _TurnSpec(t=1.5, api=0.5, hashes=(50, 51), out=4, tail=5, session=_CHILD_ID), + # Mid-chain tails: later turns extend these prefixes. + _TurnSpec(t=2.0, api=1.0, hashes=(1, 2, 3), out=16, tail=9), + _TurnSpec(t=4.0, api=1.0, hashes=(1, 2, 3, 4, 5), out=0, tail=3), + _TurnSpec(t=6.0, api=1.0, hashes=tuple(range(1, 11)), out=32, tail=11), + ] + weka, dyn = _parse_pair(tmp_path, specs) + _assert_parity(weka, dyn) + # Anti-vacuous: the recorded lengths really are unaligned, the dynamo + # recording really carried the extra partial hash (writer appends it), + # and both sides count FULL blocks only. + nodes = [ + weka.graph.nodes[n] + for n in (f"{_CHILD_ID}:0", f"{_TRACE_ID}:1", f"{_TRACE_ID}:3") + ] + assert [n.expected.input_tokens for n in nodes] == [37, 57, 171] + assert [n.theoretical_prefix_cache_total_blocks for n in nodes] == [2, 3, 10] + + +def test_tiny_sub_block_prompts_byte_identical(tmp_path: Path) -> None: + """Covered-count-0 (tiny sub-block) prompts: weka records hashes=[] with + a sub-block ``in``; dynamo records a single partial hash (stripped at + lowering). Both sides synthesize the SAME seed-sampled user message via + ``small_prompt_fallback`` -- a recorded prompt is never lowered to an + unreplayable empty messages array.""" + tiny = [ + _TurnSpec(t=0.0, api=0.5, hashes=(), out=4, tail=7), + _TurnSpec(t=1.0, api=0.5, hashes=(1,), out=4, ttft=0.25), + ] + weka, dyn = _parse_pair(tmp_path, tiny) + _assert_parity(weka, dyn) + w_node = weka.graph.nodes[f"{_TRACE_ID}:0"] + d_node = dyn.graph.nodes[f"{_TRACE_ID}:0"] + assert len(w_node.metadata["trie"]["prompt_segment_ids"]) == 1 + assert d_node.metadata["dynamo"]["small_prompt"] is True + assert w_node.expected.input_tokens == 7 + + +def test_local_hash_scope_lowers_same_shape(tmp_path: Path) -> None: + """Weka's DEFAULT ``hash_id_scope: "local"`` must still agree with dynamo + on everything content-free: ids, node fields, timing, edges, state. The + content namespaces intentionally diverge (local reseeds per trace), so + prompt segment ids differ -- pinned here so the divergence stays real.""" + weka, dyn = _parse_pair(tmp_path, _LINEAR, hash_scope="local") + _assert_shape_parity(weka, dyn) + # Anti-vacuous: the namespaces really diverged (else this test should be + # folded into the full byte compare). + w_ids = [ + weka.graph.nodes[nid].metadata["trie"]["prompt_segment_ids"] + for nid in sorted(weka.graph.nodes) + ] + d_ids = [ + dyn.graph.nodes[nid].metadata["trie"]["prompt_segment_ids"] + for nid in sorted(dyn.graph.nodes) + ] + assert w_ids != d_ids + + +def test_nested_and_sibling_children_lower_identically(tmp_path: Path) -> None: + """Two root children plus a GRANDCHILD (weka subagent-in-subagent nesting, + dynamo parent_session_id chain) lower to identical node ids and shapes.""" + specs = [ + _TurnSpec(t=0.0, api=4.0, hashes=(1, 2), out=8, ttft=0.5), + _TurnSpec(t=0.5, api=1.0, hashes=(10, 11), out=4, session="agent_001"), + _TurnSpec(t=1.0, api=0.5, hashes=(20,), out=4, session="agent_002"), + _TurnSpec(t=2.0, api=1.0, hashes=(30, 31), out=4, session="agent_003"), + _TurnSpec(t=5.0, api=1.0, hashes=(1, 2, 3), out=8, ttft=0.5), + ] + parents = { + "agent_001": "root", + "agent_002": "agent_001", # grandchild + "agent_003": "root", + } + weka, dyn = _parse_pair(tmp_path, specs, parents=parents) + _assert_parity(weka, dyn) + assert sorted(weka.graph.nodes) == [ + "agent_001:0", + "agent_002:0", + "agent_003:0", + f"{_TRACE_ID}:0", + f"{_TRACE_ID}:1", + ] + + +def test_corpus_two_traces_lower_identically(tmp_path: Path) -> None: + """Corpus level: a weka DIRECTORY of two traces vs ONE dynamo capture + holding two disjoint session trees -- the multi-graph split must agree on + graph keys, per-graph nodes/edges, and trace->graph_ref resolution.""" + specs_a = [ + _TurnSpec(t=0.0, api=1.0, hashes=(1, 2), out=8, ttft=0.5), + _TurnSpec(t=2.0, api=1.0, hashes=(1, 2, 3), out=8), + ] + specs_b = [ + _TurnSpec(t=0.0, api=2.0, hashes=(50, 51), out=8, ttft=0.5), + _TurnSpec(t=0.5, api=0.5, hashes=(60,), out=4, session="agent_001"), + ] + + weka_dir = tmp_path / "weka_corpus" + weka_dir.mkdir() + (weka_dir / "a.json").write_bytes( + orjson.dumps(_weka_trace_dict(specs_a, trace_id="trace_a")) + ) + (weka_dir / "b.json").write_bytes( + orjson.dumps(_weka_trace_dict(specs_b, trace_id="trace_b")) + ) + dynamo_path = tmp_path / "corpus.dynamo.jsonl" + dynamo_path.write_bytes( + b"\n".join( + orjson.dumps(r) + for r in ( + _dynamo_records(specs_a, trace_id="trace_a") + + _dynamo_records(specs_b, trace_id="trace_b") + ) + ) + ) + + weka = from_weka_trace( + weka_dir, + content_root_seed=_SEED, + content_tokenizer="builtin", + idle_gap_cap_seconds=60.0, + ) + dyn = from_dynamo_trace( + dynamo_path, + content_root_seed=_SEED, + content_tokenizer="builtin", + idle_gap_cap_seconds=60.0, + ) + + assert sorted(weka.graphs) == sorted(dyn.graphs) == ["trace_a", "trace_b"] + assert {t.id: t.graph_ref for t in weka.traces} == { + t.id: t.graph_ref for t in dyn.traces + } + for key in weka.graphs: + wg, dg = weka.graphs[key], dyn.graphs[key] + assert sorted(wg.nodes) == sorted(dg.nodes), key + for nid in wg.nodes: + _assert_node_parity(wg.nodes[nid], dg.nodes[nid], f"{key}/{nid}") + assert sorted(map(repr, wg.edges)) == sorted(map(repr, dg.edges)), key + assert wg.state == dg.state, key + # One corpus-wide content pool on both sides, byte-identical. + assert weka.segment_pool.by_id == dyn.segment_pool.by_id + + +def test_dynamo_record_order_invariance(tmp_path: Path) -> None: + """Dynamo lowering globally sorts turns by recorded start, so a + line-shuffled capture must parse to a byte-identical graph.""" + specs = [ + _TurnSpec(t=0.0, api=2.0, hashes=(1, 2), out=8, ttft=0.5), + _TurnSpec(t=0.5, api=0.5, hashes=(7, 8), out=4, session=_CHILD_ID), + _TurnSpec(t=3.0, api=1.0, hashes=(1, 2, 3, 4), out=8, ttft=0.5), + ] + ordered_path = _write_dynamo(tmp_path, specs) + lines = ordered_path.read_bytes().splitlines() + shuffled = list(lines) + random.Random(_SEED).shuffle(shuffled) + assert shuffled != lines, "shuffle must actually reorder the capture" + shuffled_path = tmp_path / "shuffled.dynamo.jsonl" + shuffled_path.write_bytes(b"\n".join(shuffled)) + + kwargs = dict( + content_root_seed=_SEED, + content_tokenizer="builtin", + idle_gap_cap_seconds=60.0, + ) + ordered = from_dynamo_trace(ordered_path, **kwargs) + reordered = from_dynamo_trace(shuffled_path, **kwargs) + + assert sorted(ordered.graph.nodes) == sorted(reordered.graph.nodes) + for nid in ordered.graph.nodes: + assert msgspec.structs.asdict(ordered.graph.nodes[nid]) == ( + msgspec.structs.asdict(reordered.graph.nodes[nid]) + ), nid + assert sorted(map(repr, ordered.graph.edges)) == ( + sorted(map(repr, reordered.graph.edges)) + ) + assert ordered.graph.state == reordered.graph.state + assert ordered.segment_pool.by_id == reordered.segment_pool.by_id + + +# --- randomized scenario matrix ------------------------------------------------- +# +# The hand-picked scenarios above localize known seams; this sweep validates +# the parity contract EMPIRICALLY across the recorded-shape space: random turn +# counts, prefix growth, streaming mix, zero-output turns, partial tails, +# idle gaps beyond the warp cap, and a random session TREE (0-2 children, +# each parented to root or a previous child -- i.e. possible grandchildren). +# Fixed per-case seeds keep every case reproducible standalone. + + +def _random_specs(seed: int) -> tuple[list[_TurnSpec], dict[str, str]]: + """One random recorded shape, valid in both formats (whole-ms times).""" + rng = random.Random(seed) + hash_ids = iter(range(1, 10_000)) + specs: list[_TurnSpec] = [] + parents: dict[str, str] = {} + + t_ms = 0 + prefix: tuple[int, ...] = () + root_turns = rng.randint(2, 5) + for _ in range(root_turns): + prefix += tuple(next(hash_ids) for _ in range(rng.randint(1, 4))) + api_ms = rng.randint(200, 3000) + out = rng.choice([0, 4, 8, 16, 32]) + # A recorded ttft implies at least one produced token; keep it < api. + ttft_ms = rng.randint(50, api_ms - 1) if out and rng.random() < 0.7 else None + specs.append( + _TurnSpec( + t=t_ms / 1000, + api=api_ms / 1000, + hashes=prefix, + out=out, + ttft=ttft_ms / 1000 if ttft_ms is not None else None, + tail=rng.choice([0, 0, 3, 9]), + ) + ) + # Idle gap after the turn; sometimes far beyond the 60s warp cap. + gap_ms = rng.choice([100, 500, 2000, 150_000]) + t_ms += api_ms + gap_ms + + # A random session TREE: each child overlaps its parent's first turn and + # runs its own recorded hash prefix (weka nested subagent entries / dynamo + # parent_session_id chains). + scopes = ["root"] + for i in range(rng.randint(0, 2)): + sid = f"agent_{i + 1:03d}" + parent = rng.choice(scopes) + parents[sid] = parent + parent_first = next( + s for s in specs if s.session == ("root" if parent == "root" else parent) + ) + child_t_ms = round(parent_first.t * 1000) + rng.randint( + 1, max(1, int(parent_first.api * 900)) + ) + child_prefix: tuple[int, ...] = () + for _ in range(rng.randint(1, 3)): + child_prefix += tuple(next(hash_ids) for _ in range(rng.randint(1, 3))) + api_ms = rng.randint(100, 1500) + out = rng.choice([4, 8, 16]) + specs.append( + _TurnSpec( + t=child_t_ms / 1000, + api=api_ms / 1000, + hashes=child_prefix, + out=out, + ttft=None, + session=sid, + tail=rng.choice([0, 0, 5]), + ) + ) + child_t_ms += api_ms + rng.randint(50, 800) + scopes.append(sid) + + return specs, parents + + +@pytest.mark.parametrize("seed", range(8)) +def test_randomized_recorded_shapes_lower_identically( + tmp_path: Path, seed: int +) -> None: + specs, parents = _random_specs(_SEED + seed) + weka, dyn = _parse_pair(tmp_path, specs, parents=parents) + _assert_parity(weka, dyn) + # Anti-vacuous: every spec became a node on both sides. + assert len(weka.graph.nodes) == len(specs) diff --git a/tests/unit/dataset/graph/segment_ir/__init__.py b/tests/unit/dataset/graph/segment_ir/__init__.py new file mode 100644 index 0000000000..e5725ea5a4 --- /dev/null +++ b/tests/unit/dataset/graph/segment_ir/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unit/dataset/graph/segment_ir/test_interval_order.py b/tests/unit/dataset/graph/segment_ir/test_interval_order.py new file mode 100644 index 0000000000..52ea8d8efb --- /dev/null +++ b/tests/unit/dataset/graph/segment_ir/test_interval_order.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Frontier transitive-reduction soundness for the interval-order edge rule.""" + +from __future__ import annotations + +from aiperf.dataset.graph.segment_ir.interval_order import ( + build_interval_edges, + compute_ranks, +) +from aiperf.dataset.graph.segment_ir.trie_content import TrieNode, TrieRequest + + +def _node( + nid: str, t: float, api: float, async_ancestors: set[str] | None = None +) -> TrieNode: + node = TrieNode( + node_id=nid, + request=TrieRequest( + hash_ids=[], input_length=0, output_length=0, t=t, api_time=api + ), + order=0, + async_ancestors=frozenset(async_ancestors or set()), + ) + node.warped_start = t + return node + + +def test_frontier_reduction_keeps_async_sibling_not_covered_by_main_chain() -> None: + """Dropping candidate c for a later candidate d is only sound when the + covering edge c -> d exists -- i.e. when d does NOT async-exclude c. + + Counterexample (S1): node n inside async subtree X, sibling c in X ending + t=10, main-chain d (not in X) running t=11..12, n starting t=13. No c -> d + edge exists anywhere (d excludes c), so c must stay in n's frontier or the + recorded c-before-n ordering inside the subtree is silently lost. + """ + c = _node("c", 5.0, 5.0, {"X"}) # in X, ends t=10 + d = _node("d", 11.0, 1.0) # main chain, t=11..12 + n = _node("n", 13.0, 1.0, {"X"}) # in X, starts t=13 + nodes = [c, d, n] + compute_ranks(nodes) + edges = build_interval_edges(nodes) + + n_sources = {e.source for e in edges["n"]} + assert "c" in n_sources, "c dropped despite no covering c -> d edge existing" + assert n_sources == {"c", "d"} + # d async-excludes c, so d has no predecessors at all -- the covering edge + # the unsound reduction assumed cannot exist. + assert {e.source for e in edges["d"]} == {"START"} diff --git a/tests/unit/dataset/graph/segment_ir/test_pool.py b/tests/unit/dataset/graph/segment_ir/test_pool.py new file mode 100644 index 0000000000..1ebdac58ea --- /dev/null +++ b/tests/unit/dataset/graph/segment_ir/test_pool.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Verbatim raw-wire-JSON segments through the pool and the unified store. + +Authored dag messages are arbitrary dicts the legacy path forwards verbatim +(key order and extra keys included). These tests pin the raw-segment variant: +``SegmentPool.add_raw_message`` interns ``orjson.dumps(message)`` verbatim and +the unified store persists that blob byte-for-byte instead of re-serializing a +normalized ``{"role", "content"}`` dict. +""" + +import asyncio + +import orjson +import pytest + +from aiperf.dataset.graph.segment_ir.pool import SegmentPool +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) + + +@pytest.fixture +def pool() -> SegmentPool: + return SegmentPool() + + +def test_add_raw_message_preserves_key_order_and_extra_keys(pool): + msg = {"content": "hi", "role": "user", "name": "alice"} + sid = pool.add_raw_message(message=msg, parent_id=None) + seg = pool.get(sid) + assert seg.wire_json == orjson.dumps(msg).decode() + assert orjson.loads(seg.wire_json) == msg + # key order survives round-trip + assert list(orjson.loads(seg.wire_json)) == ["content", "role", "name"] + + +def test_add_raw_message_dedups_by_content_and_prefix(pool): + m = {"role": "user", "content": "x"} + a = pool.add_raw_message(message=m, parent_id=None) + b = pool.add_raw_message(message=m, parent_id=None) + c = pool.add_raw_message(message=m, parent_id=a) + assert a == b + assert a != c # prefix-dependent id + + +def test_raw_and_text_ids_never_alias(pool): + m = {"role": "user", "content": "x"} + raw = pool.add_raw_message(message=m, parent_id=None) + text = pool.add_text(role="user", content="x", parent_id=None) + assert raw != text + + +def test_materialize_returns_verbatim_dict_for_raw(pool): + m = {"role": "user", "content": [{"type": "text", "text": "hi"}]} + sid = pool.add_raw_message(message=m, parent_id=None) + assert pool.materialize([sid]) == [m] + + +def test_store_round_trips_raw_and_text_segments(tmp_path): + """A raw segment persists verbatim; a role/content segment normalizes as before.""" + msg = {"content": "hi", "role": "user", "name": "alice"} + store = GraphSegmentUnifiedBackingStore(tmp_path, "raw") + raw_h = store.put_segment( + "raw_seg", "user", "", wire_json=orjson.dumps(msg).decode() + ) + text_h = store.put_segment("text_seg", "system", "SYS") + asyncio.run(store.finalize()) + + with GraphSegmentUnifiedClient(tmp_path, "raw").open() as c: + # materialize returns the verbatim dict for the raw segment + assert c.materialize_handles([raw_h]) == [msg] + assert c.materialize_handles([text_h]) == [{"role": "system", "content": "SYS"}] + # the request body embeds the exact orjson.dumps(msg) bytes verbatim + body = c.build_request_body_handles([raw_h], b"") + assert body == b'{"messages":[' + orjson.dumps(msg) + b"]}" diff --git a/tests/unit/dataset/graph/segment_ir/test_prefix_cache.py b/tests/unit/dataset/graph/segment_ir/test_prefix_cache.py new file mode 100644 index 0000000000..73f9b25d27 --- /dev/null +++ b/tests/unit/dataset/graph/segment_ir/test_prefix_cache.py @@ -0,0 +1,263 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared trie theoretical prefix-cache stamping. + +Locks the shared-pre-pass semantics onto the segment-trie IR: one shared per-trace +seen-set consumed in recorded global time order, leading-run hit counting +(stop at first miss), full hash-id totals, and every request's blocks entering +the infinite cache. Plus the stamping contract every consumer relies on: the +native ``theoretical_prefix_cache_hit_blocks`` / ``_total_blocks`` fields on +each hash-carrying ``LlmNode``, surviving the graph_meta sidecar strip. +""" + +from pathlib import Path + +from aiperf.dataset.graph.models import LlmNode +from aiperf.dataset.graph.segment_ir.prefix_cache import ( + compute_shared_prefix_cache_counts, + extract_prefix_cache_by_node, + stamp_theoretical_prefix_cache, +) +from aiperf.dataset.graph.segment_ir.trie_content import TrieNode, TrieRequest + +WEKA_FIXTURE = Path(__file__).parents[3] / "graph" / "fixtures" / "weka_min.json" + + +def _node(node_id: str, hash_ids: list[int], t: float, order: int) -> TrieNode: + return TrieNode( + node_id=node_id, + request=TrieRequest( + hash_ids=hash_ids, + input_length=len(hash_ids) * 4, + output_length=8, + t=t, + api_time=1.0, + ), + order=order, + ) + + +def _llm(node_id: str) -> LlmNode: + return LlmNode(prompt=[], output=f"{node_id}_out", metadata={}) + + +class TestComputeSharedPrefixCacheCounts: + def test_cold_cache_first_request_zero_hits(self): + counts = compute_shared_prefix_cache_counts([_node("a", [1, 2, 3], 0.0, 0)]) + assert counts == {"a": (0, 3)} + + def test_full_prefix_reuse_hits_leading_run(self): + counts = compute_shared_prefix_cache_counts( + [ + _node("a", [1, 2], 0.0, 0), + _node("b", [1, 2, 3], 1.0, 1), + ] + ) + assert counts == {"a": (0, 2), "b": (2, 3)} + + def test_leading_run_stops_at_first_miss(self): + # 9 is unseen: the run stops there even though 2 was cached. + counts = compute_shared_prefix_cache_counts( + [ + _node("a", [1, 2], 0.0, 0), + _node("b", [1, 9, 2], 1.0, 1), + ] + ) + assert counts["b"] == (1, 3) + + def test_global_time_order_not_list_order(self): + # "late" appears first in the list but has the LATER recorded t, so + # "early" seeds the cache first and "late" scores hits off it. + counts = compute_shared_prefix_cache_counts( + [ + _node("late", [1, 2], 5.0, 0), + _node("early", [1, 2], 1.0, 1), + ] + ) + assert counts == {"early": (0, 2), "late": (2, 2)} + + def test_equal_t_tiebreak_by_flatten_order(self): + counts = compute_shared_prefix_cache_counts( + [ + _node("second", [1], 1.0, 1), + _node("first", [1], 1.0, 0), + ] + ) + assert counts == {"first": (0, 1), "second": (1, 1)} + + def test_all_blocks_enter_cache_not_just_hits(self): + # b misses on 9, but its trailing 7 still enters the cache for c. + counts = compute_shared_prefix_cache_counts( + [ + _node("a", [1], 0.0, 0), + _node("b", [9, 7], 1.0, 1), + _node("c", [9, 7], 2.0, 2), + ] + ) + assert counts["c"] == (2, 2) + + def test_empty_hash_ids_yield_zero_total(self): + counts = compute_shared_prefix_cache_counts([_node("a", [], 0.0, 0)]) + assert counts == {"a": (0, 0)} + + +class TestStampTheoreticalPrefixCache: + def test_stamps_native_fields(self): + llm_nodes = {"a": _llm("a"), "b": _llm("b")} + stamp_theoretical_prefix_cache( + llm_nodes, + [_node("a", [1, 2], 0.0, 0), _node("b", [1, 2, 3], 1.0, 1)], + ) + assert llm_nodes["a"].theoretical_prefix_cache_hit_blocks == 0 + assert llm_nodes["a"].theoretical_prefix_cache_total_blocks == 2 + assert llm_nodes["b"].theoretical_prefix_cache_hit_blocks == 2 + assert llm_nodes["b"].theoretical_prefix_cache_total_blocks == 3 + + def test_zero_hash_node_left_unstamped(self): + llm_nodes = {"a": _llm("a")} + stamp_theoretical_prefix_cache(llm_nodes, [_node("a", [], 0.0, 0)]) + assert llm_nodes["a"].theoretical_prefix_cache_hit_blocks is None + assert llm_nodes["a"].theoretical_prefix_cache_total_blocks is None + + def test_preserves_existing_metadata_keys(self): + llm = LlmNode( + prompt=[], + output="a_out", + metadata={"trie": {"prompt_segment_ids": ["s1"]}, "dispatch": {"k": 1}}, + ) + llm_nodes = {"a": llm} + stamp_theoretical_prefix_cache(llm_nodes, [_node("a", [1], 0.0, 0)]) + stamped = llm_nodes["a"] + assert stamped.metadata["trie"] == {"prompt_segment_ids": ["s1"]} + assert stamped.metadata["dispatch"]["k"] == 1 + assert stamped.theoretical_prefix_cache_total_blocks == 1 + + def test_extract_round_trips_stamped_graph(self): + from aiperf.dataset.graph.models import GraphRecord + + llm_nodes = {"a": _llm("a"), "b": _llm("b")} + stamp_theoretical_prefix_cache( + llm_nodes, + [_node("a", [1], 0.0, 0), _node("b", [1, 2], 1.0, 1)], + ) + graph = GraphRecord(nodes=dict(llm_nodes)) + assert extract_prefix_cache_by_node(graph) == {"a": [0, 1], "b": [1, 2]} + + +class TestAdapterIntegration: + def test_weka_trie_build_stamps_every_hash_node(self): + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + from aiperf.dataset.graph.models import resolve_trace_graph + + pg = from_weka_trace(WEKA_FIXTURE, content_root_seed=0) + trace = pg.traces[0] + counts = extract_prefix_cache_by_node(resolve_trace_graph(pg, trace)) + # weka_min: three requests with growing hash-id prefixes; the exact + # counts pin the shared-seen-set walk on the recorded timeline. Node ids + # are ``{trace_id}:{k}`` (0-based turn) -- weka_min's trace is trace_03_n3. + assert counts == { + "trace_03_n3:0": [0, 2], + "trace_03_n3:1": [2, 3], + "trace_03_n3:2": [3, 4], + } + + def test_weka_counts_survive_sidecar_strip(self): + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + from aiperf.dataset.graph.graph_meta_sidecar import strip_replay_text + from aiperf.dataset.graph.models import resolve_trace_graph + + pg = from_weka_trace(WEKA_FIXTURE, content_root_seed=0) + stripped = strip_replay_text(pg) + counts = extract_prefix_cache_by_node( + resolve_trace_graph(stripped, stripped.traces[0]) + ) + assert counts == { + "trace_03_n3:0": [0, 2], + "trace_03_n3:1": [2, 3], + "trace_03_n3:2": [3, 4], + } + + def test_dynamo_stamps_recorded_hash_counts(self, tmp_path): + # Records with replay metadata use the RECORDED input_sequence_hashes: + # turn 2's hashes extend turn 1's, so its leading 2 blocks hit the + # shared cache. + import orjson + + from aiperf.dataset.graph.adapters.dynamo.trace import from_dynamo_trace + from aiperf.dataset.graph.models import resolve_trace_graph + + def rec(ts: int, hashes: list[int], input_tokens: int) -> dict: + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": {"session_id": "s1"}, + "request": { + "request_id": f"r{ts}", + "model": "m", + "input_tokens": input_tokens, + "output_tokens": 8, + "cached_tokens": 0, + "replay": { + "trace_block_size": 16, + "input_length": input_tokens, + "input_sequence_hashes": hashes, + }, + }, + } + + p = tmp_path / "dyn_kv.jsonl" + p.write_bytes( + b"\n".join( + orjson.dumps(r) + for r in [ + rec(1000, [111, 222], 32), + rec(2000, [111, 222, 333, 444], 64), + ] + ) + ) + pg = from_dynamo_trace( + p, + content_root_seed=0, + content_tokenizer="builtin", + ) + counts = extract_prefix_cache_by_node(resolve_trace_graph(pg, pg.traces[0])) + by_turn = {node_id.rsplit(":", 1)[1]: v for node_id, v in counts.items()} + assert by_turn == {"0": [0, 2], "1": [2, 4]} + + def test_dynamo_trie_build_stamps_hash_nodes(self): + from aiperf.dataset.graph.adapters.dynamo.trace import from_dynamo_trace + from aiperf.dataset.graph.models import resolve_trace_graph + + fixture = ( + Path(__file__).parents[1] + / "adapters" + / "fixtures" + / "dynamo_nested" + / "nested_2_level.jsonl.gz" + ) + pg = from_dynamo_trace(fixture, content_root_seed=0) + counts_by_trace = { + t.id: extract_prefix_cache_by_node(resolve_trace_graph(pg, t)) + for t in pg.traces + } + stamped = [c for c in counts_by_trace.values() if c] + assert stamped, "dynamo trie nodes must carry prefix-cache counts" + for counts in stamped: + assert all(t > 0 for _, t in counts.values()) + assert all(0 <= h <= t for h, t in counts.values()) + + def test_store_builder_builds_per_trace_map(self): + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + from aiperf.dataset.graph.store_build import GraphStoreBuilder + + pg = from_weka_trace(WEKA_FIXTURE, content_root_seed=0) + by_trace = GraphStoreBuilder._build_graph_prefix_cache_by_trace(pg) + assert by_trace == { + pg.traces[0].id: { + "trace_03_n3:0": [0, 2], + "trace_03_n3:1": [2, 3], + "trace_03_n3:2": [3, 4], + } + } diff --git a/tests/unit/dataset/graph/segment_ir/test_store_builder.py b/tests/unit/dataset/graph/segment_ir/test_store_builder.py new file mode 100644 index 0000000000..0a5efd0c66 --- /dev/null +++ b/tests/unit/dataset/graph/segment_ir/test_store_builder.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Streaming-split gates and unified-store read hardening.""" + +from __future__ import annotations + +import asyncio +from pathlib import Path + +import orjson +import pytest + +from aiperf.common.exceptions import MemoryMapSerializationError +from aiperf.dataset.graph.models import ( + GraphRecord, + LlmNode, + ParsedGraph, + TraceRecord, +) +from aiperf.dataset.graph.segment_ir.pool import SegmentPool +from aiperf.dataset.graph.segment_ir.store_builder import ( + _trie_envelope, + iter_trace_segment_payloads, +) +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) + +# --- S4: streaming split rejects slot-carrying nodes loudly ------------------ + + +def _parsed_with_trie_meta(extra_trie_meta: dict) -> ParsedGraph: + pool = SegmentPool() + sid = pool.add(role="user", content="hi", tokens=[1, 2], parent_id=None) + node = LlmNode( + prompt=[], + output="n0_out", + arrival_offset_us=0, + metadata={"trie": {"prompt_segment_ids": [sid], **extra_trie_meta}}, + ) + return ParsedGraph( + graph=GraphRecord(nodes={"n0": node}), + traces=[TraceRecord(id="t0")], + segment_pool=pool, + ) + + +@pytest.mark.parametrize( + "extra_trie_meta", + [ + pytest.param({"assembly": [{"s": {"src": "n0"}}]}, id="assembly_items"), + pytest.param({"capture": True}, id="capture"), + ], +) # fmt: skip +def test_streaming_split_rejects_slot_carrying_node(extra_trie_meta: dict) -> None: + """Assembly items / capture never reach the streamed envelope (it carries + neither), so the streaming split must fail loud naming the node instead of + silently persisting a manifest missing the node's dynamic slots.""" + parsed = _parsed_with_trie_meta(extra_trie_meta) + with pytest.raises(NotImplementedError, match="'n0'"): + list(iter_trace_segment_payloads(parsed)) + + +def test_streaming_split_accepts_slotless_node() -> None: + parsed = _parsed_with_trie_meta({}) + payloads = list(iter_trace_segment_payloads(parsed)) + assert len(payloads) == 1 and payloads[0].envelopes + + +# --- endpoint_extra_applied: adapter-owned extras precedence flag ------------- + + +def test_trie_envelope_omits_endpoint_extra_applied_when_unset() -> None: + """A node WITHOUT the adapter stamp yields an envelope with NO + ``endpoint_extra_applied`` key and byte-identical ``orjson.dumps`` output to + before -- weka/dynamo/native corpora envelopes stay bit-for-bit unchanged.""" + node = LlmNode(prompt=[], output="o", arrival_offset_us=0, streaming=False) + env = _trie_envelope(node, ["a", "b"]) + assert "endpoint_extra_applied" not in env + assert orjson.dumps(env) == orjson.dumps( + { + "prompt_segment_ids": ["a", "b"], + "dispatch_overrides": {}, + "stream": False, + } + ) + + +def test_trie_envelope_carries_endpoint_extra_applied_when_stamped() -> None: + """A node stamped ``metadata["dispatch"]["endpoint_extra_applied"] = True`` + (adapter folded the run's ``--extra-inputs`` into ``dispatch_overrides`` at + parse) carries the flag in its envelope so the worker skips re-merging.""" + node = LlmNode( + prompt=[], + output="o", + arrival_offset_us=0, + streaming=False, + metadata={"dispatch": {"endpoint_extra_applied": True}}, + ) + env = _trie_envelope(node, ["a"]) + assert env["endpoint_extra_applied"] is True + + +# --- native body-field fold-in -------------------------------------------------- + + +def test_trie_envelope_folds_native_body_fields() -> None: + """The native Turn-named fields (``model`` / ``max_tokens`` / ``raw_tools``) + fold into the envelope's wire-body dict after the ``extra_body`` vendor + keys; ``extra_headers`` rides the envelope top level, never the body.""" + tools = [{"type": "function", "function": {"name": "lookup"}}] + node = LlmNode( + prompt=[], + output="o", + arrival_offset_us=0, + streaming=False, + model="m", + max_tokens=25, + raw_tools=tools, + extra_headers={"x-dynamo-session-id": "s1"}, + extra_body={"temperature": 0.5}, + ) + env = _trie_envelope(node, ["a"]) + assert env["dispatch_overrides"] == { + "temperature": 0.5, + "model": "m", + "max_output_tokens": 25, + "tools": tools, + } + assert env["extra_headers"] == {"x-dynamo-session-id": "s1"} + assert "extra_headers" not in env["dispatch_overrides"] + assert env["stream"] is False + + +def test_trie_envelope_fold_skips_hand_authored_extra_body_entry() -> None: + """A hand-authored ``extra_body`` entry naming a foldable key wins: the + fold never duplicates or overwrites it.""" + node = LlmNode( + prompt=[], + output="o", + arrival_offset_us=0, + streaming=True, + model="native-model", + max_tokens=7, + extra_body={ + "model": "override-model", + "max_output_tokens": 3, + "top_p": 0.9, + }, + ) + env = _trie_envelope(node, ["a"]) + assert env["dispatch_overrides"] == { + "model": "override-model", + "max_output_tokens": 3, + "top_p": 0.9, + } + + +def test_trie_envelope_no_folds_when_fields_unset() -> None: + node = LlmNode(prompt=[], output="o", arrival_offset_us=0, streaming=False) + env = _trie_envelope(node, ["a"]) + assert env["dispatch_overrides"] == {} + assert "extra_headers" not in env + + +# --- S5: content-region spans bounds-checked at open -------------------------- + + +def test_truncated_content_blob_fails_loud_on_open(tmp_path: Path) -> None: + """A span past the end of content.blob (stale/partial store) must raise at + open instead of Python slice clamping returning truncated bytes.""" + store = GraphSegmentUnifiedBackingStore(tmp_path, "b") + handle = store.put_segment("seg_a", "user", "hello world, plenty of content") + store.add_node_manifest_interned("t0", 0, "profiling", [handle], {}, True) + asyncio.run(store.finalize()) + + blob = tmp_path / "aiperf_graph_segments_b" / "content.blob" + blob.write_bytes(blob.read_bytes()[:-4]) + + with pytest.raises(MemoryMapSerializationError, match="content.blob"): + GraphSegmentUnifiedClient(tmp_path, "b").open() diff --git a/tests/unit/dataset/graph/segment_ir/test_trie_content.py b/tests/unit/dataset/graph/segment_ir/test_trie_content.py new file mode 100644 index 0000000000..c668e80b70 --- /dev/null +++ b/tests/unit/dataset/graph/segment_ir/test_trie_content.py @@ -0,0 +1,275 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Driver-level tests for the shared trie-content core.""" + +from __future__ import annotations + +import pytest + +from aiperf.dataset.graph.segment_ir.pool import SegmentPool +from aiperf.dataset.graph.segment_ir.trie_content import ( + ReconCallbacks, + TrieISLMismatchError, + TrieNode, + TrieRequest, + build_trie_ir, + compute_asst_caps, + resolve_content_parents, +) + + +def _stub_callbacks() -> ReconCallbacks: + # Collision-free deterministic stubs: block tokens = [hash_id] * 4. + return ReconCallbacks( + decode_block_tokens=lambda hids: [t for h in hids for t in [h] * 4], + sample_partial_tail_tokens=lambda n, seed: [7] * n, + decode_tokens_to_text=lambda toks: " ".join(str(t) for t in toks), + ) + + +def _node( + nid: str, order: int, hashes: list[int], in_tok: int, out_tok: int, t: float +) -> TrieNode: + return TrieNode( + node_id=nid, + request=TrieRequest( + hash_ids=hashes, + input_length=in_tok, + output_length=out_tok, + t=t, + api_time=1.0, + ), + order=order, + ) + + +def test_extending_chain_prompt_paths_share_prefix_and_cover_isl() -> None: + nodes = [ + _node("a", 0, [1, 2], 8, 4, 0.0), + _node("b", 1, [1, 2, 3, 4], 16, 4, 5.0), + ] + pool = SegmentPool() + result = build_trie_ir( + nodes, + block_size=4, + callbacks=_stub_callbacks(), + pool=pool, + idle_gap_cap_seconds=None, + ) + p1 = result.builds["a"].prompt_path + p2 = result.builds["b"].prompt_path + assert p2[: len(p1)] == p1 + tok2 = sum(len(m["content"].split()) for m in pool.materialize(p2)) + assert tok2 == 16 # covered count == input_length (block-aligned), NOT inflated + + +def test_recorded_edge_delays_always_replay() -> None: + nodes = [ + _node("a", 0, [1], 4, 4, 0.0), + _node("b", 1, [1, 2], 8, 4, 100.0), + ] + result = build_trie_ir( + nodes, + block_size=4, + callbacks=_stub_callbacks(), + pool=SegmentPool(), + idle_gap_cap_seconds=None, + ) + delays = [ + e.delay_after_predecessor_us or 0.0 + for edges in result.edges_by_node.values() + for e in edges + ] + assert any(d > 0.0 for d in delays), ( + f"recorded end-to-start gaps must survive onto the edges; got {delays}" + ) + + +def test_isl_gate_rejects_decode_drift() -> None: + """The ISL gate must see the ACTUAL assembled token count: a decode + callback emitting 2x-block_size blocks (the W1 hazard, e.g. 64-token + default decode against a smaller recorded block size) must hard-abort the + build instead of shipping inflated prompts silently.""" + drifted = ReconCallbacks( + decode_block_tokens=lambda hids: [t for h in hids for t in [h] * 8], # 2x bs + sample_partial_tail_tokens=lambda n, seed: [7] * n, + decode_tokens_to_text=lambda toks: " ".join(str(t) for t in toks), + ) + nodes = [_node("a", 0, [1, 2], 8, 4, 0.0)] + with pytest.raises(TrieISLMismatchError): + build_trie_ir( + nodes, + block_size=4, + callbacks=drifted, + pool=SegmentPool(), + idle_gap_cap_seconds=None, + ) + + +def test_isl_gate_skipped_for_placeholder_callbacks() -> None: + """``block_exact=False`` (deliberate placeholder content, scheduling-only + timing-plane parse) opts out of the assembled-token gate; the same + mis-sized decode that aborts a content-bearing build must succeed.""" + placeholder = ReconCallbacks( + decode_block_tokens=lambda hids: [t for h in hids for t in [h] * 2], + sample_partial_tail_tokens=lambda n, seed: [7] * n, + decode_tokens_to_text=lambda toks: " ".join(str(t) for t in toks), + block_exact=False, + ) + nodes = [_node("a", 0, [1, 2], 8, 4, 0.0)] + result = build_trie_ir( + nodes, + block_size=4, + callbacks=placeholder, + pool=SegmentPool(), + idle_gap_cap_seconds=None, + ) + assert result.builds["a"].prompt_path + + +def test_compute_asst_caps_over_share_row_caps_frozen_tag_owner() -> None: + """On an over-share row (in // bs < lcp) the cap planner must clamp the + inherited count with the SAME three-way min ``assign_block_tags`` freezes + with, so the degenerate pull-back caps the owner of the re-exposed block + the tags actually see -- not the owner at the unclamped lcp boundary.""" + nodes = [ + _node("a", 0, [1], 2, 0, 0.0), # root; tiles [a] + _node("b", 1, [1, 2], 4, 0, 1.0), # tiles [a, b] + _node("c", 2, [1, 2, 3], 6, 0, 2.0), # tiles [a, b, c] + _node("d", 3, [1, 2, 3], 4, 0, 3.0), # over-share: in//2 = 2 < lcp 3 + ] + resolve_content_parents(nodes) + caps = compute_asst_caps(nodes, 2) + # d's degenerate pull-back re-exposes block index 1 (clamped boundary), + # owned by b -- the boundary the frozen tags land on. Unclamped, the + # planner would cap c (block index 2's owner) instead. + assert caps.get("b") == 0 + assert caps.get("c") is None + + +def test_small_prompt_fallback_emits_single_user_message() -> None: + nodes = [_node("tiny", 0, [], 3, 2, 0.0)] + pool = SegmentPool() + result = build_trie_ir( + nodes, + block_size=4, + callbacks=_stub_callbacks(), + pool=pool, + idle_gap_cap_seconds=None, + small_prompt_fallback=True, + ) + b = result.builds["tiny"] + assert b.small_prompt + msgs = pool.materialize(b.prompt_path) + assert len(msgs) == 1 and msgs[0]["role"] == "user" + assert len(msgs[0]["content"].split()) == 3 + + +def test_fragment_boundary_reuses_leading_whole_messages_reemits_straddler() -> None: + """A parent message straddling the child's (clamped) inherited boundary must + be re-emitted FRESH; only WHOLE parent messages inside the boundary reuse. + + ``c`` over-shares its content-parent ``b`` (``in // bs`` = 3 < lcp = 6), so + the three-way clamp lands the inherited boundary at block 3 -- strictly + inside ``b``'s second (assistant) message [2, 4). The rewrite must reuse + ``b``'s first whole message verbatim (identical content-addressed sid) and + re-emit the truncated straddling message fresh (a distinct sid, fewer + blocks). The property is byte-identical before and after the rewrite: shared + prefixes are content-addressed either way. + """ + nodes = [ + _node("a", 0, [1, 2], 8, 8, 0.0), # root -> one user message [0,1] + _node("b", 1, [1, 2, 3, 4, 5, 6], 24, 4, 1.0), # 3 msgs, ends [2,4,6] + _node("c", 2, [1, 2, 3, 4, 5, 6, 9, 10], 12, 4, 2.0), # over-share: covered 3 + ] + pool = SegmentPool() + result = build_trie_ir( + nodes, + block_size=4, + callbacks=_stub_callbacks(), + pool=pool, + idle_gap_cap_seconds=None, + ) + pa = result.builds["a"].prompt_path + pb = result.builds["b"].prompt_path + pc = result.builds["c"].prompt_path + # b reuses a's whole first message; c reuses that same first whole message. + assert pb[0] == pa[0] + assert pc[0] == pb[0] + # c covers only 3 blocks: one reused whole message + one fresh straddler. + assert len(pc) == 2 + # The straddling message is re-emitted fresh -- NOT b's full-width msg1. + assert pc[1] != pb[1] + # And the fresh straddler is exactly one block wide (block 2 only). + c_frag = pool.materialize([pc[1]])[0] + assert len(c_frag["content"].split()) == 4 # one block * block_size 4 + + +def test_reuse_boundary_uses_geometry_not_tag_prefix() -> None: + """The reuse boundary MUST be the geometric ``inherited``, never a tag-prefix + comparison: tags can coincide beyond the lcp while hash ids differ. + + ``a`` = [1,2,3,4] and ``b`` = [1,2,8,9] share only a 2-block hash prefix but + tag identically (all user), so a tag-prefix rule would wrongly splice ``a``'s + block-2 message into ``b``. Geometry stops reuse at block 2, so ``b``'s + block-2+ segment is a DISTINCT content-addressed sid. Holds before and after + the rewrite; the rewrite must not regress to a tag-prefix boundary. + """ + nodes = [ + _node("a", 0, [1, 2, 3, 4], 16, 4, 0.0), + _node("b", 1, [1, 2, 8, 9], 16, 4, 1.0), + ] + pool = SegmentPool() + result = build_trie_ir( + nodes, + block_size=4, + callbacks=_stub_callbacks(), + pool=pool, + idle_gap_cap_seconds=None, + ) + pa = result.builds["a"].prompt_path + pb = result.builds["b"].prompt_path + # No segment beyond the 2-block shared prefix may be shared: a and b are one + # user message each (all-user roots), so the whole-message sids must differ + # (blocks 3,4 vs 8,9 give distinct tokens -> distinct content-addressed ids). + assert set(pa).isdisjoint(set(pb)) + + +def test_deep_chain_decodes_each_shared_block_exactly_once() -> None: + """A deep extending chain decodes each covered block ONCE (linear), pinning + the death of the quadratic prefix re-emission. + + Counting stub: every ``decode_block_tokens`` call is recorded. With + prefix-path reuse each block is materialized exactly once at first occurrence + and spliced thereafter, so the total decode count equals the number of + DISTINCT covered blocks (the deepest node's covered count), not the per-node + re-decode sum ``n(n+1)/2``. This test FAILS on the pre-rewrite quadratic and + PASSES on the reuse rewrite -- the regression guard for the frozen file. + """ + n = 10 + nodes = [ + _node(f"x{i}", i, list(range(200, 200 + i + 1)), 4 * (i + 1), 4, float(i)) + for i in range(n) + ] + calls: list[int] = [] + + def _counting_decode(hids: list[int]) -> list[int]: + calls.append(len(hids)) + return [t for h in hids for t in [h] * 4] + + cb = ReconCallbacks( + decode_block_tokens=_counting_decode, + sample_partial_tail_tokens=lambda n, seed: [7] * n, + decode_tokens_to_text=lambda toks: " ".join(str(t) for t in toks), + ) + build_trie_ir( + nodes, + block_size=4, + callbacks=cb, + pool=SegmentPool(), + idle_gap_cap_seconds=None, + ) + assert sum(calls) == n, ( + f"decoded {sum(calls)} blocks; linear reuse decodes {n} " + f"(pre-rewrite quadratic decoded {n * (n + 1) // 2})" + ) diff --git a/tests/unit/dataset/graph/test_catalog_presence_predicate.py b/tests/unit/dataset/graph/test_catalog_presence_predicate.py new file mode 100644 index 0000000000..6e0f88bf5d --- /dev/null +++ b/tests/unit/dataset/graph/test_catalog_presence_predicate.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Catalog presence keys off the spliced ``__msgdelta`` node, not delta VALUES. + +Phase 2 strips message TEXT out of ``replay_outputs``; the per-trace node-ordinal +catalog must survive that, because its ordinals are the dispatch<->store contract +(a shift => GraphEnvelopeMissing). This proves emptying every trace's +``replay_outputs`` VALUES leaves the catalog byte-identical and non-empty. +""" + +from pathlib import Path + +import msgspec + +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.graph_path_catalog import build_graph_path_catalog + +WEKA_FIXTURE = Path(__file__).parents[2] / "graph" / "fixtures" / "weka_min.json" + +# The fixture is COMMITTED; a missing path is repo drift and must fail loudly, +# never silently vanish the test via a skipif. +assert WEKA_FIXTURE.exists(), f"committed weka fixture missing: {WEKA_FIXTURE}" + + +def test_catalog_survives_emptied_replay_outputs() -> None: + pg = from_weka_trace(WEKA_FIXTURE, content_root_seed=0) + before = build_graph_path_catalog(pg) + emptied_traces = [msgspec.structs.replace(t, replay_outputs={}) for t in pg.traces] + pg2 = msgspec.structs.replace(pg, traces=emptied_traces) + after = build_graph_path_catalog(pg2) + assert after == before + assert before # non-empty diff --git a/tests/unit/dataset/graph/test_edge_delay_exclusivity.py b/tests/unit/dataset/graph/test_edge_delay_exclusivity.py new file mode 100644 index 0000000000..f35036390b --- /dev/null +++ b/tests/unit/dataset/graph/test_edge_delay_exclusivity.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Rule 54: an edge's end-anchored and start-anchored delays are exclusive.""" + +from __future__ import annotations + +from aiperf.dataset.graph.models import ( + GraphRecord, + LlmNode, + ParsedGraph, + StaticEdge, + TraceRecord, +) +from aiperf.dataset.graph.validator import ( + _rule_54_edge_delay_exclusivity, + validate, +) + + +def graph_with(edge: StaticEdge) -> GraphRecord: + """Minimal two-LlmNode graph carrying a single edge under test.""" + nodes = { + "a": LlmNode(prompt=["hi"], output="ra"), + "b": LlmNode(prompt=["hi"], output="rb"), + } + return GraphRecord(nodes=nodes, edges=[edge], state={}) + + +def test_edge_with_both_delay_fields_rejected(): + """Rule 54: delay_after_predecessor_us and delay_after_predecessor_start_us + are mutually exclusive on one edge.""" + edge = StaticEdge( + source="a", + target="b", + delay_after_predecessor_us=1000.0, + delay_after_predecessor_start_us=1000.0, + ) + issues = _rule_54_edge_delay_exclusivity(graph_with(edge)) + assert issues, "both-set edge must produce a validation issue" + # T1-M1: the rule must also surface through the full validate() entrypoint, + # not just the direct rule call (adapter-tests-skip-validator trap). Mirror + # the Task 3 E2E test: wrap the graph in a ParsedGraph with a TraceRecord. + parsed = ParsedGraph(graph=graph_with(edge), traces=[TraceRecord(id="t")]) + assert any(i.rule_id == "rule-54" for i in validate(parsed)), ( + "rule-54 must surface end-to-end via validate()" + ) + + +def test_edge_with_single_delay_field_passes(): + for kwargs in ( + {"delay_after_predecessor_us": 1000.0}, + {"delay_after_predecessor_start_us": 1000.0}, + {}, + ): + edge = StaticEdge(source="a", target="b", **kwargs) + assert _rule_54_edge_delay_exclusivity(graph_with(edge)) == [] + + +def test_start_sourced_start_anchored_edge_rejected(): + """Rule 54: a START-sourced edge cannot be start-anchored -- the START + pseudo-node never dispatches, so the target would be silently orphaned.""" + edge = StaticEdge( + source="START", target="b", delay_after_predecessor_start_us=1000.0 + ) + issues = _rule_54_edge_delay_exclusivity(graph_with(edge)) + assert issues, "START-sourced start-anchored edge must produce an issue" + assert all(i.rule_id == "rule-54" for i in issues) + assert any("START" in i.message for i in issues) + + +def test_start_sourced_min_start_delay_only_passes(): + """A START-sourced edge with only min_start_delay_us is the correct way to + express an absolute offset from trace start; rule 54 must accept it.""" + edge = StaticEdge(source="START", target="b", min_start_delay_us=1000.0) + assert _rule_54_edge_delay_exclusivity(graph_with(edge)) == [] diff --git a/tests/unit/dataset/graph/test_first_token_anchor_rule.py b/tests/unit/dataset/graph/test_first_token_anchor_rule.py new file mode 100644 index 0000000000..85f1ed5bb8 --- /dev/null +++ b/tests/unit/dataset/graph/test_first_token_anchor_rule.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Rule 55: a first-token-anchored edge must carry its dispatch fallback and a +real source.""" + +from __future__ import annotations + +from aiperf.dataset.graph.models import ( + GraphRecord, + LlmNode, + ParsedGraph, + StaticEdge, + TraceRecord, +) +from aiperf.dataset.graph.validator import ( + _rule_55_first_token_anchor_shape, + validate, +) + + +def graph_with(edge: StaticEdge) -> GraphRecord: + """Minimal two-LlmNode graph carrying a single edge under test.""" + nodes = { + "a": LlmNode(prompt=["hi"], output="ra"), + "b": LlmNode(prompt=["hi"], output="rb"), + } + return GraphRecord(nodes=nodes, edges=[edge], state={}) + + +def _rule_55_issues_via_validate(edge: StaticEdge) -> list: + """Rule-55 issues produced by the FULL ``validate()`` entrypoint. + + Positive-path wiring proof: if the ``_rule_55...`` call were dropped from + ``validator.validate``, the rule-level assertions would still pass while + every violating graph sailed through the real entrypoint -- so the + rejection tests must also assert through here. + """ + parsed = ParsedGraph(graph=graph_with(edge), traces=[TraceRecord(id="t")]) + return [i for i in validate(parsed) if i.rule_id == "rule-55"] + + +def test_first_token_without_start_anchor_rejected(): + """Rule 55: a first-token anchor without its dispatch fallback + (delay_after_predecessor_start_us) is rejected.""" + edge = StaticEdge( + source="a", + target="b", + delay_after_predecessor_first_token_us=1000.0, + ) + assert _rule_55_first_token_anchor_shape(graph_with(edge)) + assert _rule_55_issues_via_validate(edge), ( + "violating edge must yield rule-55 through validate() (wiring proof)" + ) + + +def test_first_token_with_end_anchor_rejected(): + """Rule 55: a first-token anchor must not combine with the completion + anchor delay_after_predecessor_us.""" + edge = StaticEdge( + source="a", + target="b", + delay_after_predecessor_us=1000.0, + delay_after_predecessor_first_token_us=1000.0, + ) + assert _rule_55_first_token_anchor_shape(graph_with(edge)) + assert _rule_55_issues_via_validate(edge), ( + "violating edge must yield rule-55 through validate() (wiring proof)" + ) + + +def test_first_token_from_start_source_rejected(): + """Rule 55: a START-sourced edge cannot be first-token-anchored -- the + START pseudo-node never dispatches or streams a first token.""" + edge = StaticEdge( + source="START", + target="b", + delay_after_predecessor_start_us=2000.0, + delay_after_predecessor_first_token_us=1000.0, + ) + issues = _rule_55_first_token_anchor_shape(graph_with(edge)) + assert issues, "START-sourced first-token-anchored edge must produce an issue" + assert all(i.rule_id == "rule-55" for i in issues) + assert any("START" in i.message for i in issues) + validate_issues = _rule_55_issues_via_validate(edge) + assert validate_issues, ( + "violating edge must yield rule-55 through validate() (wiring proof)" + ) + assert any("START" in i.message for i in validate_issues) + + +def test_valid_first_token_edge_passes_full_validate(): + """A first-token anchor alongside its dispatch fallback on a real source is + the correct shape; rule 55 must accept it, end-to-end via validate().""" + edge = StaticEdge( + source="a", + target="b", + delay_after_predecessor_start_us=3000.0, + delay_after_predecessor_first_token_us=1000.0, + ) + assert _rule_55_first_token_anchor_shape(graph_with(edge)) == [] + # Mirror test_edge_delay_exclusivity.py: the rule must also stay silent + # through the full validate() entrypoint (adapter-tests-skip-validator trap). + parsed = ParsedGraph(graph=graph_with(edge), traces=[TraceRecord(id="t")]) + assert not any(i.rule_id == "rule-55" for i in validate(parsed)), ( + "valid first-token edge must not trip rule-55 via validate()" + ) diff --git a/tests/unit/dataset/graph/test_graph_meta_sidecar.py b/tests/unit/dataset/graph/test_graph_meta_sidecar.py new file mode 100644 index 0000000000..5f1665c0f3 --- /dev/null +++ b/tests/unit/dataset/graph/test_graph_meta_sidecar.py @@ -0,0 +1,62 @@ +import msgspec +import pytest + +from aiperf.dataset.graph.codecs import ( + GRAPH_META_SCHEMA_VERSION, + GRAPH_META_SIDECAR_FILENAME, + decode_graph_meta_sidecar, + encode_graph_meta_sidecar, +) +from aiperf.dataset.graph.models import GraphRecord, ParsedGraph, TraceRecord + + +def _tiny_graph() -> ParsedGraph: + return ParsedGraph( + graph=GraphRecord(), + traces=[TraceRecord(id="t-1", tags=["x"])], + ) + + +def test_sidecar_round_trips_graph_and_header(): + pg = _tiny_graph() + fp = {"kind": "file", "sha256": "abc", "size": 12} + blob = encode_graph_meta_sidecar(pg, source_fingerprint=fp, schema_version=1) + decoded, decoded_fp, version = decode_graph_meta_sidecar(blob) + assert version == 1 + assert decoded_fp == fp + assert [t.id for t in decoded.traces] == ["t-1"] + + +def test_encoder_writes_explicit_kind_and_decoder_requires_it() -> None: + pg = _tiny_graph() + frame = encode_graph_meta_sidecar(pg, source_fingerprint={"k": "v"}) + header, _blob = msgspec.msgpack.decode(frame) + assert header["kind"] == "parsed_graph" + assert header["schema_version"] == GRAPH_META_SCHEMA_VERSION + + # Kind-less frames (pre-v3 artifacts) are rejected -> caller re-parses. + del header["kind"] + stale = msgspec.msgpack.encode([header, _blob]) + with pytest.raises(ValueError, match="kind"): + decode_graph_meta_sidecar(stale) + + +def test_sidecar_filename_constant(): + assert GRAPH_META_SIDECAR_FILENAME == "graph_meta.msgpack" + + +def test_decode_rejects_garbage(): + with pytest.raises((msgspec.DecodeError, ValueError)): + decode_graph_meta_sidecar(b"\xff\xff not msgpack") + + +def test_decode_rejects_wrong_shape_frame(): + # Parseable msgpack, but not the [header, pg_bytes] frame shape. + with pytest.raises(ValueError): + decode_graph_meta_sidecar(msgspec.msgpack.encode("not-a-list")) + + +def test_decode_rejects_header_missing_keys(): + frame = msgspec.msgpack.encode([{"schema_version": 1}, b"ignored"]) + with pytest.raises(ValueError): + decode_graph_meta_sidecar(frame) diff --git a/tests/unit/dataset/graph/test_graph_meta_writer.py b/tests/unit/dataset/graph/test_graph_meta_writer.py new file mode 100644 index 0000000000..ea64cd2e05 --- /dev/null +++ b/tests/unit/dataset/graph/test_graph_meta_writer.py @@ -0,0 +1,91 @@ +from pathlib import Path + +from aiperf.dataset.graph.codecs import decode_graph_meta_sidecar +from aiperf.dataset.graph.graph_meta_sidecar import ( + catalogs_match, + sidecar_matches_index, + sidecar_path_for, + strip_replay_text, + write_graph_meta_sidecar, +) +from aiperf.dataset.graph.graph_path_catalog import build_catalog_context +from aiperf.dataset.graph.models import GraphRecord, LlmNode, ParsedGraph, TraceRecord + + +def _graph() -> ParsedGraph: + return ParsedGraph(graph=GraphRecord(), traces=[TraceRecord(id="t-1", tags=["x"])]) + + +def _graph_with_node() -> ParsedGraph: + """A graph whose catalog carries a REAL node ordinal. + + A node-free graph yields an empty per-trace catalog, which makes every + ``sidecar_matches_index`` comparison vacuously True -- the False branch + needs an ordinal that CAN go missing from the store index. + """ + graph = GraphRecord(nodes={"n0": LlmNode(prompt=["hi"], output="out")}, edges=[]) + return ParsedGraph(graph=graph, traces=[TraceRecord(id="t-1", tags=["x"])]) + + +def test_write_then_decode_round_trips(tmp_path: Path): + pg = _graph() + out = write_graph_meta_sidecar( + pg, + base_path=tmp_path, + benchmark_id="bench-9", + source_fingerprint={"kind": "file"}, + schema_version=1, + ) + assert out == sidecar_path_for(tmp_path, "bench-9") + assert out.exists() + decoded, fp, version = decode_graph_meta_sidecar(out.read_bytes()) + assert [t.id for t in decoded.traces] == ["t-1"] + assert fp == {"kind": "file"} and version == 1 + + +def test_catalogs_match_true_for_same_graph(): + pg = _graph() + catalog = build_catalog_context(pg).catalog + assert catalogs_match(pg, catalog) is True + + +def test_catalogs_match_false_for_divergent_catalog(): + pg = _graph() + assert catalogs_match(pg, {"ghost-trace": {"n": 0}}) is False + + +def test_sidecar_matches_index_true_when_store_covers_catalog(): + pg = _graph_with_node() + catalog = build_catalog_context(pg).catalog + assert catalog["t-1"], "fixture must yield real node ordinals" + # Build a fake index whose per-trace integer ordinals cover the catalog. + index = { + t: {(o, "profiling"): None for o in ords.values()} + for t, ords in catalog.items() + } + assert sidecar_matches_index(pg, index) is True + + +def test_sidecar_matches_index_false_when_catalog_ordinal_missing_from_store(): + """The False branch is the function's reason to exist: a catalog ordinal + absent from the store index means topology drift -> fall back to re-parse.""" + pg = _graph_with_node() + assert sidecar_matches_index(pg, {"t-1": {}}) is False + assert sidecar_matches_index(pg, {}) is False + + +def test_strip_replay_text_clears_only_replay_outputs(): + tr = TraceRecord( + id="t-1", + tags=["x"], + replay_outputs={ + "n__msgdelta": {"messages": [{"role": "user", "content": "hi"}]} + }, + ) + pg = ParsedGraph(graph=GraphRecord(), traces=[tr]) + stripped = strip_replay_text(pg) + assert stripped.traces[0].replay_outputs == {} + assert stripped.traces[0].id == "t-1" + assert stripped.traces[0].tags == ["x"] + # original must be untouched + assert pg.traces[0].replay_outputs != {} diff --git a/tests/unit/dataset/graph/test_graph_parse_context.py b/tests/unit/dataset/graph/test_graph_parse_context.py new file mode 100644 index 0000000000..3c562eafc3 --- /dev/null +++ b/tests/unit/dataset/graph/test_graph_parse_context.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``GraphParseContext`` carries the run's dataset-selection knobs. + +:func:`resolve_graph_parse_context` threads the run's default-dataset +``entries`` cap and ``synthesis.max_context_length`` into the context every +graph adapter parses through. Resolution reads the run config alone (outside +``DatasetResolver._resolve_one``), so neither the weka-HF ``org/name`` nor the +local-graph early-return in the dataset resolver can skip it -- these tests pin +that both knobs land on the resolved context, and that an unset ``entries`` +resolves to ``None`` (not a coalesced default). +""" + +from __future__ import annotations + +from pathlib import Path + +from aiperf.config.flags.cli_config import CLIConfig +from aiperf.dataset.graph.workload_detect import resolve_graph_parse_context +from tests.unit.conftest import make_run_from_cli + +# Any existing file makes the default dataset a FileDataset; the context +# resolver reads config only and never opens the path. +_GRAPH_FIXTURE = ( + Path(__file__).resolve().parent + / "adapters/fixtures/dynamo_nested/nested_2_level.jsonl.gz" +) + + +def test_resolve_graph_parse_context_carries_entries_and_max_context() -> None: + """Explicit ``entries`` + ``max_context_length`` reach the parse context.""" + run = make_run_from_cli( + CLIConfig( + model_names=["test-model"], + input_file=str(_GRAPH_FIXTURE), + tokenizer_name="builtin", + conversation_num_dataset_entries=50, + max_context_length=131072, + ) + ) + + ctx = resolve_graph_parse_context(run) + + assert ctx.num_dataset_entries == 50 + assert ctx.max_context_length == 131072 + + +def test_resolve_graph_parse_context_unset_entries_is_none() -> None: + """An unset ``entries`` resolves to ``None`` (no coalesced default).""" + run = make_run_from_cli( + CLIConfig( + model_names=["test-model"], + input_file=str(_GRAPH_FIXTURE), + tokenizer_name="builtin", + ) + ) + + dataset = run.cfg.get_default_dataset() + assert "entries" not in dataset.model_fields_set + + ctx = resolve_graph_parse_context(run) + + assert ctx.num_dataset_entries is None + assert ctx.max_context_length is None diff --git a/tests/unit/dataset/graph/test_native_lowering.py b/tests/unit/dataset/graph/test_native_lowering.py new file mode 100644 index 0000000000..0a4618713d --- /dev/null +++ b/tests/unit/dataset/graph/test_native_lowering.py @@ -0,0 +1,644 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Native-graph lowering onto the unified segment store (Phase 0: static only). + +Covers the canonicalization/interning happy paths (linear-chat shorthand, +explicit static prompts, per-trace init splices, content-block concatenation), +the NotImplementedError gates for un-lowerable constructs, and the +store-roundtrip parity: a lowered parse drained through +`build_unified_trie_store_interned` must materialize byte-correct messages. +""" + +from pathlib import Path + +import pytest +from pytest import param + +from aiperf.dataset.graph.models import ( + ChannelSpec, + ChannelType, + GraphRecord, + LlmNode, + ParsedGraph, + ReducerName, + StaticEdge, + TraceRecord, +) +from aiperf.dataset.graph.native_lowering import lower_native_to_unified +from aiperf.dataset.graph.parser import parse_native +from aiperf.dataset.graph.segment_ir.envelope import read_prompt_segment_ids +from aiperf.dataset.graph.segment_ir.store_builder import ( + build_unified_trie_store_interned, +) +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) + + +def _write_yaml(tmp_path: Path, text: str) -> Path: + p = tmp_path / "workload.yaml" + p.write_text(text) + return p + + +def _materialized(parsed: ParsedGraph, graph: GraphRecord, node_id: str) -> list[dict]: + node = graph.nodes[node_id] + assert isinstance(node, LlmNode) + path = read_prompt_segment_ids(node) + assert path is not None + assert parsed.segment_pool is not None + return parsed.segment_pool.materialize(path) + + +class TestNativeLoweringHappyPaths: + def test_linear_chat_shorthand_lowers_static_per_trace( + self, tmp_path: Path + ) -> None: + parsed = parse_native( + _write_yaml( + tmp_path, + """ +graph: + system: be brief +traces: + - id: t1 + messages: + - {role: user, content: hello} + - id: t2 + messages: + - {role: user, content: goodbye} +""", + ) + ) + assert parsed.segment_pool is not None + assert set(parsed.graphs) == {"t1", "t2"} + assert [t.graph_ref for t in parsed.traces] == ["t1", "t2"] + msgs = _materialized(parsed, parsed.graphs["t1"], "_llm") + assert msgs == [ + {"role": "system", "content": "be brief"}, + {"role": "user", "content": "hello"}, + ] + msgs2 = _materialized(parsed, parsed.graphs["t2"], "_llm") + assert msgs2[1] == {"role": "user", "content": "goodbye"} + + def test_static_prompts_share_one_graph(self, tmp_path: Path) -> None: + parsed = parse_native( + _write_yaml( + tmp_path, + """ +graph: + nodes: + a: + prompt: + - {role: user, content: "question one"} + output: a_out + b: + prompt: + - {role: user, content: "question two"} + output: b_out + edges: + - {source: START, target: a} + - {source: a, target: b} + - {source: b, target: END} +traces: + - id: t1 + - id: t2 +""", + ) + ) + assert parsed.segment_pool is not None + assert parsed.graphs == {} + assert all(t.graph_ref is None for t in parsed.traces) + assert _materialized(parsed, parsed.graph, "a") == [ + {"role": "user", "content": "question one"} + ] + + def test_content_block_concat_and_escape(self, tmp_path: Path) -> None: + parsed = parse_native( + _write_yaml( + tmp_path, + """ +graph: + nodes: + a: + prompt: + - role: user + content: ["Summarize: ", "@topic", " ", "@@literal"] + output: a_out + edges: + - {source: START, target: a} + - {source: a, target: END} +traces: + - id: t1 + initial_state: + topic: "graph stores" +""", + ) + ) + msgs = _materialized(parsed, parsed.graphs["t1"], "a") + assert msgs == [{"role": "user", "content": "Summarize: graph stores @literal"}] + + def test_trace_messages_shorthand_honored_with_explicit_nodes( + self, tmp_path: Path + ) -> None: + # G5: traces[].messages is documented as equivalent to + # initial_state.messages; it must be lifted for explicit-node graphs + # too, not only on the linear-chat synthesis path. + parsed = parse_native( + _write_yaml( + tmp_path, + """ +graph: + nodes: + a: + prompt: ["@messages"] + output: a_out + edges: + - {source: START, target: a} + - {source: a, target: END} +traces: + - id: t1 + messages: + - {role: user, content: authored} +""", + ) + ) + assert parsed.traces[0].messages is None + assert parsed.traces[0].initial_state["messages"] == [ + {"role": "user", "content": "authored"} + ] + assert _materialized(parsed, parsed.graphs["t1"], "a") == [ + {"role": "user", "content": "authored"} + ] + + def test_trace_messages_shorthand_defers_to_explicit_initial_state( + self, tmp_path: Path + ) -> None: + parsed = parse_native( + _write_yaml( + tmp_path, + """ +graph: + nodes: + a: + prompt: ["@messages"] + output: a_out + edges: + - {source: START, target: a} + - {source: a, target: END} +traces: + - id: t1 + messages: + - {role: user, content: shorthand} + initial_state: + messages: + - {role: user, content: explicit} +""", + ) + ) + assert _materialized(parsed, parsed.graphs["t1"], "a") == [ + {"role": "user", "content": "explicit"} + ] + + def test_self_written_channel_reads_init(self, tmp_path: Path) -> None: + parsed = parse_native( + _write_yaml( + tmp_path, + """ +graph: + state: + hist: {type: messages, reducer: add_messages} + nodes: + a: + prompt: ["@hist"] + output: hist + edges: + - {source: START, target: a} + - {source: a, target: END} +traces: + - id: t1 + initial_state: + hist: + - {role: user, content: hi} +""", + ) + ) + msgs = _materialized(parsed, parsed.graphs["t1"], "a") + assert msgs == [{"role": "user", "content": "hi"}] + + @pytest.mark.asyncio + async def test_store_roundtrip_materializes_canonical_messages( + self, tmp_path: Path + ) -> None: + parsed = parse_native( + _write_yaml( + tmp_path, + """ +graph: + system: sys +traces: + - id: t1 + messages: + - {role: user, content: hello} +""", + ) + ) + store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id="bench" + ) + catalog = await build_unified_trie_store_interned(parsed, store) + assert set(catalog) == {"t1"} + ordinal = catalog["t1"]["_llm"] + client = GraphSegmentUnifiedClient(tmp_path, "bench").open() + import orjson + + envelope = orjson.loads(client.get_node_envelope("t1", ordinal)) + assert client.materialize_handles(envelope["handles"]) == [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hello"}, + ] + + +class TestNativeLoweringGates: + @pytest.mark.parametrize( + ("yaml_text", "match"), + [ + param( + """ +graph: + nodes: + a: + prompt: [{role: user, content: q}] + output: shared + b: + prompt: ["@shared"] + output: b_out + edges: + - {source: START, target: a} + - {source: START, target: b} + - {source: a, target: END} + - {source: b, target: END} +traces: + - id: t1 +""", + "concurrent with", + id="non_ancestor_writer_gated", + ), + param( + """ +graph: + nodes: + a: + prompt: ["plain text at array level"] + output: a_out + edges: + - {source: START, target: a} + - {source: a, target: END} +traces: + - id: t1 +""", + "top-level string items", + id="literal_array_item", + ), + param( + """ +graph: + nodes: + a: + prompt: + - role: user + content: [{synth_tokens: 5}] + output: a_out + edges: + - {source: START, target: a} + - {source: a, target: END} +traces: + - id: t1 +""", + "non-string content blocks", + id="directive_block", + ), + param( + """ +graph: + nodes: + a: + prompt: ["@hist"] + output: a_out + edges: + - {source: START, target: a} + - {source: a, target: END} +traces: + - id: t1 +""", + "init-seeded content", + id="missing_init", + ), + param( + """ +graph: + nodes: + a: + prompt: [] + output: a_out + edges: + - {source: START, target: a} + - {source: a, target: END} +traces: + - id: t1 +""", + "empty prompt", + id="empty_prompt", + ), + param( + """ +graph: + nodes: + a: + prompt: ["@hist"] + output: a_out + edges: + - {source: START, target: a} + - {source: a, target: END} +traces: + - id: t1 + initial_state: + hist: "not a list" +""", + "list of message dicts", + id="messages_init_not_list", + ), + param( + """ +graph: + nodes: + a: + prompt: + - role: user + content: ["@topic"] + output: a_out + edges: + - {source: START, target: a} + - {source: a, target: END} +traces: + - id: t1 + initial_state: + topic: [not, a, string] +""", + "string initial_state value", + id="text_init_not_str", + ), + param( + """ +graph: + nodes: + a: + prompt: + - {role: user, content: q, name: alice} + output: a_out + edges: + - {source: START, target: a} + - {source: a, target: END} +traces: + - id: t1 +""", + "not representable in the unified store", + id="extra_message_keys", + ), + param( + """ +graph: + nodes: + a: + prompt: + - {role: user, content: q} + output: a_out + metadata: + replay_reducers: {a_out: overwrite} + edges: + - {source: START, target: a} + - {source: a, target: END} +traces: + - id: t1 +""", + "reducer overrides", + id="replay_reducers_metadata", + ), + param( + """ +graph: + nodes: + a: + prompt: ["@hist"] + output: a_out + edges: + - {source: START, target: a} + - {source: a, target: END} +traces: + - id: t1 + initial_state: + hist: + - {role: user, content: hi, id: m1} +""", + "not representable in the unified store", + id="extra_init_message_keys", + ), + param( + # G12: a rule-55-invalid edge (first-token anchor with no + # dispatch fallback) must fail loudly at lowering, not be + # silently lowered as a completion edge. + """ +graph: + nodes: + a: + prompt: [{role: user, content: q}] + output: a_out + b: + prompt: [{role: user, content: q}] + output: b_out + edges: + - {source: START, target: a} + - {source: a, target: b, delay_after_predecessor_first_token_us: 1000} + - {source: b, target: END} +traces: + - id: t1 +""", + r"graph\.edges\[a->b\].*sets no delay_after_predecessor_start_us", + id="first_token_anchor_without_fallback", + ), + ], + ) # fmt: skip + def test_parse_native_gates( + self, tmp_path: Path, yaml_text: str, match: str + ) -> None: + with pytest.raises(NotImplementedError, match=match): + parse_native(_write_yaml(tmp_path, yaml_text)) + + def test_retired_stream_reducer_rejected_at_parse(self, tmp_path: Path) -> None: + from aiperf.dataset.graph.decode import GraphDecodeError + + with pytest.raises(GraphDecodeError, match="stream_passthrough"): + parse_native( + _write_yaml( + tmp_path, + """ +graph: + state: + s: {type: text, reducer: stream_passthrough} + nodes: + a: + prompt: + - {role: user, content: q} + output: a_out +traces: + - id: t1 +""", + ) + ) + + def test_unknown_node_kind_rejected_at_decode(self) -> None: + from aiperf.dataset.graph.decode import GraphDecodeError, decode_node + + with pytest.raises(GraphDecodeError, match="unknown node node_type"): + decode_node({"node_type": "spawn", "graph_ref": "child"}) + + def test_conditional_edge_rejected_at_decode(self) -> None: + from aiperf.dataset.graph.decode import GraphDecodeError, decode_edge + + with pytest.raises(GraphDecodeError, match="branches"): + decode_edge({"source": "a", "branches": {"x": "END"}}) + + def test_subgraph_record_rejected_at_parse(self, tmp_path: Path) -> None: + from aiperf.dataset.graph.parser import GraphParseError + + path = tmp_path / "sub.jsonl" + path.write_text( + '{"kind":"graph","nodes":{}}\n' + '{"kind":"subgraph","name":"child","nodes":{}}\n' + ) + with pytest.raises(GraphParseError, match="unknown kind"): + parse_native(path) + + def test_refs_without_traces_gated(self) -> None: + node = LlmNode( + prompt=["@hist"], + output="a_out", + ) + graph = GraphRecord( + nodes={"a": node}, + state={ + "hist": ChannelSpec( + type=ChannelType.MESSAGES, reducer=ReducerName.ADD_MESSAGES + ) + }, + edges=[ + StaticEdge(source="START", target="a"), + StaticEdge(source="a", target="END"), + ], + ) + parsed = ParsedGraph(graph=graph, traces=[]) + with pytest.raises(NotImplementedError, match="no trace records"): + lower_native_to_unified(parsed) + + def test_duplicate_trace_ids_gated(self) -> None: + node = LlmNode(prompt=["@hist"], output="a_out") + graph = GraphRecord( + nodes={"a": node}, + edges=[ + StaticEdge(source="START", target="a"), + StaticEdge(source="a", target="END"), + ], + ) + init = {"hist": [{"role": "user", "content": "x"}]} + traces = [ + TraceRecord(id="t1", initial_state=init), + TraceRecord(id="t1", initial_state=init), + ] + parsed = ParsedGraph(graph=graph, traces=traces) + with pytest.raises(NotImplementedError, match="duplicate trace ids"): + lower_native_to_unified(parsed) + + +class TestSegmentIdDedup: + def test_shared_prefix_dedups_across_traces(self, tmp_path: Path) -> None: + parsed = parse_native( + _write_yaml( + tmp_path, + """ +graph: + system: same system +traces: + - id: t1 + messages: + - {role: user, content: alpha} + - id: t2 + messages: + - {role: user, content: beta} +""", + ) + ) + assert parsed.segment_pool is not None + p1 = read_prompt_segment_ids(parsed.graphs["t1"].nodes["_llm"]) + p2 = read_prompt_segment_ids(parsed.graphs["t2"].nodes["_llm"]) + assert p1 is not None and p2 is not None + assert p1[0] == p2[0] + assert p1[1] != p2[1] + + +class TestNativeModelFallback: + """Native nodes carry no per-node model; the run ``--model`` must be folded + into the wire body on both materialize paths, else the server 422s on a + missing ``model`` field (a Phase-0 gap the mock-server E2E surfaced).""" + + @pytest.mark.asyncio + async def test_run_model_folded_when_node_has_none(self, tmp_path: Path) -> None: + from aiperf.common.models import EndpointInfo + from aiperf.graph.worker_materialize import ( + materialize_graph_request_unified, + materialize_graph_request_unified_bytes, + ) + + parsed = parse_native( + _write_yaml( + tmp_path, + """ +graph: + system: sys +traces: + - id: t1 + messages: + - {role: user, content: hello} +""", + ) + ) + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id="m") + catalog = await build_unified_trie_store_interned(parsed, store) + ordinal = catalog["t1"]["_llm"] + client = GraphSegmentUnifiedClient(tmp_path, "m").open() + endpoint = EndpointInfo(type="chat", streaming=False, extra=[]) + + # dict path + payload = materialize_graph_request_unified( + client, "t1", ordinal, "profiling", default_model="run-model" + ) + assert payload["model"] == "run-model" + + # bytes path + import orjson + + body, model, _ = materialize_graph_request_unified_bytes( + client, + "t1", + ordinal, + "profiling", + endpoint=endpoint, + default_model="run-model", + ) + assert model == "run-model" + assert orjson.loads(body)["model"] == "run-model" + + # No default supplied => no model injected (weka/dynamo stamp their own). + bare = materialize_graph_request_unified(client, "t1", ordinal, "profiling") + assert "model" not in bare + client.close() diff --git a/tests/unit/dataset/graph/test_native_lowering_slots.py b/tests/unit/dataset/graph/test_native_lowering_slots.py new file mode 100644 index 0000000000..bea414ed06 --- /dev/null +++ b/tests/unit/dataset/graph/test_native_lowering_slots.py @@ -0,0 +1,537 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dynamic-content slots in the native lowering. + +Covers the slot composition rules: array-level splices expand to init +segments + completion-ordered writer slots (writers chained through reads); +block-level refs compose the single writer's value INTO the containing +message; producers get `capture`, readers get injected `ChannelRequirement`s; +and every legality gate (unordered writers, chain-read violations, anchored +producer edges, block multi-writer / init+writer) fails loudly. +""" + +from pathlib import Path + +import pytest +from pytest import param + +from aiperf.dataset.graph.models import LlmNode +from aiperf.dataset.graph.parser import parse_native +from aiperf.dataset.graph.segment_ir.envelope import read_prompt_segment_ids + + +def _parse(tmp_path: Path, text: str): + p = tmp_path / "workload.yaml" + p.write_text(text) + return parse_native(p) + + +def _trie(node: LlmNode) -> dict: + return (node.metadata or {}).get("trie") or {} + + +PLANNER_REVIEWER = """ +graph: + nodes: + plan: + prompt: [{role: user, content: "Make a plan."}] + output: plan_out + review: + prompt: + - role: user + content: ["Review this plan: ", "@plan_out"] + output: review_out + edges: + - {source: START, target: plan} + - {source: plan, target: review} + - {source: review, target: END} +traces: + - id: t1 +""" + + +class TestBlockSlots: + def test_planner_reviewer_composed_message(self, tmp_path: Path) -> None: + parsed = _parse(tmp_path, PLANNER_REVIEWER) + graph = parsed.graphs["t1"] + plan, review = graph.nodes["plan"], graph.nodes["review"] + + assert _trie(plan).get("capture") is True + assert "assembly" not in _trie(plan) + + assembly = _trie(review)["assembly"] + assert assembly == [ + { + "m": { + "role": "user", + "parts": [{"t": "Review this plan: "}, {"sv": "plan"}], + } + } + ] + assert read_prompt_segment_ids(review) == [] + assert [(r.channel, r.count) for r in review.inputs] == [("plan_out", 1)] + + @pytest.mark.parametrize( + ("yaml_text", "match"), + [ + param( + """ +graph: + nodes: + a: + prompt: [{role: user, content: qa}] + output: t_out + b: + prompt: [{role: user, content: qb}] + output: t_out + c: + prompt: + - role: user + content: ["@t_out"] + output: c_out + edges: + - {source: START, target: a} + - {source: a, target: b} + - {source: b, target: c} + - {source: c, target: END} +traces: + - id: t1 +""", + "multiple writers", + id="block_two_writers", + ), + param( + """ +graph: + nodes: + a: + prompt: [{role: user, content: qa}] + output: t_out + b: + prompt: + - role: user + content: ["@t_out"] + output: b_out + edges: + - {source: START, target: a} + - {source: a, target: b} + - {source: b, target: END} +traces: + - id: t1 + initial_state: + t_out: "seeded" +""", + "both init-seeded and written", + id="block_init_plus_writer", + ), + ], + ) # fmt: skip + def test_block_gates(self, tmp_path: Path, yaml_text: str, match: str) -> None: + with pytest.raises(NotImplementedError, match=match): + _parse(tmp_path, yaml_text) + + +ACCUMULATE_CHAIN = """ +graph: + state: + hist: {type: messages, reducer: add_messages} + nodes: + a: + prompt: ["@hist", {role: user, content: "turn one"}] + output: hist + b: + prompt: ["@hist", {role: user, content: "turn two"}] + output: hist + c: + prompt: ["@hist", {role: user, content: "turn three"}] + output: c_out + edges: + - {source: START, target: a} + - {source: a, target: b} + - {source: b, target: c} + - {source: c, target: END} +traces: + - id: t1 + initial_state: + hist: + - {role: system, content: "be brief"} +""" + + +class TestArraySlots: + def test_accumulate_chain_composition(self, tmp_path: Path) -> None: + parsed = _parse(tmp_path, ACCUMULATE_CHAIN) + graph = parsed.graphs["t1"] + a, b, c = graph.nodes["a"], graph.nodes["b"], graph.nodes["c"] + + # Producers referenced by a downstream slot are captured; c is a pure + # reader (writes c_out) so it is not. + assert _trie(a).get("capture") is True + assert _trie(b).get("capture") is True + assert "capture" not in _trie(c) + + # a reads @hist (= init only; b is a future/descendant writer it does + # not see), so a is fully STATIC — no slots, no assembly key. Its own + # prompt materializes [system, "turn one"]. + assert "assembly" not in _trie(a) + a_path = read_prompt_segment_ids(a) + assert parsed.segment_pool.materialize(a_path) == [ + {"role": "system", "content": "be brief"}, + {"role": "user", "content": "turn one"}, + ] + + # b reconstructs [init system, delta(a)="turn one", reply(a), + # delta(b)="turn two"]. + b_assembly = _trie(b)["assembly"] + assert b_assembly[0].keys() == {"seg"} # init: system message + assert b_assembly[1].keys() == {"seg"} # delta(a): "turn one" + assert b_assembly[2] == {"s": {"src": "a"}} # a's reply + assert b_assembly[3].keys() == {"seg"} # delta(b): "turn two" + assert [(r.channel, r.count) for r in b.inputs] == [("hist", 1)] + + # c reconstructs the full alternation of both prior turns. + c_assembly = _trie(c)["assembly"] + assert c_assembly[0].keys() == {"seg"} # init system + assert c_assembly[1].keys() == {"seg"} # delta(a) "turn one" + assert c_assembly[2] == {"s": {"src": "a"}} + assert c_assembly[3].keys() == {"seg"} # delta(b) "turn two" + assert c_assembly[4] == {"s": {"src": "b"}} + assert c_assembly[5].keys() == {"seg"} # delta(c) "turn three" + assert [(r.channel, r.count) for r in c.inputs] == [("hist", 2)] + + def test_first_writer_delta_is_its_whole_prompt(self, tmp_path: Path) -> None: + # No init: the root writer `a` need not read @hist (Gate B only applies + # to init-bearing channels). Its delta is its whole prompt. + parsed = _parse( + tmp_path, + """ +graph: + nodes: + a: + prompt: [{role: user, content: q}] + output: hist + b: + prompt: ["@hist", {role: user, content: next}] + output: b_out + edges: + - {source: START, target: a} + - {source: a, target: b} + - {source: b, target: END} +traces: + - id: t1 +""", + ) + assembly = _trie(parsed.graphs["t1"].nodes["b"])["assembly"] + assert assembly[0].keys() == {"seg"} # delta(a) = whole prompt "q" + assert assembly[1] == {"s": {"src": "a"}} # a's reply + assert assembly[2].keys() == {"seg"} # delta(b) = "next" + + def test_init_seeded_write_only_channel_not_gated(self, tmp_path: Path) -> None: + # G4: hist is init-seeded and written by `a`, but NOTHING reconstructs + # it (no downstream @hist reader with committed writers), so Gate B + # must not fire — the divergence it guards against cannot occur. + parsed = _parse( + tmp_path, + """ +graph: + state: + hist: {type: messages, reducer: add_messages} + nodes: + a: + prompt: [{role: user, content: "turn one"}] + output: hist + edges: + - {source: START, target: a} + - {source: a, target: END} +traces: + - id: t1 + initial_state: + hist: + - {role: system, content: "be brief"} +""", + ) + a = parsed.graph.nodes["a"] + assert parsed.segment_pool.materialize(read_prompt_segment_ids(a)) == [ + {"role": "user", "content": "turn one"} + ] + + def test_rewrite_own_draft_block_ref_not_gated(self, tmp_path: Path) -> None: + # G4 (review repro): a node block-refs @draft, writes draft, and the + # trace seeds it. No node reconstructs the channel, so the workload + # must lower (previously Gate B rejected it and the suggested fix + # dead-ended on the messages-splice list gate). + parsed = _parse( + tmp_path, + """ +graph: + nodes: + a: + prompt: + - role: user + content: ["Rewrite this draft: ", "@draft"] + output: draft + edges: + - {source: START, target: a} + - {source: a, target: END} +traces: + - id: t1 + initial_state: + draft: "first draft" +""", + ) + a = parsed.graphs["t1"].nodes["a"] + assert parsed.segment_pool.materialize(read_prompt_segment_ids(a)) == [ + {"role": "user", "content": "Rewrite this draft: first draft"} + ] + + def test_init_bearing_root_must_read_channel(self, tmp_path: Path) -> None: + # Gate B: hist is init-seeded and written by root `a`, but `a` doesn't + # read @hist -> rejected (it would dispatch without the seed). `b` + # reconstructs @hist (array slot on a committed writer), so the gate + # still applies after the G4 narrowing. + with pytest.raises(NotImplementedError, match="does not splice"): + _parse( + tmp_path, + """ +graph: + state: + hist: {type: messages, reducer: add_messages} + nodes: + a: + prompt: [{role: user, content: "turn one"}] + output: hist + b: + prompt: ["@hist", {role: user, content: "turn two"}] + output: b_out + edges: + - {source: START, target: a} + - {source: a, target: b} + - {source: b, target: END} +traces: + - id: t1 + initial_state: + hist: + - {role: system, content: "be brief"} +""", + ) + + def test_duplicate_channel_splice_gated(self, tmp_path: Path) -> None: + with pytest.raises(NotImplementedError, match="appears 2 times"): + _parse( + tmp_path, + """ +graph: + nodes: + a: {prompt: [{role: user, content: q}], output: hist} + b: {prompt: ["@hist", "@hist"], output: b_out} + edges: + - {source: START, target: a} + - {source: a, target: b} + - {source: b, target: END} +traces: + - id: t1 +""", + ) + + def test_reader_writer_nonleading_splice_gated(self, tmp_path: Path) -> None: + # b writes hist AND reads @hist, but @hist is not the first item. + with pytest.raises(NotImplementedError, match="must be the first prompt item"): + _parse( + tmp_path, + """ +graph: + nodes: + a: {prompt: [{role: user, content: q}], output: hist} + b: + prompt: [{role: user, content: pre}, "@hist"] + output: hist + c: {prompt: ["@hist"], output: c_out} + edges: + - {source: START, target: a} + - {source: a, target: b} + - {source: b, target: c} + - {source: c, target: END} +traces: + - id: t1 +""", + ) + + def test_root_writer_nonleading_splice_gated(self, tmp_path: Path) -> None: + # N6: w is the ROOT writer of init-seeded conv (no ancestor writers), + # so the reader-side Gate A never sees it; its non-leading "@conv" + # would survive to _delta_messages, which drops only a LEADING splice + # and would re-expand the init seed — r would reconstruct + # [SEED, SYS, SEED, Q1, ...], duplicating SEED and displacing SYS. + with pytest.raises(NotImplementedError, match="must be the first prompt item"): + _parse( + tmp_path, + """ +graph: + nodes: + w: + prompt: [{role: system, content: SYS}, "@conv", {role: user, content: Q1}] + output: conv + r: + prompt: ["@conv", {role: user, content: Q2}] + output: answer + edges: + - {source: START, target: w} + - {source: w, target: r} + - {source: r, target: END} +traces: + - id: t1 + initial_state: + conv: + - {role: user, content: SEED} +""", + ) + + def test_root_writer_leading_splice_no_seed_duplication( + self, tmp_path: Path + ) -> None: + # The leading-splice form of the graph above must still parse, and the + # reader's reconstruction must contain the init seed exactly once, in + # the same order the writer dispatched it. + parsed = _parse( + tmp_path, + """ +graph: + nodes: + w: + prompt: ["@conv", {role: system, content: SYS}, {role: user, content: Q1}] + output: conv + r: + prompt: ["@conv", {role: user, content: Q2}] + output: answer + edges: + - {source: START, target: w} + - {source: w, target: r} + - {source: r, target: END} +traces: + - id: t1 + initial_state: + conv: + - {role: user, content: SEED} +""", + ) + graph = parsed.graphs["t1"] + w, r = graph.nodes["w"], graph.nodes["r"] + assert parsed.segment_pool.materialize(read_prompt_segment_ids(w)) == [ + {"role": "user", "content": "SEED"}, + {"role": "system", "content": "SYS"}, + {"role": "user", "content": "Q1"}, + ] + assembly = _trie(r)["assembly"] + assert assembly[3] == {"s": {"src": "w"}} # w's live reply slot + assert parsed.segment_pool.materialize(read_prompt_segment_ids(r)) == [ + {"role": "user", "content": "SEED"}, + {"role": "system", "content": "SYS"}, + {"role": "user", "content": "Q1"}, + {"role": "user", "content": "Q2"}, + ] + + @pytest.mark.parametrize( + ("yaml_text", "match"), + [ + param( + """ +graph: + nodes: + a: {prompt: [{role: user, content: q}], output: fan} + b: {prompt: [{role: user, content: q}], output: hist} + c: {prompt: [{role: user, content: q}], output: hist} + d: {prompt: ["@hist"], output: d_out} + edges: + - {source: START, target: a} + - {source: a, target: b} + - {source: a, target: c} + - {source: b, target: d} + - {source: c, target: d} + - {source: d, target: END} +traces: + - id: t1 +""", + "mutually unordered", + id="parallel_writers_gated", + ), + param( + """ +graph: + nodes: + a: {prompt: [{role: user, content: q}], output: hist} + b: {prompt: [{role: user, content: q}], output: hist} + c: {prompt: ["@hist"], output: c_out} + edges: + - {source: START, target: a} + - {source: a, target: b} + - {source: b, target: c} + - {source: c, target: END} +traces: + - id: t1 +""", + "must chain through reads", + id="writer_chain_violation", + ), + ], + ) # fmt: skip + def test_array_gates(self, tmp_path: Path, yaml_text: str, match: str) -> None: + with pytest.raises(NotImplementedError, match=match): + _parse(tmp_path, yaml_text) + + +def test_anchored_producer_to_reader_edge_gated(tmp_path: Path) -> None: + with pytest.raises(NotImplementedError, match="contradictory timing intent"): + _parse( + tmp_path, + """ +graph: + nodes: + a: {prompt: [{role: user, content: q}], output: a_out} + x: {prompt: [{role: user, content: q}], output: x_out} + b: + prompt: + - role: user + content: ["@a_out"] + output: b_out + edges: + - {source: START, target: a} + - {source: a, target: x} + - {source: x, target: b} + - {source: a, target: b, delay_after_predecessor_start_us: 1000} + - {source: b, target: END} +traces: + - id: t1 +""", + ) + + +def test_tstar_gate_rejects_slot_workloads(tmp_path: Path) -> None: + from aiperf.dataset.graph.workload_detect import ( + _gate_dynamic_slots_vs_tstar, + ) + + parsed = _parse(tmp_path, PLANNER_REVIEWER) + + with pytest.raises(ValueError, match="t\\* snapshot window"): + _gate_dynamic_slots_vs_tstar(parsed, 0.75) + + _gate_dynamic_slots_vs_tstar(parsed, 0.0) + + +def test_static_workloads_pass_tstar_gate(tmp_path: Path) -> None: + from aiperf.dataset.graph.workload_detect import ( + _gate_dynamic_slots_vs_tstar, + ) + + parsed = _parse( + tmp_path, + """ +graph: + system: sys +traces: + - id: t1 + messages: + - {role: user, content: hello} +""", + ) + _gate_dynamic_slots_vs_tstar(parsed, 0.75) diff --git a/tests/unit/dataset/graph/test_parser_decode_strictness.py b/tests/unit/dataset/graph/test_parser_decode_strictness.py new file mode 100644 index 0000000000..7815da2ade --- /dev/null +++ b/tests/unit/dataset/graph/test_parser_decode_strictness.py @@ -0,0 +1,341 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Parser / decoder strictness for hand-authored native graph input. + +Locks the strictness contract: unknown top-level YAML +keys and unknown record fields fail loudly instead of silently dropping data; +a node with both ``prompt`` and ``messages`` is rejected; non-finite delay +values are rejected at decode time; the dead synth-token prompt fabrication +is a clear error; a non-dict ``provenance.extra`` gets file/loc context. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from pytest import param + +from aiperf.dataset.graph.codecs import ( + decode_parsed_graph_msgpack, + encode_parsed_graph_msgpack, +) +from aiperf.dataset.graph.decode import ( + GraphDecodeError, + decode_edge, + decode_graph, + decode_node, +) +from aiperf.dataset.graph.models import ( + GraphRecord, + LlmNode, + ParsedGraph, + StaticEdge, + TraceRecord, +) +from aiperf.dataset.graph.parser import GraphParseError, parse_native + + +def _write_yaml(tmp_path: Path, text: str) -> Path: + p = tmp_path / "workload.yaml" + p.write_text(text) + return p + + +class TestUnknownTopLevelKeys: + """G1: unknown top-level single-document keys must fail, not vanish.""" + + def test_typo_trace_singular_rejected_with_suggestion(self, tmp_path: Path) -> None: + with pytest.raises(GraphParseError, match=r"'trace'.*did you mean 'traces'"): + parse_native( + _write_yaml( + tmp_path, + """ +graph: + system: s +trace: + - id: t1 +""", + ) + ) + + def test_bare_unwrapped_graph_doc_rejected(self, tmp_path: Path) -> None: + # nodes:/edges: at top level (missing the graph: wrapper) must fail + # loudly rather than be discarded wholesale. + with pytest.raises(GraphParseError, match=r"unknown top-level key.*'nodes'"): + parse_native( + _write_yaml( + tmp_path, + """ +nodes: + a: + prompt: [{role: user, content: q}] + output: a_out +edges: + - {source: START, target: a} +""", + ) + ) + + def test_doc_with_none_of_the_known_sections_rejected(self, tmp_path: Path) -> None: + with pytest.raises(GraphParseError, match="contains none of"): + parse_native(_write_yaml(tmp_path, "{}\n")) + + @pytest.mark.parametrize( + "section", + [ + param("mix", id="mix"), + param("subgraphs", id="subgraphs"), + ], + ) # fmt: skip + def test_retired_mix_subgraphs_sections_rejected( + self, tmp_path: Path, section: str + ) -> None: + # mix:/subgraphs: were retired with their record kinds; they now fail + # at the top-level-key gate instead of expanding to a kind that + # _assemble rejects one step later. + with pytest.raises(GraphParseError, match=rf"unknown top-level key.*{section}"): + parse_native( + _write_yaml( + tmp_path, + f""" +graph: + system: s +{section}: + child: {{}} +""", + ) + ) + + def test_known_sections_still_parse(self, tmp_path: Path) -> None: + parsed = parse_native( + _write_yaml( + tmp_path, + """ +graph: + system: s +traces: + - id: t1 + messages: + - {role: user, content: hi} +""", + ) + ) + assert [t.id for t in parsed.traces] == ["t1"] + + +class TestUnknownRecordFields: + """G2: LlmNode/GraphRecord/TraceRecord forbid unknown fields like edges do.""" + + @pytest.mark.parametrize( + ("yaml_text", "match"), + [ + param( + """ +graph: + nodes: + a: + prompt: [{role: user, content: q}] + output: a_out + streming: false +traces: + - id: t1 +""", + "streming", + id="node_field_typo", + ), + param( + """ +graph: + verzion: "2.0" + nodes: + a: + prompt: [{role: user, content: q}] + output: a_out +traces: + - id: t1 +""", + "verzion", + id="graph_field_typo", + ), + param( + """ +graph: + nodes: + a: + prompt: [{role: user, content: q}] + output: a_out +traces: + - id: t1 + replay_output: {} +""", + "replay_output", + id="trace_field_typo", + ), + ], + ) # fmt: skip + def test_unknown_fields_rejected_at_parse( + self, tmp_path: Path, yaml_text: str, match: str + ) -> None: + with pytest.raises((GraphParseError, GraphDecodeError), match=match): + parse_native(_write_yaml(tmp_path, yaml_text)) + + def test_node_with_both_prompt_and_messages_rejected(self) -> None: + with pytest.raises(GraphDecodeError, match="both 'prompt' and 'messages'"): + decode_node( + { + "prompt": [{"role": "user", "content": "a"}], + "messages": [{"role": "user", "content": "b"}], + "output": "o", + }, + "n1", + ) + + def test_messages_alias_still_accepted_alone(self) -> None: + node = decode_node( + {"messages": [{"role": "user", "content": "a"}], "output": "o"} + ) + assert isinstance(node, LlmNode) + assert node.prompt == [{"role": "user", "content": "a"}] + + def test_msgpack_codec_round_trips_typed_structs(self) -> None: + # forbid_unknown_fields must not break the cross-process msgpack path + # (the encoder emits the node_type tag; the decoder must accept it). + node = LlmNode( + prompt=[{"role": "user", "content": "q"}], + output="o", + metadata={"trie": {"prompt_segment_ids": ["ab"]}}, + ) + pg = ParsedGraph( + graph=GraphRecord( + nodes={"a": node}, + edges=[ + StaticEdge(source="START", target="a"), + StaticEdge(source="a", target="END"), + ], + ), + traces=[TraceRecord(id="t1", replay_outputs={"a": {"o": "y"}})], + ) + assert decode_parsed_graph_msgpack(encode_parsed_graph_msgpack(pg)) == pg + + +class TestNonFiniteDelaysAtDecode: + """G3: +/-inf and NaN delay values must fail at decode, not hang at run.""" + + @pytest.mark.parametrize( + "field", + [ + param("delay_after_predecessor_us", id="completion"), + param("min_start_delay_us", id="min_start"), + param("delay_after_predecessor_start_us", id="start_anchor"), + param("delay_after_predecessor_first_token_us", id="first_token"), + ], + ) # fmt: skip + def test_inf_edge_delay_rejected(self, field: str) -> None: + raw: dict = {"source": "a", "target": "b", field: float("inf")} + if field == "delay_after_predecessor_first_token_us": + raw["delay_after_predecessor_start_us"] = 1.0 + with pytest.raises(GraphDecodeError, match=f"{field} must be finite"): + decode_edge(raw) + + def test_nan_edge_delay_rejected(self) -> None: + # NaN fails the msgspec ge=0 constraint; either way it must not decode. + with pytest.raises(GraphDecodeError): + decode_edge( + {"source": "a", "target": "b", "min_start_delay_us": float("nan")} + ) + + def test_inf_node_min_start_delay_rejected(self) -> None: + with pytest.raises(GraphDecodeError, match="min_start_delay_us must be finite"): + decode_node( + {"prompt": [], "output": "o", "min_start_delay_us": float("inf")}, + "n1", + ) + + def test_inf_delay_in_yaml_rejected_at_parse(self, tmp_path: Path) -> None: + with pytest.raises((GraphParseError, GraphDecodeError), match="must be finite"): + parse_native( + _write_yaml( + tmp_path, + """ +graph: + nodes: + a: + prompt: [{role: user, content: q}] + output: a_out + b: + prompt: [{role: user, content: q}] + output: b_out + edges: + - {source: START, target: a} + - {source: a, target: b, delay_after_predecessor_us: .inf} + - {source: b, target: END} +traces: + - id: t1 +""", + ) + ) + + def test_finite_delays_still_decode(self) -> None: + edge = decode_edge( + {"source": "a", "target": "b", "delay_after_predecessor_us": 1500.0} + ) + assert edge.delay_after_predecessor_us == 1500.0 + + +class TestSynthTokenFabricationRemoved: + """G6: prompt-less nodes with expected.input_tokens fail clearly.""" + + def test_promptless_node_with_expected_input_tokens_errors(self) -> None: + with pytest.raises( + GraphDecodeError, + match=( + r"graph\.nodes\.n1: node has no prompt; synth-token fabrication " + r"is not supported" + ), + ): + decode_node({"expected": {"input_tokens": 128}, "output": "o"}, "n1") + + def test_promptless_node_via_parse_names_the_node(self, tmp_path: Path) -> None: + with pytest.raises( + (GraphParseError, GraphDecodeError), match=r"graph\.nodes\.a.*no prompt" + ): + parse_native( + _write_yaml( + tmp_path, + """ +graph: + nodes: + a: + expected: {input_tokens: 64} + output: a_out +traces: + - id: t1 +""", + ) + ) + + +class TestProvenanceExtraCoercion: + """G10: a non-dict provenance.extra fails with loc context, not a raw + ValueError/TypeError.""" + + @pytest.mark.parametrize( + "bad_extra", + [ + param("a string", id="str"), + param(["vendor", "keys"], id="list"), + param(7, id="int"), + ], + ) # fmt: skip + def test_non_dict_extra_rejected_with_loc(self, bad_extra: object) -> None: + with pytest.raises( + GraphDecodeError, match=r"graph\.provenance\.extra: must be a mapping" + ): + decode_graph({"provenance": {"source": "weka_trace", "extra": bad_extra}}) + + def test_vendor_keys_still_fold_into_extra(self) -> None: + graph = decode_graph( + {"provenance": {"source": "weka_trace", "tool": "x/1", "vendor_key": 3}} + ) + assert graph.provenance.extra == {"vendor_key": 3} diff --git a/tests/unit/dataset/graph/test_parser_detect.py b/tests/unit/dataset/graph/test_parser_detect.py new file mode 100644 index 0000000000..f47ed198c6 --- /dev/null +++ b/tests/unit/dataset/graph/test_parser_detect.py @@ -0,0 +1,281 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Format-detection tests for Dynamo agent-trace JSONL/JSONL.gz/segmented dirs.""" + +from __future__ import annotations + +import gzip +import json +from pathlib import Path +from typing import Any + +import pytest + +from aiperf.dataset.graph.parser import ( + WorkloadFormatError, + detect_format, +) +from aiperf.plugin import plugins +from aiperf.plugin.enums import PluginType + + +def _dynamo_record( + *, + event_type: str = "request_end", + agent_context: dict[str, Any] | None = None, + **overrides: Any, +) -> dict[str, Any]: + rec: dict[str, Any] = { + "schema": "dynamo.request.trace.v1", + "event_type": event_type, + "agent_context": {} if agent_context is None else agent_context, + "ts": "2026-05-06T00:00:00.000Z", + } + rec.update(overrides) + return rec + + +def _write_jsonl(path: Path, recs: list[dict[str, Any]]) -> None: + path.write_text("\n".join(json.dumps(r) for r in recs) + "\n") + + +def _write_jsonl_gz(path: Path, recs: list[dict[str, Any]]) -> None: + payload = ("\n".join(json.dumps(r) for r in recs) + "\n").encode() + path.write_bytes(gzip.compress(payload)) + + +def test_graph_adapter_registry_includes_dynamo_trace() -> None: + assert plugins.has_entry(PluginType.GRAPH_ADAPTER, "dynamo_trace") + + +def test_detects_dynamo_trace_jsonl(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl" + _write_jsonl(p, [_dynamo_record(), _dynamo_record(event_type="tool_start")]) + assert detect_format(p) == "dynamo_trace" + + +def test_detects_dynamo_trace_jsonl_gz(tmp_path: Path) -> None: + p = tmp_path / "trace.jsonl.gz" + _write_jsonl_gz( + p, + [ + _dynamo_record(), + _dynamo_record(event_type="tool_start"), + _dynamo_record(event_type="tool_end"), + ], + ) + assert detect_format(p) == "dynamo_trace" + + +def test_detects_segmented_dynamo_trace_dir(tmp_path: Path) -> None: + d = tmp_path / "agent-traces" + d.mkdir() + _write_jsonl_gz(d / "agent.000000.jsonl.gz", [_dynamo_record()]) + _write_jsonl_gz( + d / "agent.000001.jsonl.gz", [_dynamo_record(event_type="tool_start")] + ) + assert detect_format(d) == "dynamo_trace" + + +def test_detects_plain_jsonl_dir(tmp_path: Path) -> None: + """Un-gzipped capture dirs are parseable (discover_segments takes *.jsonl), + so detection must claim them too.""" + d = tmp_path / "agent-traces-plain" + d.mkdir() + _write_jsonl(d / "a.jsonl", [_dynamo_record()]) + _write_jsonl(d / "b.jsonl", [_dynamo_record(event_type="tool_end")]) + assert detect_format(d) == "dynamo_trace" + + +def test_plain_jsonl_dir_with_foreign_records_not_claimed(tmp_path: Path) -> None: + """The dir branch still sniffs the first record: foreign .jsonl dirs fall + through instead of being claimed by extension alone.""" + d = tmp_path / "not-dynamo" + d.mkdir() + _write_jsonl(d / "a.jsonl", [{"schema": "other.v1", "event_type": "request_end"}]) + with pytest.raises(WorkloadFormatError): + detect_format(d) + + +# valid 10-byte gzip header + non-deflate payload -> zlib.error on read. +_GZ_HEADER_PLUS_GARBAGE = ( + b"\x1f\x8b\x08\x00\x00\x00\x00\x00\x00\x03" + b"this-is-not-deflate-data" +) + + +def test_can_load_corrupt_gzip_returns_false(tmp_path: Path) -> None: + """zlib.error from corrupt deflate bytes means "not ours", never a crash.""" + from aiperf.dataset.graph.adapters.dynamo.trace import DynamoTraceAdapter + + p = tmp_path / "trace.jsonl.gz" + p.write_bytes(_GZ_HEADER_PLUS_GARBAGE) + assert DynamoTraceAdapter.can_load(p) is False + + +def test_can_load_truncated_gzip_returns_false(tmp_path: Path) -> None: + """A gz truncated before the first line completes must not crash sniffing.""" + from aiperf.dataset.graph.adapters.dynamo.trace import DynamoTraceAdapter + + p = tmp_path / "trace.jsonl.gz" + _write_jsonl_gz(p, [_dynamo_record()]) + p.write_bytes(p.read_bytes()[:20]) # header + a few deflate bytes + assert DynamoTraceAdapter.can_load(p) is False + + +def test_detect_format_contains_adapter_can_load_crash( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """One adapter's can_load raising must not abort detection for the rest.""" + from aiperf.dataset.graph.adapters.dynamo.trace import DynamoTraceAdapter + + def _boom(cls: type, path: Path) -> bool: + raise RuntimeError("sniff crashed") + + monkeypatch.setattr(DynamoTraceAdapter, "can_load", classmethod(_boom)) + p = tmp_path / "some.jsonl" + _write_jsonl(p, [{"kind": "note", "text": "not a dynamo record"}]) + assert detect_format(p) == "native" + + +def test_empty_agent_context_still_routes_dynamo(tmp_path: Path) -> None: + """The predicate just checks isinstance(dict); empty dict is fine.""" + p = tmp_path / "trace.jsonl" + _write_jsonl(p, [_dynamo_record(agent_context={})]) + assert detect_format(p) == "dynamo_trace" + + +def test_wrong_schema_falls_through(tmp_path: Path) -> None: + """A different schema string should not match dynamo_trace.""" + p = tmp_path / "other.jsonl" + _write_jsonl( + p, + [ + { + "schema": "other.schema.v1", + "event_type": "request_end", + "agent_context": {}, + } + ], + ) + # Falls through every JSONL predicate and lands on "native". + assert detect_format(p) == "native" + + +# A genuine minimal weka trace object: the exact discriminator key set weka's +# file sniff requires ({id, models, block_size, hash_id_scope, requests}). +_WEKA_TRACE_DOC: dict[str, Any] = { + "id": "trace_ordering_guard", + "models": ["m"], + "block_size": 64, + "hash_id_scope": "local", + "requests": [{"t": 0.0, "type": "n", "model": "m", "in": 10, "out": 5}], +} + + +def test_dynamo_predicate_does_not_eat_weka_trace(tmp_path: Path) -> None: + """Detection ordering: a genuine weka trace routes to weka_trace. + + dynamo_trace has the HIGHER detection_priority (100 vs weka's 85), so if + its predicate ever claimed a weka trace file it would win the tie-break + and silently reroute the workload. Lock both halves: dynamo's can_load + rejects the file, and full registry detection resolves it to weka_trace. + """ + from aiperf.dataset.graph.adapters.dynamo.trace import DynamoTraceAdapter + + p = tmp_path / "trace_01.json" + p.write_text(json.dumps(_WEKA_TRACE_DOC)) + assert DynamoTraceAdapter.can_load(p) is False + assert detect_format(p) == "weka_trace" + + +def test_dynamo_predicate_does_not_eat_native_graph_jsonl(tmp_path: Path) -> None: + """Detection ordering: a native graph .jsonl doc falls through to native. + + Native graph workloads share dynamo's ``.jsonl`` extension; dynamo's + higher-priority predicate must reject the doc so detection lands on the + lowest-priority native fallback instead of hijacking it. + """ + from aiperf.dataset.graph.adapters.dynamo.trace import DynamoTraceAdapter + + p = tmp_path / "workload.jsonl" + _write_jsonl( + p, + [ + { + "graph": { + "nodes": { + "a": { + "node_type": "llm", + "prompt": [{"role": "user", "content": "hi"}], + "output": "out", + } + }, + "edges": [], + }, + "traces": [{"id": "t1"}], + } + ], + ) + assert DynamoTraceAdapter.can_load(p) is False + assert detect_format(p) == "native" + + +def test_unknown_event_type_falls_through(tmp_path: Path) -> None: + """Right schema but wrong event_type — does not match dynamo_trace.""" + p = tmp_path / "trace.jsonl" + _write_jsonl( + p, + [ + { + "schema": "dynamo.request.trace.v1", + "event_type": "stream_chunk", # not in the known set + "agent_context": {}, + } + ], + ) + assert detect_format(p) == "native" + + +def test_unknown_directory_raises(tmp_path: Path) -> None: + d = tmp_path / "mystery-dir" + d.mkdir() + (d / "unrelated.txt").write_text("hello") + with pytest.raises(WorkloadFormatError): + detect_format(d) + + +def test_parser_dispatch_routes_to_from_dynamo_trace(tmp_path: Path) -> None: + """Ensure parser.parse_graph dispatches dynamo_trace → from_dynamo_trace. + + Gated on the adapter module's existence — parallel agents are building + `dynamo/trace.py`. When their commits land on the merge target this test + will execute against the real adapter; until then it skips at import time. + """ + pytest.importorskip("aiperf.dataset.graph.adapters.dynamo.trace") + + from aiperf.dataset.graph import parser as parser_mod + + p = tmp_path / "trace.jsonl" + _write_jsonl(p, [_dynamo_record()]) + + called: dict[str, Any] = {} + + def _fake_from_dynamo_trace(path: Path, **kwargs): + called["path"] = Path(path) + from aiperf.dataset.graph.models import GraphRecord, ParsedGraph + + return ParsedGraph( + graph=GraphRecord(), + traces=[], + ) + + import aiperf.dataset.graph.adapters.dynamo.trace as dt_mod + + monkey_orig = dt_mod.from_dynamo_trace # type: ignore[attr-defined] + dt_mod.from_dynamo_trace = _fake_from_dynamo_trace # type: ignore[attr-defined] + try: + parser_mod.parse_graph(p) + finally: + dt_mod.from_dynamo_trace = monkey_orig # type: ignore[attr-defined] + + assert called["path"] == p diff --git a/tests/unit/dataset/graph/test_validator_rules.py b/tests/unit/dataset/graph/test_validator_rules.py new file mode 100644 index 0000000000..7c078cb47c --- /dev/null +++ b/tests/unit/dataset/graph/test_validator_rules.py @@ -0,0 +1,252 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Validator rule coverage. + +Covers: rule-1 iterative cycle detection at corpus scale (no RecursionError), +``validate()`` visiting every ``parsed.graphs`` entry, rule-56 dangling edge +endpoints, rule-57 non-finite delay values, and the rule-13 default-"manual" +provenance warning for adapter-emitted graphs. +""" + +from __future__ import annotations + +import pytest +from pytest import param + +from aiperf.dataset.graph.models import ( + GraphRecord, + LlmNode, + ParsedGraph, + ProvenanceSpec, + StaticEdge, + TraceRecord, +) +from aiperf.dataset.graph.validator import ( + ValidationSeverity, + _rule_01_cycles, + _rule_13_provenance_tool, + _rule_56_edge_endpoints, + _rule_57_finite_delays, + validate, +) + + +def _node() -> LlmNode: + return LlmNode(prompt=[{"role": "user", "content": "q"}], output="o") + + +def _chain_graph(n: int) -> GraphRecord: + nodes = {f"n{i}": _node() for i in range(n)} + edges = [StaticEdge(source="START", target="n0")] + edges += [StaticEdge(source=f"n{i}", target=f"n{i + 1}") for i in range(n - 1)] + edges.append(StaticEdge(source=f"n{n - 1}", target="END")) + return GraphRecord(nodes=nodes, edges=edges) + + +class TestRule01IterativeCycles: + """G8: rule-1 must not recurse -- recorded corpora exceed 100k-node chains.""" + + def test_100k_node_chain_validates_without_recursion_error(self) -> None: + assert _rule_01_cycles(_chain_graph(100_000)) == [] + + def test_long_cycle_still_detected(self) -> None: + n = 5_000 + graph = _chain_graph(n) + edges = [*graph.edges, StaticEdge(source=f"n{n - 1}", target="n0")] + import msgspec + + issues = _rule_01_cycles(msgspec.structs.replace(graph, edges=edges)) + assert [i.rule_id for i in issues] == ["rule-1"] + + def test_small_cycle_detected(self) -> None: + graph = GraphRecord( + nodes={"a": _node(), "b": _node()}, + edges=[ + StaticEdge(source="START", target="a"), + StaticEdge(source="a", target="b"), + StaticEdge(source="b", target="a"), + ], + ) + assert [i.rule_id for i in _rule_01_cycles(graph)] == ["rule-1"] + + def test_diamond_is_not_a_cycle(self) -> None: + graph = GraphRecord( + nodes={k: _node() for k in ("a", "b", "c", "d")}, + edges=[ + StaticEdge(source="START", target="a"), + StaticEdge(source="a", target="b"), + StaticEdge(source="a", target="c"), + StaticEdge(source="b", target="d"), + StaticEdge(source="c", target="d"), + StaticEdge(source="d", target="END"), + ], + ) + assert _rule_01_cycles(graph) == [] + + +class TestValidateVisitsAllGraphs: + """G7: validate() must run every rule over parsed.graphs values too.""" + + def test_cycle_in_secondary_graph_reported_with_graph_name(self) -> None: + clean = _chain_graph(2) + cyclic = GraphRecord( + nodes={"a": _node(), "b": _node()}, + edges=[ + StaticEdge(source="START", target="a"), + StaticEdge(source="a", target="b"), + StaticEdge(source="b", target="a"), + ], + ) + parsed = ParsedGraph( + graph=clean, + graphs={"t1": clean, "t2": cyclic}, + traces=[ + TraceRecord(id="t1", graph_ref="t1"), + TraceRecord(id="t2", graph_ref="t2"), + ], + ) + cycle_issues = [i for i in validate(parsed) if i.rule_id == "rule-1"] + assert cycle_issues, "cycle in a parsed.graphs entry must surface" + assert all(i.location.startswith("graphs[t2]") for i in cycle_issues) + + def test_aliased_main_graph_not_double_reported(self) -> None: + # The native lowering aliases parsed.graph to the first graphs entry; + # its issues must not be duplicated. + cyclic = GraphRecord( + nodes={"a": _node()}, + edges=[ + StaticEdge(source="START", target="a"), + StaticEdge(source="a", target="a"), + ], + ) + parsed = ParsedGraph( + graph=cyclic, + graphs={"t1": cyclic}, + traces=[TraceRecord(id="t1", graph_ref="t1")], + ) + assert len([i for i in validate(parsed) if i.rule_id == "rule-1"]) == 1 + + def test_new_rules_run_via_validate(self) -> None: + graph = GraphRecord( + nodes={"a": _node()}, + edges=[ + StaticEdge(source="START", target="a"), + StaticEdge(source="a", target="ghost"), + StaticEdge( + source="a", target="END", delay_after_predecessor_us=float("inf") + ), + ], + ) + parsed = ParsedGraph(graph=graph, traces=[TraceRecord(id="t")]) + rule_ids = {i.rule_id for i in validate(parsed)} + assert {"rule-56", "rule-57"} <= rule_ids + + +class TestRule56EdgeEndpoints: + """G9: edges must reference declared nodes or the matching sentinel.""" + + @pytest.mark.parametrize( + ("source", "target", "bad"), + [ + param("START", "ghost", "ghost", id="dangling_target"), + param("ghost", "END", "ghost", id="dangling_source"), + param("END", "a", "END", id="end_as_source"), + param("a", "START", "START", id="start_as_target"), + ], + ) # fmt: skip + def test_dangling_endpoint_flagged( + self, source: str, target: str, bad: str + ) -> None: + graph = GraphRecord( + nodes={"a": _node()}, + edges=[StaticEdge(source=source, target=target)], + ) + issues = _rule_56_edge_endpoints(graph) + assert [i.rule_id for i in issues] == ["rule-56"] + assert bad in issues[0].message + + def test_declared_endpoints_and_sentinels_pass(self) -> None: + graph = GraphRecord( + nodes={"a": _node(), "b": _node()}, + edges=[ + StaticEdge(source="START", target="a"), + StaticEdge(source="a", target="b"), + StaticEdge(source="b", target="END"), + ], + ) + assert _rule_56_edge_endpoints(graph) == [] + + +class TestRule57FiniteDelays: + """G3 (validator half): already-decoded graphs with non-finite delays are + caught even though typed construction bypasses decode.py.""" + + @pytest.mark.parametrize( + "field", + [ + param("delay_after_predecessor_us", id="completion"), + param("min_start_delay_us", id="min_start"), + param("delay_after_predecessor_start_us", id="start_anchor"), + ], + ) # fmt: skip + def test_inf_edge_delay_flagged(self, field: str) -> None: + graph = GraphRecord( + nodes={"a": _node(), "b": _node()}, + edges=[StaticEdge(source="a", target="b", **{field: float("inf")})], + ) + issues = _rule_57_finite_delays(graph) + assert [i.rule_id for i in issues] == ["rule-57"] + assert issues[0].location.endswith(field) + + def test_inf_node_min_start_delay_flagged(self) -> None: + node = LlmNode( + prompt=[{"role": "user", "content": "q"}], + output="o", + min_start_delay_us=float("inf"), + ) + graph = GraphRecord(nodes={"a": node}, edges=[]) + issues = _rule_57_finite_delays(graph) + assert [i.rule_id for i in issues] == ["rule-57"] + assert issues[0].location == "graph.nodes.a.min_start_delay_us" + + def test_finite_delays_pass(self) -> None: + graph = GraphRecord( + nodes={"a": _node(), "b": _node()}, + edges=[ + StaticEdge( + source="a", + target="b", + delay_after_predecessor_us=1.0, + min_start_delay_us=2.0, + ) + ], + ) + assert _rule_57_finite_delays(graph) == [] + + +class TestRule13DefaultManualWarning: + """G11: an adapter-emitted graph still carrying the default tool 'manual' + is treated as unstamped and warned about.""" + + def test_adapter_source_with_default_manual_warns(self) -> None: + graph = GraphRecord( + provenance=ProvenanceSpec(source="weka_trace") # tool defaults 'manual' + ) + issues = _rule_13_provenance_tool(graph) + assert [i.rule_id for i in issues] == ["rule-13"] + assert issues[0].severity is ValidationSeverity.WARNING + + def test_adapter_source_with_stamped_tool_passes(self) -> None: + graph = GraphRecord( + provenance=ProvenanceSpec(source="weka_trace", tool="aiperf-weka-trie/1") + ) + assert _rule_13_provenance_tool(graph) == [] + + def test_adapter_source_with_empty_tool_still_errors(self) -> None: + graph = GraphRecord(provenance=ProvenanceSpec(source="weka_trace", tool=" ")) + issues = _rule_13_provenance_tool(graph) + assert [i.rule_id for i in issues] == ["rule-13"] + assert issues[0].severity is ValidationSeverity.ERROR + + def test_hand_authored_with_manual_tool_passes(self) -> None: + assert _rule_13_provenance_tool(GraphRecord()) == [] diff --git a/tests/unit/dataset/graph/test_workload_detect_path_predicate.py b/tests/unit/dataset/graph/test_workload_detect_path_predicate.py new file mode 100644 index 0000000000..269958f10a --- /dev/null +++ b/tests/unit/dataset/graph/test_workload_detect_path_predicate.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Path-level graph-workload predicate tests. + +:func:`is_graph_workload_path` is the path-level companion to +:func:`resolve_graph_workload` (which takes a run). It uses the SAME registry +detection (``_detect_graph_workload_format``, which excludes ``native``), so a +local trace file (dynamo ``.jsonl.gz``) is recognized while a plain conversation +``.jsonl`` is not. +""" + +from __future__ import annotations + +from pathlib import Path + +from aiperf.dataset.graph.workload_detect import is_graph_workload_path + +_DYNAMO_FIXTURE = ( + Path(__file__).resolve().parent + / "adapters/fixtures/dynamo_nested/nested_2_level.jsonl.gz" +) + + +def test_is_graph_workload_path_true_for_dynamo_fixture() -> None: + assert is_graph_workload_path(_DYNAMO_FIXTURE) is True + + +def test_is_graph_workload_path_false_for_plain_jsonl(tmp_path: Path) -> None: + plain = tmp_path / "plain.jsonl" + plain.write_text('{"messages": [{"role": "user", "content": "hi"}]}\n') + assert is_graph_workload_path(plain) is False diff --git a/tests/unit/dataset/loader/test_exgentic.py b/tests/unit/dataset/loader/test_exgentic.py index b6c4ddb86f..f2b6513848 100644 --- a/tests/unit/dataset/loader/test_exgentic.py +++ b/tests/unit/dataset/loader/test_exgentic.py @@ -217,10 +217,6 @@ async def test_convert_preserves_snapshots_tools_osl_order_and_delays() -> None: }, } ] - assert all( - turn.extra_headers == {"x-dynamo-session-id": "session-1"} - for turn in conversation.turns - ) assert all( "recorded output" not in orjson.dumps(turn.raw_messages).decode() for turn in conversation.turns @@ -265,10 +261,6 @@ async def test_fixed_schedule_preserves_overlapping_start_times() -> None: ] assert all(len(conversation.turns) == 1 for conversation in conversations) assert all(conversation.turns[0].delay is None for conversation in conversations) - assert all( - conversation.turns[0].extra_headers == {"x-dynamo-session-id": "session-1"} - for conversation in conversations - ) @pytest.mark.asyncio diff --git a/tests/unit/dataset/test_dag_jsonl_streaming_store_parity.py b/tests/unit/dataset/test_dag_jsonl_streaming_store_parity.py new file mode 100644 index 0000000000..9f5b55b57e --- /dev/null +++ b/tests/unit/dataset/test_dag_jsonl_streaming_store_parity.py @@ -0,0 +1,273 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""dag_jsonl store-build parity: a slot-free dag graph's streamed payload +drain must equal the interned build byte-for-byte (the byte-parity oracle that +proves the in-process interned drain and the weka worker-pool trie drain build +identical stores), and ``_build_graph_store_streaming`` now routes EVERY +non-weka dag_jsonl parse -- slot-carrying or not -- through the in-process +interned drain (parse once, drain the same parse, return the full parse). +``graph_carries_assembly_slots`` is retained for the ``workload_detect`` +t*-gate; its lineage detection is pinned here too.""" + +from __future__ import annotations + +from pathlib import Path +from types import MethodType, SimpleNamespace + +import orjson +import pytest + +from aiperf.dataset.graph.adapters.dag_jsonl.trace import from_dag_jsonl +from aiperf.dataset.graph.graph_meta_sidecar import strip_replay_text +from aiperf.dataset.graph.merge import merge_parsed_graphs +from aiperf.dataset.graph.models import resolve_trace_graph +from aiperf.dataset.graph.segment_ir.store_builder import ( + build_unified_trie_store_from_payloads, + build_unified_trie_store_interned, + graph_carries_assembly_slots, + iter_trace_segment_payloads, +) +from aiperf.dataset.graph.store_build import GraphStoreBuilder +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) + +DAG_FIXTURES = Path(__file__).parents[2] / "fixtures" / "dag" +# Spawn-only dag graphs are the slot-free shape: spawned children start fresh +# histories, so no node assembles a live-reply (lineage) slot. Fork children +# inherit the parent's history including its assistant reply, so every other +# fixture in the dir carries slots. +SLOT_FREE_FIXTURE = DAG_FIXTURES / "spawn_minimal.dag.jsonl" +LINEAGE_FIXTURE = DAG_FIXTURES / "full.dag.jsonl" + +# Two INDEPENDENT root sessions (neither spawns the other) with distinct +# topologies: conv-a is a lone turn, conv-b spawns a child. One file therefore +# parses to ONE ParsedGraph carrying TWO traces whose graphs live in +# ``parsed.graphs`` keyed by graph_ref -- the multi-trace-per-source shape the +# structural merge must preserve. +MULTI_TRACE_DAG_LINES = """\ +{"session_id":"conv-a","turns":[{"model":"Qwen3-0.6B","messages":[{"role":"system","content":"a-sys"},{"role":"user","content":"a-u"}],"max_tokens":20}]} +{"session_id":"conv-b","turns":[{"model":"Qwen3-0.6B","messages":[{"role":"system","content":"b-sys"},{"role":"user","content":"b-u"}],"max_tokens":20,"spawns":["conv-b-child"]}]} +{"session_id":"conv-b-child","turns":[{"model":"Qwen3-0.6B","messages":[{"role":"system","content":"bc-sys"},{"role":"user","content":"bc-u"}],"max_tokens":20}]} +""" + + +@pytest.mark.asyncio +async def test_dag_jsonl_slot_free_payload_stream_matches_interned_oracle( + tmp_path: Path, +) -> None: + parsed = from_dag_jsonl(str(SLOT_FREE_FIXTURE)) + assert parsed.segment_pool is not None + # Slot-free premise: this fixture must stream with no eager fallback. + assert not graph_carries_assembly_slots(parsed) + + eager_store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id="dag-eager" + ) + eager_catalog = await build_unified_trie_store_interned(parsed, eager_store) + + stream_store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id="dag-stream" + ) + stream_catalog = await build_unified_trie_store_from_payloads( + iter_trace_segment_payloads(parsed), stream_store + ) + + assert stream_catalog == eager_catalog and eager_catalog + + # Strong equivalence: the persisted unified store is byte-for-byte identical + # on disk (mirrors ``test_dynamo_streaming_store_parity``'s eager-vs-stream + # oracle), which subsumes the interned handle map, node-manifest region, and + # content pool -- the whole store, not just the materialized messages. + eager_dir = tmp_path / "aiperf_graph_segments_dag-eager" + stream_dir = tmp_path / "aiperf_graph_segments_dag-stream" + eager_files = sorted(p.name for p in eager_dir.iterdir()) + stream_files = sorted(p.name for p in stream_dir.iterdir()) + assert eager_files == stream_files and eager_files, ( + f"unified store file sets differ: {eager_files} vs {stream_files}" + ) + for name in eager_files: + assert (eager_dir / name).read_bytes() == (stream_dir / name).read_bytes(), ( + f"unified store file {name!r} differs between streaming and eager" + ) + + # Semantic equivalence through the worker read face: the byte-identical + # stores materialize identical content and agree on the non-handle envelope + # fields (handles are store-local ints; compare MATERIALIZED content). + with ( + GraphSegmentUnifiedClient(tmp_path, "dag-eager").open() as ec, + GraphSegmentUnifiedClient(tmp_path, "dag-stream").open() as sc, + ): + for trace_id, ordinals in eager_catalog.items(): + for ordinal in ordinals.values(): + e_raw = ec.get_node_envelope(trace_id, ordinal, "profiling") + s_raw = sc.get_node_envelope(trace_id, ordinal, "profiling") + assert e_raw is not None and s_raw is not None + e_env = orjson.loads(e_raw) + s_env = orjson.loads(s_raw) + assert ec.materialize_handles( + e_env["handles"] + ) == sc.materialize_handles(s_env["handles"]) + assert {k: v for k, v in e_env.items() if k != "handles"} == { + k: v for k, v in s_env.items() if k != "handles" + } + + +def test_dag_jsonl_lineage_fixture_carries_slots() -> None: + """Real live-reply dag workloads carry assembly slots. + + Fork children replay the parent's assistant response as a live-reply + assembly slot, so a lineage-carrying dag parse must trip + ``graph_carries_assembly_slots``. The store route does not branch on + this (every non-weka format takes the interned drain, which persists slot + envelopes), but ``workload_detect``'s t*-gate still consults the detector, + so its lineage detection is pinned here. + """ + parsed = from_dag_jsonl(str(LINEAGE_FIXTURE)) + assert parsed.segment_pool is not None + assert graph_carries_assembly_slots(parsed) + + +@pytest.mark.asyncio +async def test_configure_route_dag_jsonl_slot_free_takes_interned_drain( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A slot-free dag_jsonl workload takes the in-process interned drain. + + Routing pin: ``_build_graph_store_streaming`` serves dag_jsonl (like every + non-weka format) by parsing once in-process and draining that SAME parse + through the interned builder. The second return IS the full parse + (identity), NOT a content-free structural merge, and the weka trie + drain/structural merge are never touched. + """ + from aiperf.dataset.graph import workload_detect + + parsed = from_dag_jsonl(str(SLOT_FREE_FIXTURE)) + assert not graph_carries_assembly_slots(parsed) + monkeypatch.setattr( + workload_detect, "parse_graph_workload", lambda run, path: parsed + ) + + # Just the attributes the in-process interned branch reads from self, with + # the REAL interned-drain/sidecar helpers bound; the weka trie drain and + # structural merge fail loudly if reached (no non-weka format may take them). + stub = SimpleNamespace( + run=SimpleNamespace(benchmark_id="bench"), + info=lambda *a, **k: None, + warning=lambda *a, **k: None, + _sidecar_path=None, + ) + for name in ("_write_graph_sidecar", "_build_interned_unified_store"): + setattr(stub, name, MethodType(getattr(GraphStoreBuilder, name), stub)) + + async def _fail_trie(payloads, base_path): # noqa: ANN001, ARG001 + raise AssertionError("dag_jsonl must not take the weka trie payload drain") + + def _fail_merge(structural_sink): # noqa: ANN001, ARG001 + raise AssertionError("dag_jsonl must not merge a structural stream") + + stub._build_graph_store_streaming_trie = _fail_trie + stub._merge_structural_graphs = _fail_merge + + catalog, returned = await GraphStoreBuilder._build_graph_store_streaming( + stub, SLOT_FREE_FIXTURE, tmp_path, "dag_jsonl" + ) + + assert set(catalog) == {t.id for t in parsed.traces} and catalog + # The interned drain hands back the FULL parse, not a merged structural graph. + assert returned is parsed + # The interned store is the real unified store the worker opens. + with GraphSegmentUnifiedClient(tmp_path, "bench").open() as client: + for trace_id, ordinals in catalog.items(): + for ordinal in ordinals.values(): + assert client.get_node_envelope(trace_id, ordinal, "profiling") + + +def test_merge_preserves_per_trace_graphs_multi_trace_source( + tmp_path: Path, +) -> None: + """A multi-trace SOURCE graph keeps each trace's own topology through merge. + + The structural merge used to assume one trace per source and keyed every + trace to ``pg.graph`` (the FIRST tree's record), dropping ``pg.graphs`` + entirely -- so conv-b resolved to conv-a's topology and the mandatory + sidecar's ``catalogs_match`` gate hard-failed the build. + """ + src = tmp_path / "multi.dag.jsonl" + src.write_text(MULTI_TRACE_DAG_LINES) + parsed = from_dag_jsonl(str(src)) + assert len(parsed.traces) == 2 + # Premise: the two traces genuinely carry DIFFERENT topologies. + node_sets = {t.id: set(resolve_trace_graph(parsed, t).nodes) for t in parsed.traces} + assert node_sets["conv-a"] != node_sets["conv-b"] + + merged = merge_parsed_graphs([strip_replay_text(parsed)]) + + assert {t.id for t in merged.traces} == set(node_sets) + for trace in merged.traces: + assert set(resolve_trace_graph(merged, trace).nodes) == node_sets[trace.id], ( + f"trace {trace.id!r} lost its own topology through the merge" + ) + + +@pytest.mark.asyncio +async def test_configure_route_dag_jsonl_multi_trace_slot_free_takes_interned_drain( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A multi-trace slot-free dag_jsonl file builds through the interned route. + + Regression pin: the mandatory ``_write_graph_sidecar`` cross-checks the + build catalog against the structural graph (``catalogs_match``). The two + traces carry DIFFERENT topologies, so a route that collapsed them to one + tree's graph (the historical merge bug) would fail that gate. The in-process + interned drain writes the sidecar from the FULL parse directly, so the build + must SUCCEED with both traces on their own topologies. + """ + from aiperf.dataset.graph import workload_detect + + src = tmp_path / "multi.dag.jsonl" + src.write_text(MULTI_TRACE_DAG_LINES) + parsed = from_dag_jsonl(str(src)) + assert len(parsed.traces) == 2 + assert not graph_carries_assembly_slots(parsed) + # Premise: the two traces genuinely carry DIFFERENT topologies. + node_sets = {t.id: set(resolve_trace_graph(parsed, t).nodes) for t in parsed.traces} + assert node_sets["conv-a"] != node_sets["conv-b"] + monkeypatch.setattr( + workload_detect, "parse_graph_workload", lambda run, path: parsed + ) + + stub = SimpleNamespace( + run=SimpleNamespace(benchmark_id="bench-multi"), + info=lambda *a, **k: None, + warning=lambda *a, **k: None, + _sidecar_path=None, + ) + for name in ("_write_graph_sidecar", "_build_interned_unified_store"): + setattr(stub, name, MethodType(getattr(GraphStoreBuilder, name), stub)) + + async def _fail_trie(payloads, base_path): # noqa: ANN001, ARG001 + raise AssertionError("dag_jsonl must not take the weka trie payload drain") + + stub._build_graph_store_streaming_trie = _fail_trie + + catalog, returned = await GraphStoreBuilder._build_graph_store_streaming( + stub, src, tmp_path, "dag_jsonl" + ) + + assert set(catalog) == {t.id for t in parsed.traces} and len(catalog) == 2 + # The interned drain returns the FULL parse. The per-trace topology is + # cross-checked STORE-SIDE (against the persisted index) in + # ``test_nonweka_interned_route`` rather than re-derived from ``returned is + # parsed`` here (which would be tautological -- ``node_sets`` came from the + # same parse). + assert returned is parsed + # The mandatory sidecar's catalogs_match gate passed (path recorded), so the + # multi-topology build was NOT collapsed to a single tree's graph. + assert stub._sidecar_path is not None + # The store carries manifests for BOTH traces' real nodes. + with GraphSegmentUnifiedClient(tmp_path, "bench-multi").open() as client: + for trace_id, ordinals in catalog.items(): + for ordinal in ordinals.values(): + assert client.get_node_envelope(trace_id, ordinal, "profiling") diff --git a/tests/unit/dataset/test_dynamo_direct_store_route.py b/tests/unit/dataset/test_dynamo_direct_store_route.py new file mode 100644 index 0000000000..d5c1f556c7 --- /dev/null +++ b/tests/unit/dataset/test_dynamo_direct_store_route.py @@ -0,0 +1,387 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dynamo direct write-through store route. + +Covers the route face the byte-parity oracle does not: + +* the ``StoreBackedSegmentPool`` shim -- ``add()`` dedups identically to + ``SegmentPool.add`` and every non-``add`` method fails loud (dynamo-only); +* the ``**adapter_kwargs`` seam -- ``direct_store`` threads + ``parse_graph_workload -> parse_graph -> _parse_via_adapter -> adapter.parse`` + with the two documented fail-loud modes (``GraphParseError`` for native, + ``TypeError`` for an adapter that does not accept the kwarg); +* the store-before-parse cleanup -- a mid-parse failure (a block-inconsistent + record, or any exception once content has spilled) leaves NO store files, i.e. + the abort()+rmtree cleanup composes with the parse, not just the + drain; +* the ``GraphStoreBuilder`` route on a REAL resolved ``BenchmarkRun`` (not a + MagicMock -- the documented path-drift trap): the dynamo route hands a live + store to the parse and the produced store + facet match the eager route. +""" + +from __future__ import annotations + +from pathlib import Path +from types import MethodType, SimpleNamespace + +import orjson +import pytest + +from aiperf.config.flags.cli_config import CLIConfig +from aiperf.dataset.graph import workload_detect +from aiperf.dataset.graph.adapters.dynamo.store_backed_pool import ( + InterningSegmentPool, + StoreBackedSegmentPool, +) +from aiperf.dataset.graph.codecs import GRAPH_META_SCHEMA_VERSION +from aiperf.dataset.graph.graph_meta_sidecar import write_graph_meta_sidecar +from aiperf.dataset.graph.parser import GraphParseError, parse_graph +from aiperf.dataset.graph.segment_ir.pool import SegmentPool, segment_id +from aiperf.dataset.graph.segment_ir.store_builder import ( + build_unified_trie_store_interned, +) +from aiperf.dataset.graph.store_build import GraphStoreBuilder +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, +) +from tests.unit.conftest import make_run_from_cli + + +def _dynamo_record(ts: int, sid: str, input_tokens: int, hashes: list[int]) -> dict: + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": {"session_id": sid}, + "request": { + "request_id": f"r{ts}", + "model": "m", + "input_tokens": input_tokens, + "output_tokens": 8, + "cached_tokens": 0, + "replay": { + "trace_block_size": 16, + "input_length": input_tokens, + "input_sequence_hashes": hashes, + }, + }, + } + + +@pytest.fixture +def dyn_trace(tmp_path: Path) -> Path: + p = tmp_path / "dyn_route.jsonl" + records = [ + _dynamo_record(1000, "s1", 32, [111, 222]), + _dynamo_record(2000, "s1", 64, [111, 222, 333, 444]), + _dynamo_record(3000, "s2", 48, [555, 666, 777]), + ] + p.write_bytes(b"\n".join(orjson.dumps(r) for r in records)) + return p + + +# --- StoreBackedSegmentPool shim ------------------------------------------------ + + +_ADD_CALLS = [ + {"role": "user", "content": "one", "tokens": [1, 2, 3], "parent_id": None}, + {"role": "assistant", "content": "two", "tokens": [4, 5], "parent_id": "p"}, + # A verbatim repeat of the first call: first-occurrence dedup must return the + # SAME sid and intern nothing new. + {"role": "user", "content": "one", "tokens": [1, 2, 3], "parent_id": None}, +] + + +def test_shim_add_dedups_identically_to_segmentpool(tmp_path: Path) -> None: + """The shim's ``add()`` returns the SAME sid stream as ``SegmentPool.add`` -- + same content-addressed ids, same first-occurrence dedup -- and interns exactly + the unique segments into the store (the repeat is a no-op).""" + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id="shim-dd") + shim = StoreBackedSegmentPool(store) + pool = SegmentPool() + + shim_sids = [shim.add(**call) for call in _ADD_CALLS] + pool_sids = [pool.add(**call) for call in _ADD_CALLS] + + assert shim_sids == pool_sids + # First-occurrence dedup: the repeated call yields the first call's sid. + assert shim_sids[0] == shim_sids[2] + # The store interned exactly the two UNIQUE segments (dedup no-op on repeat), + # matching the plain pool's two-entry _by_id. + assert len(store._ids) == 2 + assert len(pool.by_id) == 2 + + +def test_interning_pool_add_returns_canonical_sid_object(tmp_path: Path) -> None: + """``InterningSegmentPool.add`` returns the FIRST-BORN ``Segment.id`` object on + a dedup hit (not a fresh equal string), while values and ``_by_id`` insertion + order stay identical to a plain ``SegmentPool``. + + The eager route relies on this canonical-object return so every node's + ``prompt_segment_ids`` re-listing shares one string object per unique value. + """ + interning = InterningSegmentPool() + plain = SegmentPool() + + interning_sids = [interning.add(**call) for call in _ADD_CALLS] + plain_sids = [plain.add(**call) for call in _ADD_CALLS] + + # Values are byte-identical to the plain pool (interning changes identity only). + assert interning_sids == plain_sids + assert list(interning.by_id.keys()) == list(plain.by_id.keys()) + # The repeat (call 2 == call 0) returns the SAME object as the first add... + assert interning_sids[0] is interning_sids[2] + # ...which is exactly the first-born Segment.id stored on first occurrence. + assert all(sid is interning.by_id[sid].id for sid in interning_sids) + # A plain SegmentPool does NOT canonicalize -- the dedup hit is a fresh object. + assert plain_sids[0] is not plain_sids[2] + + +def test_shim_add_returns_canonical_sid_object(tmp_path: Path) -> None: + """``StoreBackedSegmentPool.add`` returns the canonical (first-born) sid object + on a repeat via its handle-indexed ``_sids`` list, values unchanged, and + ``_sids`` is dense (one entry per unique segment, in first-occurrence order).""" + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id="shim-can") + shim = StoreBackedSegmentPool(store) + plain = SegmentPool() + + shim_sids = [shim.add(**call) for call in _ADD_CALLS] + plain_sids = [plain.add(**call) for call in _ADD_CALLS] + + # Values identical to the plain pool's sid stream (identity-only change). + assert shim_sids == plain_sids + # The repeat returns the SAME object as the first add (canonical). + assert shim_sids[0] is shim_sids[2] + # _sids is dense: exactly the two unique sids, in first-occurrence order, and + # holding the canonical objects the shim returns. + assert shim._sids == [shim_sids[0], shim_sids[1]] + assert shim._sids[0] is shim_sids[0] + assert len(shim._sids) == len(store._ids) == 2 + + +def test_shim_add_pre_populated_store_falls_through_to_fresh_sid( + tmp_path: Path, +) -> None: + """When the store was written before this shim existed (``handle`` outruns + ``_sids``), ``add`` degrades to returning the fresh value-correct sid and + interns NOTHING into ``_sids`` -- the defensive third branch, never hit in + production (the shim is the store's sole writer during a parse).""" + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id="shim-pre") + # Pre-populate the store with two segments the shim never recorded in _sids. + store.put_segment("pre-a", "user", "a") + store.put_segment("pre-b", "user", "b") + + shim = StoreBackedSegmentPool(store) + assert shim._sids == [] + + got = shim.add(role="user", content="one", tokens=[1, 2, 3], parent_id=None) + + # The handle (2) outran the empty _sids, so add fell through: value is still the + # correct content-addressed sid, and _sids stayed empty (nothing canonicalized). + assert got == segment_id(None, "user", [1, 2, 3]) + assert shim._sids == [] + + +def test_shim_non_add_methods_fail_loud(tmp_path: Path) -> None: + """Every non-``add`` pool method raises ``NotImplementedError`` naming the + dynamo-only write-through contract -- so a non-dynamo adopter cannot silently + intern into an empty ``_by_id`` the store never sees.""" + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id="shim-fl") + shim = StoreBackedSegmentPool(store) + + with pytest.raises(NotImplementedError, match="add_text"): + shim.add_text(role="user", content="x", parent_id=None) + with pytest.raises(NotImplementedError, match="add_raw_message"): + shim.add_raw_message(message={"role": "user"}, parent_id=None) + with pytest.raises(NotImplementedError, match="get"): + shim.get("some-sid") + with pytest.raises(NotImplementedError, match="materialize"): + shim.materialize(["some-sid"]) + + +# --- **adapter_kwargs seam: fail-loud contract ---------------------------------- + + +def test_parse_graph_native_rejects_adapter_kwargs(tmp_path: Path) -> None: + """``adapter_kwargs`` with ``format="native"`` raise ``GraphParseError`` up + front: the native parse path has no adapter ``parse`` to receive them.""" + p = tmp_path / "native.yaml" + p.write_text("graph:\n nodes: {}\n") + with pytest.raises(GraphParseError, match="native parse path accepts no"): + parse_graph(p, format="native", direct_store=object()) + + +def test_parse_graph_unknown_adapter_kwarg_fails_loud_typeerror( + tmp_path: Path, +) -> None: + """A kwarg reaching an adapter whose ``parse`` does not accept it fails loud + with ``TypeError`` (never silently dropped). weka's ``parse(path, ctx)`` takes + no ``direct_store``, so threading one is a build-plane wiring bug that must + surface, not vanish.""" + p = tmp_path / "t.json" + p.write_text("{}") + with pytest.raises(TypeError): + parse_graph(p, format="weka_trace", direct_store=object()) + + +# --- store-before-parse cleanup ------------------------------------------------- + + +def _cleanup_stub(benchmark_id: str) -> SimpleNamespace: + """Stub self carrying only what the dynamo direct branch reads. + + Binds the REAL interned drain + sidecar helpers so the branch runs + end-to-end; ``parse_graph_workload`` is monkeypatched per-test. + """ + stub = SimpleNamespace( + run=SimpleNamespace(benchmark_id=benchmark_id), + info=lambda *a, **k: None, + warning=lambda *a, **k: None, + _sidecar_path=None, + ) + for name in ("_write_graph_sidecar", "_build_interned_unified_store"): + setattr(stub, name, MethodType(getattr(GraphStoreBuilder, name), stub)) + return stub + + +@pytest.mark.asyncio +async def test_direct_route_cleans_partial_content_on_mid_parse_failure( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A mid-parse failure AFTER content has spilled leaves NO store files. + + The store is now constructed BEFORE the parse and content.blob spills + incrementally, so the abort()+rmtree cleanup MUST cover the parse. A fake + parse writes a real segment through the passed ``direct_store`` (proving the + write-through is live and content was written) then raises; the store dir must + be gone afterward. + """ + captured: dict[str, object] = {} + + def _fake_parse(run, path, *, direct_store): # noqa: ANN001, ANN202 + direct_store.put_segment("sid-partial", "user", "partial content") + captured["bytes_written"] = direct_store._content_bytes_written + captured["data_dir"] = direct_store.data_dir + raise GraphParseError(f"{path}: boom mid-parse") + + monkeypatch.setattr(workload_detect, "parse_graph_workload", _fake_parse) + + stub = _cleanup_stub("bench-partial") + with pytest.raises(GraphParseError, match="boom mid-parse"): + await GraphStoreBuilder._build_graph_store_streaming( + stub, tmp_path / "dyn.jsonl", tmp_path, "dynamo_trace" + ) + + # Content really was written through before the failure... + assert captured["bytes_written"] > 0 + # ...and the whole store dir (partial content.blob included) is gone. + assert not Path(captured["data_dir"]).exists() + + +@pytest.mark.asyncio +async def test_direct_route_block_inconsistent_record_aborts_store( + tmp_path: Path, +) -> None: + """A block-inconsistent record through the direct route surfaces as + ``GraphParseError`` (the seam re-wraps ``DynamoISLMismatchError``) AND leaves + no store dir -- the real end-to-end proof the parse-failure cleanup composes + with Task-5's abort path. + + Built on a REAL resolved ``BenchmarkRun`` (the MagicMock path-drift trap: + a stub config silently no-ops the wrong attribute read).""" + bad = tmp_path / "bad.jsonl" + # input_length=100 with 2 hashes at block_size=16 violates + # (n-1)*bs < input_length <= n*bs (16 < 100 <= 32 is False). + bad.write_bytes(orjson.dumps(_dynamo_record(1000, "s1", 100, [111, 222]))) + + run = make_run_from_cli( + CLIConfig( + model_names=["test-model"], + input_file=str(bad), + tokenizer_name="builtin", + ) + ) + builder = GraphStoreBuilder(run) + store_dir = tmp_path / f"aiperf_graph_segments_{run.benchmark_id}" + + with pytest.raises(GraphParseError, match="not.*block-aligned"): + await builder._build_graph_store_streaming(bad, tmp_path, "dynamo_trace") + + assert not store_dir.exists() + + +# --- REAL-BenchmarkRun route parity --------------------------------------------- + + +@pytest.mark.asyncio +async def test_real_run_dynamo_route_hands_store_and_matches_eager( + dyn_trace: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """On a REAL resolved run the dynamo route (i) hands a live store to the parse + and (ii) produces a store + facet byte/value-identical to the eager route. + + Spies the real ``parse_graph_workload`` to capture the threaded + ``direct_store``; the eager reference re-parses the SAME run (same ctx -> same + seed) so the byte comparison needs no knowledge of the resolved seed.""" + run = make_run_from_cli( + CLIConfig( + model_names=["test-model"], + input_file=str(dyn_trace), + tokenizer_name="builtin", + ) + ) + + real_parse = workload_detect.parse_graph_workload + captured: dict[str, object] = {} + + def _spy(run_arg, path_arg, **kwargs): # noqa: ANN001, ANN202 + captured["direct_store"] = kwargs.get("direct_store") + return real_parse(run_arg, path_arg, **kwargs) + + # Direct route (real parse under the spy). The eager reference below calls + # ``real_parse`` directly, so the spy never intercepts it. + monkeypatch.setattr(workload_detect, "parse_graph_workload", _spy) + builder = GraphStoreBuilder(run) + catalog, prefix_source = await builder._build_graph_store_streaming( + dyn_trace, tmp_path, "dynamo_trace" + ) + + # (i) The dynamo route threaded a LIVE unified store into the parse. + assert isinstance(captured["direct_store"], GraphSegmentUnifiedBackingStore) + assert builder._sidecar_path is not None and builder._sidecar_path.exists() + + # Eager reference: the SAME run, no direct_store. + eager_parsed = real_parse(run, dyn_trace) + eager_store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id="eager-route-ref" + ) + eager_catalog = await build_unified_trie_store_interned(eager_parsed, eager_store) + + # (ii-a) Store bytes match the eager route. + direct_dir = tmp_path / f"aiperf_graph_segments_{run.benchmark_id}" + eager_dir = tmp_path / "aiperf_graph_segments_eager-route-ref" + direct_files = sorted(p.name for p in direct_dir.iterdir()) + eager_files = sorted(p.name for p in eager_dir.iterdir()) + assert direct_files == eager_files and direct_files + for name in direct_files: + assert (direct_dir / name).read_bytes() == (eager_dir / name).read_bytes(), ( + f"direct-route store file {name!r} diverged from the eager route" + ) + assert catalog == eager_catalog + + # (ii-b) Facet (per-node prefix-cache map) matches the eager route. + assert GraphStoreBuilder._build_graph_prefix_cache_by_trace( + prefix_source + ) == GraphStoreBuilder._build_graph_prefix_cache_by_trace(eager_parsed) + + # (ii-c) Sidecar bytes match the eager route (benchmark_id affects only the + # path, never the content-free encoded structural graph). + eager_sidecar = write_graph_meta_sidecar( + eager_parsed, + base_path=tmp_path, + benchmark_id="eager-sidecar-ref", + source_fingerprint={}, + schema_version=GRAPH_META_SCHEMA_VERSION, + ) + assert builder._sidecar_path.read_bytes() == eager_sidecar.read_bytes() diff --git a/tests/unit/dataset/test_dynamo_store_golden_digest.py b/tests/unit/dataset/test_dynamo_store_golden_digest.py new file mode 100644 index 0000000000..502429797d --- /dev/null +++ b/tests/unit/dataset/test_dynamo_store_golden_digest.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Golden store-digest pin: committed blake2b digests of every finalized +unified-store file for a deterministic dynamo fixture, built at the CURRENT tree. + +This is the uniform-drift detector the streaming/eager parity suites provably +CANNOT be: every parity test builds BOTH stores from a single in-test parse, so a +change that shifts the produced bytes uniformly (a re-hashed corpus, a reshaped +envelope, a different span packing) leaves those tests green while silently +moving the store. A digest committed as a literal at this tree is the time-pin +they lack -- it compares today's bytes against bytes frozen in the commit. + +The fixture records are COPIED (not imported) from +``test_dynamo_streaming_store_parity.py`` so this pin is self-contained: an edit +to that suite's fixture cannot silently re-point the golden. +""" + +from __future__ import annotations + +import hashlib +from pathlib import Path + +import orjson +import pytest + +from aiperf.dataset.graph.adapters.dynamo.trace import from_dynamo_trace +from aiperf.dataset.graph.segment_ir.store_builder import ( + build_unified_trie_store_interned, +) +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, +) + +# The four finalized unified-store files, in the fixed order finalize() writes +# them. blake2b(digest_size=16) is a 128-bit content fingerprint: collision-free +# for a drift detector, compact enough to read in a diff. +_STORE_FILES = ("content.blob", "content.idx", "nodes.blob", "nodes.idx") + +# Committed golden digests, pinned at this tree. Regenerating them is DELIBERATE: +# run this test, copy the reported ``got`` values here, and state in the commit +# WHY the store bytes moved (a corpus/trie/envelope change). An UNEXPLAINED drift +# is a frozen-behavior break -- investigate, do not re-pin blindly. +# +# nodes.blob / nodes.idx re-pinned for the dynamo multi-graph restructure: this +# fixture's ``s1`` + ``s2`` sessions are two independent session-trees, so the +# adapter now emits TWO per-tree ``TraceRecord``s (``s1`` with nodes s1_a1/s1_a2 +# at ordinals 0/1; ``s2`` with node s2_a1 at ordinal 0) instead of ONE union +# trace ``s1`` carrying all three nodes at ordinals 0/1/2. The per-node manifest +# region is thus re-keyed by (trace_id, ordinal) -- a LAYOUT-only shift. The +# content pool is byte-UNCHANGED (content.blob / content.idx below are identical +# to the pre-restructure pins), proving no node/segment CONTENT changed; the +# determinism guard still holds (two independent builds agree). +# +# nodes.blob / nodes.idx re-pinned again 2026-07-07 for the recorded-output +# pinning unification: every dynamo node envelope now carries +# ``dispatch_overrides["max_output_tokens"]`` (recorded output_tokens, weka +# parity; recorded 0 upgrades to 1). Content pool digests unchanged once more. +# +# nodes.blob / nodes.idx re-pinned once more 2026-07-07: model / stream / the +# cap all moved to NATIVE LlmNode fields (Turn naming); the redundant +# ``dispatch_overrides["stream"]`` entry is gone (the envelope's top-level +# ``stream`` was always authoritative) and the envelope folds the native +# model + cap back in, so overrides became {model, max_output_tokens} -- same +# effective wire values, envelope bytes only. +# +# content.blob / content.idx re-pinned 2026-07-07 for the data-inherent node-id +# scheme ({session_id}:{k}): response/tiny synthesis seeds are node-id-derived +# and now trace-scoped, so those pool entries' bytes moved. Hash-block prompt +# content is unchanged (bare-hash-id namespace untouched); nodes.* digests are +# unchanged because manifests key by ordinal, not node id. +GOLDEN_DIGESTS: dict[str, str] = { + "content.blob": "24bf57403c7920ec1e769f21f365612c", + "content.idx": "6533bb19bdf9c1ec99e78ca704d564f3", + "nodes.blob": "da6297b1a2346fbed7922efb534db79b", + "nodes.idx": "0d8f88a7f105f60da2112c1fea7cd8e2", +} + +_REPIN = ( + "an INTENTIONAL trie/corpus/envelope change re-pins this digest deliberately " + "(regenerate the literal and explain the byte shift in the commit); an " + "unexplained drift is a frozen-behavior break -- do not delete the check" +) + + +def _dynamo_record(ts: int, sid: str, input_tokens: int, hashes: list[int]) -> dict: + """A single ``dynamo.request.trace.v1`` record. + + Byte-for-byte identical to ``test_dynamo_streaming_store_parity._dynamo_record`` + -- copied, not imported, so the golden is self-contained. + """ + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": {"session_id": sid}, + "request": { + "request_id": f"r{ts}", + "model": "m", + "input_tokens": input_tokens, + "output_tokens": 8, + "cached_tokens": 0, + "replay": { + "trace_block_size": 16, + "input_length": input_tokens, + "input_sequence_hashes": hashes, + }, + }, + } + + +def _write_dynamo_fixture(path: Path) -> Path: + """Write the 3-record dynamo trace fixture with literals copied verbatim from + ``test_dynamo_streaming_store_parity.dyn_trace``.""" + records = [ + _dynamo_record(1000, "s1", 32, [111, 222]), + _dynamo_record(2000, "s1", 64, [111, 222, 333, 444]), + _dynamo_record(3000, "s2", 48, [555, 666, 777]), + ] + path.write_bytes(b"\n".join(orjson.dumps(r) for r in records)) + return path + + +async def _build_store_digests( + fixture: Path, tmp_path: Path, bid: str +) -> dict[str, str]: + """Parse the fixture fresh, drain it interned into a store keyed by ``bid``, + and return ``{filename: blake2b-hexdigest}`` for the four finalized files. + + Each call re-parses from disk (``from_dynamo_trace`` reads fresh chains), so + two calls are fully independent builds -- the basis for the determinism guard. + """ + parsed = from_dynamo_trace( + fixture, content_root_seed=1234, content_tokenizer="builtin" + ) + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id=bid) + await build_unified_trie_store_interned(parsed, store) + store_dir = tmp_path / f"aiperf_graph_segments_{bid}" + return { + name: hashlib.blake2b( + (store_dir / name).read_bytes(), digest_size=16 + ).hexdigest() + for name in _STORE_FILES + } + + +@pytest.mark.asyncio +async def test_golden_store_digest_pin(tmp_path: Path) -> None: + """The finalized unified-store files digest to their committed golden values. + + Two guards in one test: + + 1. Determinism guard -- two independent in-test builds (fresh parse each) must + produce byte-identical files. A disagreement means the build is + nondeterministic (e.g. a digest depends on dict/set iteration order across + runs); NOTHING downstream can be gated on a golden that is not itself + stable, so this fails FIRST and hard. + 2. Golden pin -- the (deterministic) digests must equal the committed + literals. This is the uniform-drift tripwire for the frozen-code rewrites + that follow. + """ + fixture = _write_dynamo_fixture(tmp_path / "dyn_golden.jsonl") + + first = await _build_store_digests(fixture, tmp_path, "golden-a") + second = await _build_store_digests(fixture, tmp_path, "golden-b") + + assert first == second, ( + "nondeterministic store build: two independent in-test builds produced " + f"different digests ({first} vs {second}). A golden pin over a " + "nondeterministic build cannot gate anything downstream -- STOP and fix " + "the ordering nondeterminism before pinning." + ) + + for name in _STORE_FILES: + assert first[name] == GOLDEN_DIGESTS[name], ( + f"golden store-digest drift for {name!r}: got {first[name]!r}, " + f"pinned {GOLDEN_DIGESTS[name]!r}. {_REPIN}" + ) diff --git a/tests/unit/dataset/test_dynamo_streaming_store_parity.py b/tests/unit/dataset/test_dynamo_streaming_store_parity.py new file mode 100644 index 0000000000..a59167c28c --- /dev/null +++ b/tests/unit/dataset/test_dynamo_streaming_store_parity.py @@ -0,0 +1,464 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dynamo/native store-build parity: the streamed payload drain must equal the +eager interned build, and the prefix-cache map must survive the structural +handoff -- the invariants the dynamo/native streaming gate flip relies on.""" + +from __future__ import annotations + +from pathlib import Path + +import msgspec +import orjson +import pytest + +from aiperf.dataset.graph.adapters.dynamo.trace import from_dynamo_trace +from aiperf.dataset.graph.codecs import decode_parsed_graph_msgpack +from aiperf.dataset.graph.merge import merge_parsed_graphs +from aiperf.dataset.graph.models import GraphRecord, LlmNode, ParsedGraph +from aiperf.dataset.graph.parser import parse_native +from aiperf.dataset.graph.segment_ir.store_builder import ( + build_unified_trie_store_from_payloads, + build_unified_trie_store_interned, + graph_carries_assembly_slots, + iter_trace_segment_payloads, +) +from aiperf.dataset.graph.store_build import GraphStoreBuilder +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) + + +def _dynamo_record(ts: int, sid: str, input_tokens: int, hashes: list[int]) -> dict: + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": {"session_id": sid}, + "request": { + "request_id": f"r{ts}", + "model": "m", + "input_tokens": input_tokens, + "output_tokens": 8, + "cached_tokens": 0, + "replay": { + "trace_block_size": 16, + "input_length": input_tokens, + "input_sequence_hashes": hashes, + }, + }, + } + + +@pytest.fixture +def dyn_trace(tmp_path: Path) -> Path: + p = tmp_path / "dyn_parity.jsonl" + records = [ + _dynamo_record(1000, "s1", 32, [111, 222]), + _dynamo_record(2000, "s1", 64, [111, 222, 333, 444]), + _dynamo_record(3000, "s2", 48, [555, 666, 777]), + ] + p.write_bytes(b"\n".join(orjson.dumps(r) for r in records)) + return p + + +@pytest.mark.asyncio +async def test_dynamo_payload_stream_store_matches_interned_oracle( + dyn_trace: Path, tmp_path: Path +) -> None: + parsed = from_dynamo_trace( + dyn_trace, content_root_seed=1234, content_tokenizer="builtin" + ) + assert parsed.segment_pool is not None + + eager_store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id="eager" + ) + eager_catalog = await build_unified_trie_store_interned(parsed, eager_store) + + stream_store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id="stream" + ) + stream_catalog = await build_unified_trie_store_from_payloads( + iter_trace_segment_payloads(parsed), stream_store + ) + + assert stream_catalog == eager_catalog and eager_catalog + + # Strong equivalence: the persisted unified store is byte-for-byte identical + # on disk (mirrors ``test_hf_streaming_trie_stores`` eager-vs-streaming + # oracle), which subsumes the interned handle map, node-manifest region, and + # content pool -- the whole store, not just the materialized messages. + eager_dir = tmp_path / "aiperf_graph_segments_eager" + stream_dir = tmp_path / "aiperf_graph_segments_stream" + eager_files = sorted(p.name for p in eager_dir.iterdir()) + stream_files = sorted(p.name for p in stream_dir.iterdir()) + assert eager_files == stream_files and eager_files, ( + f"unified store file sets differ: {eager_files} vs {stream_files}" + ) + for name in eager_files: + assert (eager_dir / name).read_bytes() == (stream_dir / name).read_bytes(), ( + f"unified store file {name!r} differs between streaming and eager" + ) + + # Semantic equivalence through the worker read face: the byte-identical + # stores materialize identical content and agree on the non-handle envelope + # fields (handles are store-local ints; compare MATERIALIZED content). + with ( + GraphSegmentUnifiedClient(tmp_path, "eager").open() as ec, + GraphSegmentUnifiedClient(tmp_path, "stream").open() as sc, + ): + for trace_id, ordinals in eager_catalog.items(): + for ordinal in ordinals.values(): + e_raw = ec.get_node_envelope(trace_id, ordinal, "profiling") + s_raw = sc.get_node_envelope(trace_id, ordinal, "profiling") + assert e_raw is not None and s_raw is not None + e_env = orjson.loads(e_raw) + s_env = orjson.loads(s_raw) + assert ec.materialize_handles( + e_env["handles"] + ) == sc.materialize_handles(s_env["handles"]) + assert {k: v for k, v in e_env.items() if k != "handles"} == { + k: v for k, v in s_env.items() if k != "handles" + } + + +def test_dynamo_prefix_cache_from_structural_matches_eager(dyn_trace: Path) -> None: + parsed = from_dynamo_trace( + dyn_trace, content_root_seed=42, content_tokenizer="builtin" + ) + eager_map = GraphStoreBuilder._build_graph_prefix_cache_by_trace(parsed) + assert eager_map, "dynamo fixture must stamp a non-empty prefix-cache map" + + structural_blobs = [ + p.structural_graph + for p in iter_trace_segment_payloads(parsed) + if p.structural_graph + ] + merged = merge_parsed_graphs( + decode_parsed_graph_msgpack(b) for b in structural_blobs + ) + assert GraphStoreBuilder._build_graph_prefix_cache_by_trace(merged) == eager_map + + +def _with_sentinel_prompts(parsed: ParsedGraph) -> ParsedGraph: + """Return a copy of ``parsed`` whose every ``LlmNode.prompt`` is a non-empty + sentinel, across ``parsed.graph`` AND every ``parsed.graphs`` value. + + A drain that read the inline ``node.prompt`` (rather than the segment pool + + trie envelope) would emit different bytes for this copy than for the + real-content baseline. Sentinelling both graph surfaces keeps multi-graph + corpus shapes pinned even though the dynamo fixture is single-graph. + """ + sentinel: list = [{"role": "user", "content": "SENTINEL"}] + + def _stamp(graph: GraphRecord) -> GraphRecord: + nodes = { + nid: ( + msgspec.structs.replace(node, prompt=sentinel) + if isinstance(node, LlmNode) + else node + ) + for nid, node in graph.nodes.items() + } + return msgspec.structs.replace(graph, nodes=nodes) + + return msgspec.structs.replace( + parsed, + graph=_stamp(parsed.graph), + graphs={ref: _stamp(g) for ref, g in parsed.graphs.items()}, + ) + + +def _all_llm_nodes(parsed: ParsedGraph) -> list[LlmNode]: + return [ + node + for graph in (parsed.graph, *parsed.graphs.values()) + for node in graph.nodes.values() + if isinstance(node, LlmNode) + ] + + +async def _build_interned_dir(parsed: ParsedGraph, tmp_path: Path, bid: str) -> Path: + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id=bid) + await build_unified_trie_store_interned(parsed, store) + return tmp_path / f"aiperf_graph_segments_{bid}" + + +async def _build_streamed_dir(parsed: ParsedGraph, tmp_path: Path, bid: str) -> Path: + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id=bid) + await build_unified_trie_store_from_payloads( + iter_trace_segment_payloads(parsed), store + ) + return tmp_path / f"aiperf_graph_segments_{bid}" + + +async def _build_direct_dir( + dyn_trace: Path, tmp_path: Path, bid: str +) -> tuple[Path, ParsedGraph]: + """Build the store via the DIRECT write-through route and return (dir, parsed). + + The store is constructed FIRST and threaded as ``direct_store`` so + ``build_trie_ir``'s ``pool.add`` interns each segment straight into it during + the parse; the returned pool is an EMPTY pool (the per-tree write-through + shims carry no ``by_id`` and the multi-graph merge normalizes them to a plain + empty ``SegmentPool``) and the interned drain's put loop no-ops over it + (segments already resident). The same content_root_seed / tokenizer as the + eager/streaming legs so all three parses are the identical trie. + """ + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id=bid) + parsed = from_dynamo_trace( + dyn_trace, + content_root_seed=1234, + content_tokenizer="builtin", + direct_store=store, + ) + await build_unified_trie_store_interned(parsed, store) + return tmp_path / f"aiperf_graph_segments_{bid}", parsed + + +def _assert_store_dirs_identical(dir_a: Path, dir_b: Path) -> None: + files_a = sorted(p.name for p in dir_a.iterdir()) + files_b = sorted(p.name for p in dir_b.iterdir()) + assert files_a == files_b and files_a, ( + f"unified store file sets differ: {files_a} vs {files_b}" + ) + for name in files_a: + assert (dir_a / name).read_bytes() == (dir_b / name).read_bytes(), ( + f"unified store file {name!r} differs -- a drain read inline " + f"node.prompt instead of the segment pool + trie envelope" + ) + + +@pytest.mark.asyncio +async def test_store_bytes_independent_of_inline_prompt( + dyn_trace: Path, tmp_path: Path +) -> None: + """The persisted store bytes are a function of (segment pool, trie envelope) + ONLY -- never the inline ``LlmNode.prompt``. + + Pins the trie-route invariant that no drain reads ``node.prompt``: a + real-content baseline parse and a copy whose every node carries a sentinel + non-empty prompt build BYTE-IDENTICAL stores through BOTH the eager + (``build_unified_trie_store_interned``) and streaming + (``build_unified_trie_store_from_payloads``) drains. Holds before AND after + adapters stamp ``prompt=[]``, so the pin survives the change permanently. + """ + parsed = from_dynamo_trace( + dyn_trace, content_root_seed=1234, content_tokenizer="builtin" + ) + sentinel = _with_sentinel_prompts(parsed) + # Guard the guard: the sentinel copy must actually carry non-empty inline + # prompts, else the byte-equality below would be vacuously true. + assert _all_llm_nodes(sentinel) and all( + node.prompt == [{"role": "user", "content": "SENTINEL"}] + for node in _all_llm_nodes(sentinel) + ) + + base_interned = await _build_interned_dir(parsed, tmp_path, "base-interned") + sent_interned = await _build_interned_dir(sentinel, tmp_path, "sent-interned") + _assert_store_dirs_identical(base_interned, sent_interned) + + base_stream = await _build_streamed_dir(parsed, tmp_path, "base-stream") + sent_stream = await _build_streamed_dir(sentinel, tmp_path, "sent-stream") + _assert_store_dirs_identical(base_stream, sent_stream) + + +@pytest.mark.asyncio +async def test_dynamo_release_replay_store_bytes_identical( + dyn_trace: Path, tmp_path: Path +) -> None: + """The ``release_replay`` adjunct is a pure build-time RAM optimization: a + parse with replay-release ON produces a BYTE-IDENTICAL unified store to one + with it OFF. + + ``dynamo_trie_nodes`` copies each record's recorded hashes / input_length + into the ``TrieRequest`` BEFORE freeing ``req.replay``, so the trie IR -- + and therefore the persisted store (content pool + node manifests) -- is + unchanged. Each ``from_dynamo_trace`` call reads fresh chains from disk, so + the release on one parse never affects the other. + """ + keep = from_dynamo_trace( + dyn_trace, content_root_seed=1234, content_tokenizer="builtin" + ) + release = from_dynamo_trace( + dyn_trace, + content_root_seed=1234, + content_tokenizer="builtin", + release_replay=True, + ) + keep_dir = await _build_interned_dir(keep, tmp_path, "replay-keep") + release_dir = await _build_interned_dir(release, tmp_path, "replay-release") + _assert_store_dirs_identical(keep_dir, release_dir) + + +@pytest.mark.asyncio +async def test_dynamo_direct_store_route_matches_eager_bytes( + dyn_trace: Path, tmp_path: Path +) -> None: + """THREE-WAY parity: the direct write-through store is byte-for-byte identical + to the eager interned store. + + ``test_dynamo_payload_stream_store_matches_interned_oracle`` already proves + eager == streaming on the SAME store files; this proves direct == eager on + the same files, so by transitivity **direct == eager == streaming**. Both the + direct and eager routes intern in ``build_trie_ir``'s content-loop + first-occurrence order (the single ordering authority), so the write-through + sink assigns the same handle stream -- the parity holds by construction, not + by luck. + """ + eager_dir = await _build_interned_dir( + from_dynamo_trace( + dyn_trace, content_root_seed=1234, content_tokenizer="builtin" + ), + tmp_path, + "eager-3way", + ) + direct_dir, _parsed = await _build_direct_dir(dyn_trace, tmp_path, "direct-3way") + _assert_store_dirs_identical(eager_dir, direct_dir) + + +@pytest.mark.asyncio +async def test_dynamo_direct_store_route_mechanism( + dyn_trace: Path, tmp_path: Path +) -> None: + """The direct route returns an empty pool and content-free nodes: segments + live in the store, not the pool. + + The DM-level mechanism pins (original §1 ruling): the returned pool is EMPTY + -- the write-through shim carries no ``by_id`` and the multi-graph merge + (:func:`merge_parsed_graphs`) normalizes the per-tree shims to a plain empty + ``SegmentPool`` -- so the interned drain's put loop no-ops over it, every + ``LlmNode.prompt`` was stamped ``[]`` at lowering, and ``metadata["trie"]`` + carries ONLY ``prompt_segment_ids`` (the load-bearing store path -- + ``hash_ids`` and inline content are gone). + """ + _direct_dir, parsed = await _build_direct_dir(dyn_trace, tmp_path, "direct-mech") + + assert parsed.segment_pool is not None + assert parsed.segment_pool.by_id == {} + nodes = _all_llm_nodes(parsed) + assert nodes, "dynamo fixture must lower to at least one LlmNode" + for node in nodes: + assert node.prompt == [] + trie_meta = (node.metadata or {}).get("trie") or {} + assert set(trie_meta) == {"prompt_segment_ids"}, ( + f"node {node.node_id!r} trie metadata must carry only " + f"prompt_segment_ids, got {sorted(trie_meta)}" + ) + + +def test_dynamo_direct_route_prefix_cache_matches_eager(dyn_trace: Path) -> None: + """The direct route stamps the SAME per-node prefix-cache counts as the eager + parse (template: ``test_dynamo_prefix_cache_from_structural_matches_eager``). + + ``stamp_theoretical_prefix_cache`` reads ``node.request.hash_ids``, never the + pool, so swapping in the write-through shim leaves the prefix-cache map + untouched -- pinned here on the direct route so a future shim change that + perturbed lowering would be caught. + """ + eager = from_dynamo_trace( + dyn_trace, content_root_seed=1234, content_tokenizer="builtin" + ) + eager_map = GraphStoreBuilder._build_graph_prefix_cache_by_trace(eager) + assert eager_map, "dynamo fixture must stamp a non-empty prefix-cache map" + + # A throwaway store is fine: we only read the direct-route parse's metadata. + class _NullStore: + def put_segment(self, *a: object, **k: object) -> int: + return 0 + + direct = from_dynamo_trace( + dyn_trace, + content_root_seed=1234, + content_tokenizer="builtin", + direct_store=_NullStore(), # type: ignore[arg-type] + ) + assert GraphStoreBuilder._build_graph_prefix_cache_by_trace(direct) == eager_map + + +@pytest.fixture +def native_workload(tmp_path: Path) -> Path: + p = tmp_path / "native_parity.yaml" + p.write_text( + """ +graph: + nodes: + a: + prompt: + - {role: user, content: "question one"} + output: a_out + b: + prompt: + - {role: user, content: "question two"} + output: b_out + edges: + - {source: START, target: a} + - {source: a, target: b} + - {source: b, target: END} +traces: + - id: t1 + - id: t2 +""" + ) + return p + + +@pytest.mark.asyncio +async def test_native_payload_stream_store_matches_interned_oracle( + native_workload: Path, tmp_path: Path +) -> None: + parsed = parse_native(native_workload) + assert parsed.segment_pool is not None + # Slot-free premise: plain native graphs stream with no eager fallback. + assert not graph_carries_assembly_slots(parsed) + + eager_store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id="native-eager" + ) + eager_catalog = await build_unified_trie_store_interned(parsed, eager_store) + + stream_store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id="native-stream" + ) + stream_catalog = await build_unified_trie_store_from_payloads( + iter_trace_segment_payloads(parsed), stream_store + ) + + assert stream_catalog == eager_catalog and eager_catalog + + eager_dir = tmp_path / "aiperf_graph_segments_native-eager" + stream_dir = tmp_path / "aiperf_graph_segments_native-stream" + eager_files = sorted(p.name for p in eager_dir.iterdir()) + stream_files = sorted(p.name for p in stream_dir.iterdir()) + assert eager_files == stream_files and eager_files, ( + f"unified store file sets differ: {eager_files} vs {stream_files}" + ) + for name in eager_files: + assert (eager_dir / name).read_bytes() == (stream_dir / name).read_bytes(), ( + f"unified store file {name!r} differs between streaming and eager" + ) + + with ( + GraphSegmentUnifiedClient(tmp_path, "native-eager").open() as ec, + GraphSegmentUnifiedClient(tmp_path, "native-stream").open() as sc, + ): + for trace_id, ordinals in eager_catalog.items(): + for ordinal in ordinals.values(): + e_raw = ec.get_node_envelope(trace_id, ordinal, "profiling") + s_raw = sc.get_node_envelope(trace_id, ordinal, "profiling") + assert e_raw is not None and s_raw is not None + e_env = orjson.loads(e_raw) + s_env = orjson.loads(s_raw) + assert ec.materialize_handles( + e_env["handles"] + ) == sc.materialize_handles(s_env["handles"]) + assert {k: v for k, v in e_env.items() if k != "handles"} == { + k: v for k, v in s_env.items() if k != "handles" + } + + # Native never stamps prefix-cache counts (hash-id-free graphs); pin the + # documented absence rather than a non-empty map. + assert GraphStoreBuilder._build_graph_prefix_cache_by_trace(parsed) == {} diff --git a/tests/unit/dataset/test_graph_client_metadata.py b/tests/unit/dataset/test_graph_client_metadata.py new file mode 100644 index 0000000000..ed74b18346 --- /dev/null +++ b/tests/unit/dataset/test_graph_client_metadata.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Graph-native dataset broadcast types. + +Graph runs broadcast a graph facet on ``DatasetMetadata`` and a graph-typed +``client_metadata`` (store + sidecar locations) instead of fabricating stub +conversations. These tests lock the discriminated-union wire round-trip and +the plugin registration the worker's generic client-store construction +resolves through. +""" + +from pathlib import Path + +import pytest + +from aiperf.common.exceptions import InvalidOperationError +from aiperf.common.messages import DatasetConfiguredNotification +from aiperf.common.models.dataset_models import ( + DatasetMetadata, + GraphDatasetMetadata, + GraphSegmentClientMetadata, +) +from aiperf.plugin.enums import DatasetClientStoreType, DatasetSamplingStrategy + + +def _notification(tmp_path: Path) -> DatasetConfiguredNotification: + return DatasetConfiguredNotification( + service_id="dataset_manager", + metadata=DatasetMetadata( + conversations=[], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + graph=GraphDatasetMetadata( + trace_ids=["t-1", "t-2"], + prefix_cache_by_trace={"t-1": {"n0": [3, 7]}}, + ), + ), + client_metadata=GraphSegmentClientMetadata( + store_base_path=tmp_path, + benchmark_id="bench-x", + sidecar_path=tmp_path / "aiperf_graph_meta_bench-x" / "graph_meta.msgpack", + ), + ) + + +def test_graph_client_metadata_wire_round_trip(tmp_path): + msg = _notification(tmp_path) + decoded = DatasetConfiguredNotification.model_validate(msg.model_dump(mode="json")) + assert isinstance(decoded.client_metadata, GraphSegmentClientMetadata) + assert decoded.client_metadata.client_type == DatasetClientStoreType.GRAPH_SEGMENT + assert decoded.client_metadata.benchmark_id == "bench-x" + assert decoded.metadata.graph is not None + assert decoded.metadata.graph.trace_ids == ["t-1", "t-2"] + assert decoded.metadata.graph.prefix_cache_by_trace["t-1"]["n0"] == [3, 7] + + +def test_graph_client_store_plugin_resolves(tmp_path): + from aiperf.plugin import plugins + from aiperf.plugin.enums import PluginType + + StoreClass = plugins.get_class( + PluginType.DATASET_CLIENT_STORE, DatasetClientStoreType.GRAPH_SEGMENT + ) + store = StoreClass(client_metadata=_notification(tmp_path).client_metadata) + assert store.client_metadata.benchmark_id == "bench-x" + + +@pytest.mark.asyncio +async def test_graph_client_store_get_conversation_hard_errors(tmp_path): + from aiperf.dataset.graph_client_store import GraphSegmentDatasetClientStore + + store = GraphSegmentDatasetClientStore( + client_metadata=_notification(tmp_path).client_metadata + ) + with pytest.raises(InvalidOperationError, match="no conversation store"): + await store.get_conversation("t-1") diff --git a/tests/unit/dataset/test_graph_segment_unified_store.py b/tests/unit/dataset/test_graph_segment_unified_store.py new file mode 100644 index 0000000000..5eec7708ed --- /dev/null +++ b/tests/unit/dataset/test_graph_segment_unified_store.py @@ -0,0 +1,311 @@ +import asyncio + +import orjson +import pytest + +from aiperf.common.exceptions import MemoryMapSerializationError +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) + + +def test_unified_backing_store_writes_four_files(tmp_path): + store = GraphSegmentUnifiedBackingStore(tmp_path, "b1") + store.put_segment("a", "system", "SYS") + store.put_segment("b", "user", "hi") + env = orjson.dumps({"prompt_segment_ids": ["a", "b"], "stream": True}) + store.add_node_manifest("t0", 0, "profiling", env) + asyncio.run(store.finalize()) + + d = tmp_path / "aiperf_graph_segments_b1" + assert (d / "content.blob").exists() + assert (d / "content.idx").exists() + assert (d / "nodes.blob").exists() + assert (d / "nodes.idx").exists() + + client = GraphSegmentUnifiedClient(tmp_path, "b1") + assert client.data_dir == d + + +def test_unified_client_duck_types_both_stores(tmp_path): + store = GraphSegmentUnifiedBackingStore(tmp_path, "b2") + ha = store.put_segment("a", "system", "SYS") + hb = store.put_segment("b", "user", "hi") + store.add_node_manifest_interned("t0", 0, "profiling", [ha, hb], {}, True) + asyncio.run(store.finalize()) + + with GraphSegmentUnifiedClient(tmp_path, "b2").open() as c: + # envelope addressing face: + raw = c.get_node_envelope("t0", 0, "profiling") + assert orjson.loads(raw) == { + "handles": [ha, hb], + "dispatch_overrides": {}, + "stream": True, + } + assert c.get_node_envelope("t0", 9, "profiling") is None + # Segment-content face: + assert c.materialize_handles([ha, hb]) == [ + {"role": "system", "content": "SYS"}, + {"role": "user", "content": "hi"}, + ] + assert c.build_request_body_handles([ha, hb], b"") == ( + b'{"messages":[{"role":"system","content":"SYS"},' + b'{"role":"user","content":"hi"}]}' + ) + + +def test_endpoint_extra_applied_flag_round_trips_when_set(tmp_path): + """A manifest written with ``endpoint_extra_applied=True`` reads back with the + flag; the key is written only when True.""" + store = GraphSegmentUnifiedBackingStore(tmp_path, "eea1") + ha = store.put_segment("a", "user", "hi") + store.add_node_manifest_interned( + "t0", 0, "profiling", [ha], {}, True, endpoint_extra_applied=True + ) + asyncio.run(store.finalize()) + + with GraphSegmentUnifiedClient(tmp_path, "eea1").open() as c: + env = orjson.loads(c.get_node_envelope("t0", 0, "profiling")) + assert env["endpoint_extra_applied"] is True + + +def test_endpoint_extra_applied_flag_omitted_leaves_envelope_bytes_unchanged(tmp_path): + """Default (unset) leaves the manifest envelope byte-identical to a store + written without the parameter at all -- byte-neutrality when the flag is + off is the whole contract.""" + store = GraphSegmentUnifiedBackingStore(tmp_path, "eea2") + ha = store.put_segment("a", "user", "hi") + store.add_node_manifest_interned("t0", 0, "profiling", [ha], {}, True) + asyncio.run(store.finalize()) + + with GraphSegmentUnifiedClient(tmp_path, "eea2").open() as c: + raw = c.get_node_envelope("t0", 0, "profiling") + assert b"endpoint_extra_applied" not in raw + assert orjson.loads(raw) == { + "handles": [ha], + "dispatch_overrides": {}, + "stream": True, + } + + +def test_put_segment_returns_stable_insertion_index_handles(tmp_path): + store = GraphSegmentUnifiedBackingStore(tmp_path, "h1") + assert store.put_segment("a", "system", "SYS") == 0 + assert store.put_segment("b", "user", "hi") == 1 + assert store.put_segment("a", "system", "SYS") == 0 # dedup -> same handle + assert store.segment_handle("b") == 1 + assert store.segment_handle("missing") is None + + +def test_unified_client_empty_pool_and_unknown_handle(tmp_path): + store = GraphSegmentUnifiedBackingStore(tmp_path, "b3") + asyncio.run(store.finalize()) + with GraphSegmentUnifiedClient(tmp_path, "b3").open() as c: + assert c._content_mv is None # empty pool leaves mapping unset, no raise + assert c.build_request_body_handles([], b"") == b'{"messages":[]}' + store2 = GraphSegmentUnifiedBackingStore(tmp_path, "b4") + store2.put_segment("a", "system", "SYS") + asyncio.run(store2.finalize()) + with ( + GraphSegmentUnifiedClient(tmp_path, "b4").open() as c, + pytest.raises(MemoryMapSerializationError), + ): + c.materialize_handles([1]) + + +def test_a1_json_content_idx_is_rejected_with_reparse_error(tmp_path): + """Pre-retirement A1 (non-interned JSON content.idx) stores fail loud.""" + # Hand-rolled A1-shaped store: the JSON content.idx the retired + # non-interned builder used to write ({"ids": hex->handle, "spans": ...}). + d = tmp_path / "aiperf_graph_segments_a1" + d.mkdir() + seg = orjson.dumps({"role": "user", "content": "hi"}) + (d / "content.blob").write_bytes(seg) + (d / "content.idx").write_bytes( + orjson.dumps({"ids": {"a": 0}, "spans": [[0, len(seg)]]}) + ) + (d / "nodes.blob").write_bytes(b"") + (d / "nodes.idx").write_bytes(orjson.dumps({})) + + with pytest.raises(ValueError, match="re-parse"): + GraphSegmentUnifiedClient(tmp_path, "a1").open() + + +def test_finalize_releases_write_side_buffers(tmp_path): + """finalize() drops every write-side buffer at its END. + + The store object outlives finalize through the structural-merge / sidecar / + prefix-cache build tail, where NOTHING reads these accumulation buffers + (``build_stats`` is already snapshotted at finalize ENTRY). Holding the + load-bearing ones (``_ids`` dedup map, ``_spans`` content.idx source, + ``_node_offsets``, ``_nodes_buf``) just pins RAM for zero readers. Content + is spilled to disk at ``put`` time, so ``_content_buf`` is already empty + (its clear is a formality) and the running counter tracks the footprint. + This pins the memory-release contract: an edit that stops clearing a buffer + trips here. + """ + store = GraphSegmentUnifiedBackingStore(tmp_path, "rel1") + ha = store.put_segment("a", "system", "SYS") + hb = store.put_segment("b", "user", "hi") + store.add_node_manifest_interned("t0", 0, "profiling", [ha, hb], {}, True) + + # Before finalize: content is already spilled to disk (tracked by the running + # counter, never accumulated in _content_buf); the index buffers hold state. + assert store._content_bytes_written > 0 + assert store._content_buf == bytearray() # spill: content never sits in RAM + assert len(store._nodes_buf) > 0 + assert store._ids and store._spans and store._node_offsets + + asyncio.run(store.finalize()) + + assert store._content_buf == bytearray() + assert store._nodes_buf == bytearray() + assert store._ids == {} + assert store._spans == [] + assert store._node_offsets == {} + # The snapshot survives the release -- it is taken at finalize ENTRY, before + # the writes, so post-finalize log sites keep reading real numbers. + assert store.build_stats is not None + assert store.build_stats.segment_count == 2 + + +def test_store_fully_readable_after_finalize_buffer_release(tmp_path): + """The persisted store stays fully readable AFTER finalize releases the write + buffers: open with the client and materialize a node envelope + its content + end-to-end. Byte parity is proven by the dynamo/dag_jsonl parity suites; here + we prove readability survives the release, since readers work off the flushed + files, never the writer's dropped buffers.""" + store = GraphSegmentUnifiedBackingStore(tmp_path, "rel2") + ha = store.put_segment("a", "system", "SYS") + hb = store.put_segment("b", "user", "hi") + store.add_node_manifest_interned("t0", 0, "profiling", [ha, hb], {}, True) + asyncio.run(store.finalize()) + + assert store._content_buf == bytearray() # writer state is gone + + with GraphSegmentUnifiedClient(tmp_path, "rel2").open() as c: + assert orjson.loads(c.get_node_envelope("t0", 0, "profiling")) == { + "handles": [ha, hb], + "dispatch_overrides": {}, + "stream": True, + } + assert c.materialize_handles([ha, hb]) == [ + {"role": "system", "content": "SYS"}, + {"role": "user", "content": "hi"}, + ] + + +def test_write_and_handle_reads_raise_after_finalize(tmp_path): + """Post-finalize writes AND handle lookups fail loud, not silently no-op. + + ``put_segment`` / ``add_node_manifest`` already guarded; ``segment_handle`` + now guards too because its ``_ids`` map is released at finalize END -- an + unguarded lookup would silently return ``None`` (misresolving a handle) where + it previously returned the real handle. Double-finalize stays a hard error. + """ + store = GraphSegmentUnifiedBackingStore(tmp_path, "rel3") + store.put_segment("a", "system", "SYS") + asyncio.run(store.finalize()) + + with pytest.raises(RuntimeError, match="after finalize"): + store.put_segment("b", "user", "hi") + with pytest.raises(RuntimeError, match="after finalize"): + store.add_node_manifest("t0", 0, "profiling", b"{}") + with pytest.raises(RuntimeError, match="after finalize"): + store.segment_handle("a") + with pytest.raises(RuntimeError, match="finalize called twice"): + asyncio.run(store.finalize()) + + +def test_interned_content_idx_is_packed_and_client_reads_int_handles(tmp_path): + store = GraphSegmentUnifiedBackingStore(tmp_path, "h2") + ha = store.put_segment("a", "system", "SYS") + hb = store.put_segment("b", "user", "hi") + store.add_node_manifest_interned("t0", 0, "profiling", [ha, hb], {}, True) + asyncio.run(store.finalize()) + + # content.idx is packed binary (8-byte 'Q' pairs), not JSON: + raw = (tmp_path / "aiperf_graph_segments_h2" / "content.idx").read_bytes() + assert len(raw) == 2 * 2 * 8 # 2 segments * (off,size) * 8 bytes + with pytest.raises(orjson.JSONDecodeError): + orjson.loads(raw) # not JSON + + with GraphSegmentUnifiedClient(tmp_path, "h2").open() as c: + assert c.materialize_handles([ha, hb]) == [ + {"role": "system", "content": "SYS"}, + {"role": "user", "content": "hi"}, + ] + assert c.build_request_body_handles([ha, hb], b"") == ( + b'{"messages":[{"role":"system","content":"SYS"},' + b'{"role":"user","content":"hi"}]}' + ) + with pytest.raises(MemoryMapSerializationError): + c.materialize_handles([999]) # unknown handle raises + + +def test_content_bytes_counter_equals_blob_size_and_span_sum(tmp_path): + """The incremental-spill counter is the single source of truth for content + size: ``_content_bytes_written`` == sum of span lengths (pre-finalize) == + the finalized ``content.blob`` size == ``build_stats.content_bytes``. + + This is the spill's core invariant -- a divergence would mean a segment's + bytes were counted but not written (or vice versa), silently corrupting + every span offset after it. + """ + store = GraphSegmentUnifiedBackingStore(tmp_path, "spill-eq") + store.put_segment("a", "system", "SYS") + store.put_segment("b", "user", "hi") + store.put_segment("c", "assistant", "there") + store.put_segment("a", "system", "SYS") # dedup: counted/written once + + counter = store._content_bytes_written + assert counter == sum(size for _off, size in store._spans) + + asyncio.run(store.finalize()) + + assert store.build_stats is not None + assert store.build_stats.content_bytes == counter + blob = (tmp_path / "aiperf_graph_segments_spill-eq" / "content.blob").read_bytes() + assert len(blob) == counter + + +def test_abort_removes_partial_store_files_and_is_idempotent(tmp_path): + """A store that errors before finalize leaves no half-written files. + + put_segment streams straight to ``content.blob``, so an aborted build would + otherwise leave a partial blob for a later open to trip on. ``abort()`` + closes the write handle and unlinks the four files; it is idempotent and + non-raising (safe to call twice, and after finalize). + """ + store = GraphSegmentUnifiedBackingStore(tmp_path, "abort1") + store.put_segment("a", "user", "hi") + d = tmp_path / "aiperf_graph_segments_abort1" + assert (d / "content.blob").exists() # eager-opened + written at put time + + store.abort() + + assert not (d / "content.blob").exists() + assert not (d / "content.idx").exists() + assert not (d / "nodes.blob").exists() + assert not (d / "nodes.idx").exists() + + store.abort() # idempotent second call must not raise + + +def test_abort_after_finalize_preserves_the_finalized_store(tmp_path): + """``abort()`` after a SUCCESSFUL finalize is a safe no-op: the complete + on-disk store is never unlinked (the ``_finalized`` guard skips the unlink), + so a stray abort on the builder's success path cannot delete a good store.""" + store = GraphSegmentUnifiedBackingStore(tmp_path, "abort2") + ha = store.put_segment("a", "user", "hi") + store.add_node_manifest_interned("t0", 0, "profiling", [ha], {}, True) + asyncio.run(store.finalize()) + + store.abort() + + d = tmp_path / "aiperf_graph_segments_abort2" + assert (d / "content.blob").exists() + assert (d / "content.idx").exists() + with GraphSegmentUnifiedClient(tmp_path, "abort2").open() as c: + assert c.materialize_handles([ha]) == [{"role": "user", "content": "hi"}] diff --git a/tests/unit/dataset/test_graph_store_build_offload.py b/tests/unit/dataset/test_graph_store_build_offload.py new file mode 100644 index 0000000000..5713c834d0 --- /dev/null +++ b/tests/unit/dataset/test_graph_store_build_offload.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DM3: graph store builds must not freeze the DatasetManager event loop. + +The HF streaming drain consumes a multiprocessing-backed iterator whose +``__next__`` blocks in ``result.get()`` per trace. Run synchronously on the +service loop that freezes heartbeats for the entire corpus-scale build (and +lets ``PROFILE_CONFIGURE_TIMEOUT`` kill a still-progressing configure). +``_build_graph_store_streaming_trie`` therefore drains the stream in a worker +thread; this test drives a deliberately slow payload iterator and asserts a +concurrent ticker keeps ticking while the stream is mid-drain. + +The EAGER (local file/dir) path has the same failure shape: the interned +unified-store build is a zero-yield CPU drain (orjson-encode + pool copy of +ALL synthesized content) until its trailing ``finalize`` await, so +``_build_interned_unified_store`` runs it in a worker thread too. The second +test drives a deliberately blocking builder through that seam and asserts the +same ticker liveness. +""" + +from __future__ import annotations + +import asyncio +import time +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from aiperf.dataset.graph.segment_ir.store_builder import TraceSegmentPayload +from aiperf.dataset.graph.store_build import GraphStoreBuilder + + +class _StubManager: + """Just the attributes ``_build_graph_store_streaming_trie`` reads from self.""" + + def __init__(self) -> None: + self.run = SimpleNamespace(benchmark_id="bench-offload-test") + self.infos: list[object] = [] + self.merge_calls: list[list[bytes]] = [] + self.sidecar_calls: list[tuple] = [] + self.merged_sentinel = object() + + def info(self, msg: object) -> None: + self.infos.append(msg() if callable(msg) else msg) + + def _merge_structural_graphs(self, structural_sink: list[bytes]) -> object: + self.merge_calls.append(structural_sink) + return self.merged_sentinel + + def _write_graph_sidecar( + self, + merged: object, + catalog: dict[str, dict[str, int]], + base_path: Path, + ) -> None: + self.sidecar_calls.append((merged, catalog, base_path)) + + +@pytest.mark.asyncio +async def test_streaming_trie_drain_keeps_event_loop_responsive( + tmp_path: Path, +) -> None: + ticks = 0 + + async def ticker() -> None: + nonlocal ticks + while True: + ticks += 1 + await asyncio.sleep(0) + + ticks_at_payload: list[int] = [] + + def slow_payloads(): + for i in range(3): + # Stands in for the blocking multiprocessing result.get() the + # real payload iterator performs per trace. + time.sleep(0.05) + ticks_at_payload.append(ticks) + yield TraceSegmentPayload( + trace_id=f"t{i}", + node_ordinals={f"n{i}": 0}, + envelopes=[], + ) + + manager = _StubManager() + ticker_task = asyncio.create_task(ticker()) + try: + catalog, merged = await GraphStoreBuilder._build_graph_store_streaming_trie( + manager, slow_payloads(), tmp_path + ) + finally: + ticker_task.cancel() + + assert set(catalog) == {"t0", "t1", "t2"} + assert merged is manager.merged_sentinel + assert manager.sidecar_calls, "sidecar write must still run after the drain" + # A synchronous drain freezes the loop for the whole stream: the ticker + # never gets scheduled and every sample stays 0. + assert ticks_at_payload[-1] > 0, ( + "event loop made no progress while the payload stream was draining; " + "the drain is blocking the DatasetManager loop" + ) + + +@pytest.mark.asyncio +async def test_eager_interned_store_build_keeps_event_loop_responsive( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The eager interned build runs off-loop; a concurrent ticker keeps ticking.""" + from aiperf.dataset.graph.segment_ir import store_builder + + ticks = 0 + + async def ticker() -> None: + nonlocal ticks + while True: + ticks += 1 + await asyncio.sleep(0) + + ticks_at_step: list[int] = [] + + async def slow_interned_build(parsed: object, unified: object) -> dict: + # Stands in for the zero-yield orjson-encode + pool-copy drain the + # real builder performs before its trailing ``store.finalize()``. + for _ in range(3): + time.sleep(0.05) + ticks_at_step.append(ticks) + return {"t0": {"n0": 0}} + + monkeypatch.setattr( + store_builder, "build_unified_trie_store_interned", slow_interned_build + ) + + manager = _StubManager() + ticker_task = asyncio.create_task(ticker()) + try: + catalog = await GraphStoreBuilder._build_interned_unified_store( + manager, SimpleNamespace(), SimpleNamespace() + ) + finally: + ticker_task.cancel() + + assert catalog == {"t0": {"n0": 0}} + # A synchronous build freezes the loop for the whole drain: the ticker + # never gets scheduled and every sample stays 0. + assert ticks_at_step[-1] > 0, ( + "event loop made no progress while the eager interned store build was " + "running; the build is blocking the DatasetManager loop" + ) diff --git a/tests/unit/dataset/test_graph_store_build_stats.py b/tests/unit/dataset/test_graph_store_build_stats.py new file mode 100644 index 0000000000..399c6e1d41 --- /dev/null +++ b/tests/unit/dataset/test_graph_store_build_stats.py @@ -0,0 +1,266 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Graph store build-stats snapshot: exact-field pins, cross-payload dedup, and +the GraphStoreBuilder build-complete log line. + +The stats are a cheap ``O(traces) + O(1)`` snapshot computed at +``GraphSegmentUnifiedBackingStore.finalize()`` ENTRY -- the measurement baseline +that turns pool/envelope-size regressions into a visible log delta and a hard +CI failure instead of mystery RSS at corpus scale. +""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +import orjson +import pytest + +from aiperf.dataset import graph_segment_unified_store as store_mod +from aiperf.dataset.graph.adapters.dag_jsonl.trace import from_dag_jsonl +from aiperf.dataset.graph.segment_ir.store_builder import ( + TraceSegmentPayload, + build_unified_trie_store_from_payloads, + build_unified_trie_store_interned, +) +from aiperf.dataset.graph.store_build import ( + GraphStoreBuilder, + _format_store_build_stats, +) +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphStoreBuildStats, + NodeEnvelope, +) + +DAG_FIXTURES = Path(__file__).parents[2] / "fixtures" / "dag" +# Spawn-only dag graph: a fast, slot-free trie parse (one trace carrying the +# root plus its spawned child) drained through the frozen trie pipeline. +SPAWN_MINIMAL_FIXTURE = DAG_FIXTURES / "spawn_minimal.dag.jsonl" + + +def _envelope_bytes(prompt_segment_ids: list[str]) -> bytes: + """A minimal slot-free manifest envelope the streaming drain resolves.""" + return orjson.dumps( + { + "prompt_segment_ids": prompt_segment_ids, + "dispatch_overrides": {}, + "stream": False, + } + ) + + +@pytest.mark.asyncio +async def test_finalize_computes_build_stats_exact_fields(tmp_path: Path) -> None: + """Pin the exact drained-store field values for a deterministic trie parse. + + These literals are structural products of the FROZEN trie pipeline over the + spawn_minimal fixture. Tripping this assertion on an INTENTIONAL corpus or + trie change is the DESIRED tripwire: update the literal deliberately after + confirming the new build shape, do not delete the check. + """ + parsed = from_dag_jsonl(str(SPAWN_MINIMAL_FIXTURE)) + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id="stats") + + assert store.build_stats is None, "build_stats must be None until finalize" + + await build_unified_trie_store_interned(parsed, store) + + stats = store.build_stats + assert stats is not None + assert stats.segment_count == 4, ( + "spawn_minimal interns 4 unique content segments (system+user per node, " + "two nodes); a change here means the trie/corpus shape moved -- re-pin " + "deliberately" + ) + assert stats.content_bytes == 146, ( + "total interned content blob bytes for spawn_minimal; drifts with an " + "intentional corpus-text change -- re-pin deliberately" + ) + assert stats.node_manifest_count == 2, ( + "one manifest per trie LlmNode (root + spawned child); a change means " + "the node set moved -- re-pin deliberately" + ) + + +@pytest.mark.asyncio +async def test_cross_payload_overlapping_segment_counted_once(tmp_path: Path) -> None: + """A segment id shipped in TWO payloads is interned once (put_segment dedup). + + This is the memory property the snapshot guards: at weka corpus scale + per-trace workers re-ship shared-prefix segments, so ``segment_count`` and + ``content_bytes`` must count each content-addressed id exactly once no matter + how many payloads carry it. + """ + shared = ("seg-shared", "user", "shared-content", None) + a_only = ("seg-a", "user", "a-only", None) + b_only = ("seg-b", "user", "b-only", None) + + payloads = [ + TraceSegmentPayload( + trace_id="t0", + node_ordinals={"n0": 0}, + envelopes=[ + NodeEnvelope(0, "profiling", _envelope_bytes(["seg-shared", "seg-a"])) + ], + segments=[shared, a_only], + ), + TraceSegmentPayload( + trace_id="t1", + node_ordinals={"n0": 0}, + envelopes=[ + NodeEnvelope(0, "profiling", _envelope_bytes(["seg-shared", "seg-b"])) + ], + segments=[shared, b_only], + ), + ] + + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id="dedup") + await build_unified_trie_store_from_payloads(payloads, store) + + stats = store.build_stats + assert stats is not None + # 4 segment tuples supplied across the two payloads, but seg-shared repeats. + assert stats.segment_count == 3, "overlapping segment must be interned once" + + expected_content_bytes = sum( + len(orjson.dumps({"role": role, "content": content})) + for _id, role, content, _wire in (shared, a_only, b_only) + ) + assert stats.content_bytes == expected_content_bytes, ( + "content_bytes must count the shared segment's blob once, not twice" + ) + + +class _LogCaptureStub: + """Minimal GraphStoreBuilder stand-in for the streaming build log site. + + Carries only what ``_build_graph_store_streaming_trie`` reads from ``self``, + capturing ``info`` calls; the merge/sidecar hooks are no-ops so the drain + runs and its build-complete log line (which fires before the merge) lands. + """ + + def __init__(self) -> None: + self.run = SimpleNamespace(benchmark_id="bench-stats-log") + self.infos: list[str] = [] + + def info(self, msg: object) -> None: + self.infos.append(msg() if callable(msg) else msg) + + def _merge_structural_graphs(self, structural_sink: list[bytes]) -> object: + return object() + + def _write_graph_sidecar( + self, + merged: object, + catalog: dict[str, dict[str, int]], + base_path: Path, + ) -> None: + return None + + +@pytest.mark.asyncio +async def test_streaming_drain_logs_build_stats(tmp_path: Path) -> None: + """The streaming build-complete log line carries the formatted store stats.""" + payload = TraceSegmentPayload( + trace_id="t0", + node_ordinals={"n0": 0}, + envelopes=[NodeEnvelope(0, "profiling", _envelope_bytes(["seg-x"]))], + segments=[("seg-x", "user", "hello", None)], + ) + stub = _LogCaptureStub() + + await GraphStoreBuilder._build_graph_store_streaming_trie(stub, [payload], tmp_path) + + build_logs = [m for m in stub.infos if "UNIFIED store built (streaming)" in m] + assert len(build_logs) == 1, stub.infos + line = build_logs[0] + assert "segments=1" in line, line + assert "content_bytes=" in line, line + assert "node_manifests=1" in line, line + assert "peak_rss_mib=" in line, line + + +@pytest.mark.asyncio +async def test_streaming_drain_failure_aborts_and_removes_store_dir( + tmp_path: Path, +) -> None: + """A mid-drain exception aborts the store and removes its dir. + + The store spills ``content.blob`` incrementally, so a drain that raises + before finalize leaves a partial blob on disk. ``GraphStoreBuilder`` wraps + the drain in ``abort()`` + ``rmtree``, so no half-written store survives for + a later open. The error re-surfaces unchanged. + """ + + def _raising_payloads(): + yield TraceSegmentPayload( + trace_id="t0", + node_ordinals={"n0": 0}, + envelopes=[NodeEnvelope(0, "profiling", _envelope_bytes(["seg-x"]))], + segments=[("seg-x", "user", "hello", None)], + ) + raise RuntimeError("boom mid-drain") + + stub = _LogCaptureStub() + stub.run = SimpleNamespace(benchmark_id="drain-abort") + + with pytest.raises(RuntimeError, match="boom mid-drain"): + await GraphStoreBuilder._build_graph_store_streaming_trie( + stub, _raising_payloads(), tmp_path + ) + + assert not (tmp_path / "aiperf_graph_segments_drain-abort").exists() + + +def test_format_store_build_stats_none_renders_unavailable() -> None: + """A pre-finalize store (``build_stats is None``) logs a marker, not a crash.""" + assert "build_stats=unavailable" in _format_store_build_stats(None) + + +def test_format_store_build_stats_peak_rss_none_renders_na() -> None: + """``peak_rss_mib=None`` (Windows / unavailable) renders as ``n/a``.""" + stats = GraphStoreBuildStats( + segment_count=0, + content_bytes=0, + node_manifest_count=0, + manifest_bytes=0, + trace_count=0, + peak_rss_mib=None, + ) + assert "n/a" in _format_store_build_stats(stats) + + +@pytest.mark.skipif( + store_mod.IS_WINDOWS, reason="resource module is unavailable on Windows" +) +def test_peak_rss_mib_macos_converts_bytes_to_mib( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """macOS ``ru_maxrss`` is BYTES; the Darwin branch divides by 1024*1024.""" + import resource + + monkeypatch.setattr(store_mod, "IS_WINDOWS", False) + monkeypatch.setattr(store_mod, "IS_MACOS", True) + monkeypatch.setattr( + resource, "getrusage", lambda who: SimpleNamespace(ru_maxrss=10 * 1024 * 1024) + ) + assert store_mod._peak_rss_mib() == pytest.approx(10.0) + + +@pytest.mark.skipif( + store_mod.IS_WINDOWS, reason="resource module is unavailable on Windows" +) +def test_peak_rss_mib_linux_converts_kib_to_mib( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Linux ``ru_maxrss`` is KiB; the non-Darwin branch divides by 1024.""" + import resource + + monkeypatch.setattr(store_mod, "IS_WINDOWS", False) + monkeypatch.setattr(store_mod, "IS_MACOS", False) + monkeypatch.setattr( + resource, "getrusage", lambda who: SimpleNamespace(ru_maxrss=10 * 1024) + ) + assert store_mod._peak_rss_mib() == pytest.approx(10.0) diff --git a/tests/unit/dataset/test_graph_store_builder.py b/tests/unit/dataset/test_graph_store_builder.py new file mode 100644 index 0000000000..db5350420e --- /dev/null +++ b/tests/unit/dataset/test_graph_store_builder.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``GraphStoreBuilder.build`` end-to-end contract on a real run. + +The builder owns the ONE graph store build the DatasetManager used to inline: +given a run and a graph path it must build the unified segment store at the +run's store root, write the mandatory graph_meta sidecar, and hand back the +graph facet + both locations as a ``GraphStoreBuildResult``. These tests drive +a REAL ``GraphStoreBuilder(run)`` (``make_run_from_cli``, no stubs) over the +weka_min fixture and check every field of that result against the on-disk +artifacts. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from aiperf.common.environment import Environment +from aiperf.common.exceptions import DatasetError +from aiperf.config.flags.cli_config import CLIConfig +from aiperf.dataset.graph import store_build +from aiperf.dataset.graph.graph_meta_sidecar import sidecar_path_for +from aiperf.dataset.graph.store_build import GraphStoreBuilder +from aiperf.dataset.graph_segment_unified_store import GraphSegmentUnifiedClient +from tests.unit.conftest import make_run_from_cli + +WEKA_MIN = Path(__file__).parents[1] / "graph" / "fixtures" / "weka_min.json" + + +@pytest.fixture +def store_root(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + """Redirect the store root to tmp_path so build artifacts land in a known dir.""" + monkeypatch.setattr(Environment.DATASET, "MMAP_BASE_PATH", tmp_path) + return tmp_path + + +@pytest.mark.asyncio +async def test_build_weka_min_returns_facet_sidecar_and_openable_store( + store_root: Path, +) -> None: + """A real weka_min build yields the trace universe, the sidecar, and a store + the worker-side unified client can open and read node envelopes from.""" + run = make_run_from_cli( + CLIConfig( + model_names=["test-model"], + input_file=str(WEKA_MIN), + tokenizer_name="builtin", + ) + ) + + result = await GraphStoreBuilder(run).build(WEKA_MIN) + + assert result.base_path == store_root + assert result.facet.trace_ids == ["trace_03_n3"] + assert result.facet.prefix_cache_by_trace == { + "trace_03_n3": { + "trace_03_n3:0": [0, 2], + "trace_03_n3:1": [2, 3], + "trace_03_n3:2": [3, 4], + } + } + assert result.sidecar_path == sidecar_path_for(store_root, run.benchmark_id) + assert result.sidecar_path.exists() + with GraphSegmentUnifiedClient(store_root, run.benchmark_id).open() as client: + assert "trace_03_n3" in client._node_offsets + assert client.get_node_envelope("trace_03_n3", 0, "profiling") + + +@pytest.mark.asyncio +async def test_build_without_sidecar_write_raises( + store_root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A drain that never records the sidecar path is a hard build failure. + + The sidecar is mandatory for graph runs (the TimingManager only ingests + it), so ``build()`` enforces the recorded path instead of handing back a + result with a dead location. + """ + del store_root + run = make_run_from_cli( + CLIConfig(model_names=["test-model"], input_file=str(WEKA_MIN)) + ) + builder = GraphStoreBuilder(run) + monkeypatch.setattr( + store_build, "publish_graph_loader_tokenizer_env", lambda run: None + ) + monkeypatch.setattr(builder, "_prestart_loader_forkserver", lambda: None) + + async def _no_sidecar_drain( + graph_path: Path, base_path: Path, fmt: str | None + ) -> tuple: + return {"t": {"n": 0}}, None + + monkeypatch.setattr(builder, "_build_graph_store_streaming", _no_sidecar_drain) + monkeypatch.setattr( + builder, "_build_graph_prefix_cache_by_trace", lambda prefix_source: {} + ) + + with pytest.raises(DatasetError, match="sidecar"): + await builder.build(WEKA_MIN) diff --git a/tests/unit/dataset/test_loader_mp_context.py b/tests/unit/dataset/test_loader_mp_context.py new file mode 100644 index 0000000000..62069a5c32 --- /dev/null +++ b/tests/unit/dataset/test_loader_mp_context.py @@ -0,0 +1,257 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Loader mp-context platform branch (DM6) + preload env triple (DM4). + +Platform-conditional code must branch on ``IS_LINUX`` from +``aiperf.common.constants`` -- never on ``platform.system()`` inline. The +constant is also what makes the branch patchable here. + +The preload env triple ``(name, trust_remote_code, revision)`` is how the run +config reaches the forkserver helper AND every pool worker; DM4 was the loader +silently preloading/synthesizing with ``(False, "main")`` for runs pinned to a +non-``main`` tokenizer revision, and resolving aliases differently on a preload +hit vs an on-demand fallback. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from types import SimpleNamespace + +import pytest + +import aiperf.dataset._mp_context as mp_context +import aiperf.dataset._tokenizer_preload as tokenizer_preload +from aiperf.common.constants import IS_WINDOWS +from aiperf.common.tokenizer import Tokenizer +from aiperf.config.tokenizer import TokenizerConfig +from aiperf.dataset.graph.adapters.shared.content import CorpusContentSynthesizer +from aiperf.dataset.graph.workload_detect import publish_graph_loader_tokenizer_env + +_ENV_VARS = ( + mp_context._ENV_PRELOAD_NAME, + mp_context._ENV_PRELOAD_TRUST, + mp_context._ENV_PRELOAD_REVISION, +) + + +@pytest.fixture +def clean_loader_env(monkeypatch: pytest.MonkeyPatch) -> pytest.MonkeyPatch: + """Fresh context cache + no preload env vars, spawn context (no forkserver).""" + monkeypatch.setattr(mp_context, "_loader_ctx", None) + monkeypatch.setattr(mp_context, "IS_LINUX", False) + for var in _ENV_VARS: + monkeypatch.delenv(var, raising=False) + return monkeypatch + + +def test_get_loader_mp_context_uses_is_linux_constant( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """With IS_LINUX False the loader context must be a spawn context (no + forkserver helper). Patching the module constant only works because the + module imports IS_LINUX; the inline platform.system() call it replaced was + not patchable and violated the constants contract.""" + monkeypatch.setattr(mp_context, "_loader_ctx", None) + monkeypatch.setattr(mp_context, "IS_LINUX", False) + + ctx = mp_context.get_loader_mp_context() + + assert ctx.get_start_method() == "spawn" + + +def test_get_loader_mp_context_defaults_env_triple_when_unconfigured( + clean_loader_env: pytest.MonkeyPatch, +) -> None: + mp_context.get_loader_mp_context(preload_tokenizer="some/tok") + + assert os.environ[mp_context._ENV_PRELOAD_NAME] == "some/tok" + assert os.environ[mp_context._ENV_PRELOAD_TRUST] == "false" + assert os.environ[mp_context._ENV_PRELOAD_REVISION] == "main" + + +def test_get_loader_mp_context_respects_configured_trust_and_revision( + clean_loader_env: pytest.MonkeyPatch, +) -> None: + """The name-only call (the only production call shape, via + ``_loader_pool_context``) must NOT clobber the run-config trust/revision + published by ``configure_loader_tokenizer_env``.""" + mp_context.configure_loader_tokenizer_env( + trust_remote_code=True, revision="deadbeef" + ) + mp_context.get_loader_mp_context(preload_tokenizer="some/tok") + + assert os.environ[mp_context._ENV_PRELOAD_NAME] == "some/tok" + assert os.environ[mp_context._ENV_PRELOAD_TRUST] == "true" + assert os.environ[mp_context._ENV_PRELOAD_REVISION] == "deadbeef" + + +def test_get_loader_mp_context_explicit_args_override_configured_env( + clean_loader_env: pytest.MonkeyPatch, +) -> None: + mp_context.configure_loader_tokenizer_env(trust_remote_code=True, revision="r1") + mp_context.get_loader_mp_context( + preload_tokenizer="some/tok", trust_remote_code=False, revision="r2" + ) + + assert os.environ[mp_context._ENV_PRELOAD_TRUST] == "false" + assert os.environ[mp_context._ENV_PRELOAD_REVISION] == "r2" + + +@pytest.mark.skipif(IS_WINDOWS, reason="forkserver context does not exist on Windows") +def test_get_loader_mp_context_second_call_skips_forkserver_start( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The cached context must not re-enter the stdio dup2 swap. + + The GraphStoreBuilder pre-starts the forkserver on the event loop (a + known-quiet point) precisely so the pool's later call from inside the + offloaded ``asyncio.to_thread`` build is a cached no-op; a second + ``_eagerly_start_forkserver`` there would race the process-wide dup2 + stdio swap against live event-loop logging. + """ + starts: list[None] = [] + monkeypatch.setattr(mp_context, "_loader_ctx", None) + monkeypatch.setattr(mp_context, "IS_LINUX", True) + monkeypatch.setattr( + mp_context, "_eagerly_start_forkserver", lambda: starts.append(None) + ) + + first = mp_context.get_loader_mp_context() + second = mp_context.get_loader_mp_context(preload_tokenizer="some/tok") + + assert first is second + assert len(starts) == 1 + + +def test_parse_graph_workload_publishes_env_triple( + clean_loader_env: pytest.MonkeyPatch, +) -> None: + """parse_graph_workload is the one seam every graph parse goes through, so + it must publish the run's tokenizer trust/revision triple itself: a direct + caller (tooling, tests) has no DatasetManager configure step to do it, and + the forkserver helper snapshots the env once at spawn. + + Real resolved config on purpose (a mock would hide wrong-path reads); only + the HF tokenizer load itself is mocked. + """ + monkeypatch = clean_loader_env + from aiperf.config.flags.cli_config import CLIConfig + from aiperf.dataset.graph.workload_detect import parse_graph_workload + from tests.unit.conftest import make_run_from_cli + + weka_min = Path(__file__).parent.parent / "graph" / "fixtures" / "weka_min.json" + run = make_run_from_cli( + CLIConfig( + model_names=["test-model"], + input_file=str(weka_min), + tokenizer_name="builtin", + trust_remote_code=True, + tokenizer_revision="pinned-rev", + ) + ) + + # A direct caller with no DatasetManager configure step (tooling / tests). + parse_graph_workload(run, weka_min) + + assert os.environ[mp_context._ENV_PRELOAD_TRUST] == "true" + assert os.environ[mp_context._ENV_PRELOAD_REVISION] == "pinned-rev" + + # ...and the synthesizer layer loads with exactly that triple: the + # worker-role fallback in ``content._build_generator`` on a preload miss. + calls: list[dict] = [] + + def recording_from_pretrained(name: str, **kwargs) -> SimpleNamespace: + calls.append({"name": name, **kwargs}) + return SimpleNamespace() + + class _StubGenerator: + def __init__(self, config, tokenizer) -> None: + self.config = config + self.tokenizer = tokenizer + self._hash_id_corpus_rng = None + + monkeypatch.setattr(Tokenizer, "from_pretrained", recording_from_pretrained) + monkeypatch.setattr( + "aiperf.dataset._tokenizer_preload.get_preloaded", lambda *a, **k: None + ) + monkeypatch.setattr( + "aiperf.dataset.generator.coding_content.CodingContentGenerator", + _StubGenerator, + ) + CorpusContentSynthesizer._build_generator("gpt2", "coding") + + assert calls, "fallback tokenizer load must fire" + assert calls[0]["trust_remote_code"] is True + assert calls[0]["revision"] == "pinned-rev" + + +def test_store_builder_publishes_loader_tokenizer_env_from_run_config( + clean_loader_env: pytest.MonkeyPatch, +) -> None: + """Real TokenizerConfig on purpose: a mock would hide attribute-path drift.""" + tokenizer_config = TokenizerConfig(trust_remote_code=True, revision="pinned-rev") + run = SimpleNamespace(cfg=SimpleNamespace(tokenizer=tokenizer_config)) + + publish_graph_loader_tokenizer_env(run) + mp_context.get_loader_mp_context(preload_tokenizer="some/tok") + + assert os.environ[mp_context._ENV_PRELOAD_TRUST] == "true" + assert os.environ[mp_context._ENV_PRELOAD_REVISION] == "pinned-rev" + + +def test_preload_and_fallback_resolve_the_same_tokenizer( + clean_loader_env: pytest.MonkeyPatch, +) -> None: + """A preload hit and an on-demand fallback miss must load with the SAME + ``(name, trust_remote_code, revision, resolve_alias)`` arguments, or an + aliased name silently synthesizes with a different tokenizer depending on + whether the forkserver preload fired.""" + monkeypatch = clean_loader_env + calls: list[dict] = [] + + def recording_from_pretrained(name: str, **kwargs) -> SimpleNamespace: + calls.append({"name": name, **kwargs}) + return SimpleNamespace() + + # One patch covers both call sites: content.py's module-level ``Tokenizer`` + # and _tokenizer_preload's local import are the same class object. + monkeypatch.setattr(Tokenizer, "from_pretrained", recording_from_pretrained) + monkeypatch.setattr(tokenizer_preload, "_LOADED", {}) + monkeypatch.setenv(mp_context._ENV_PRELOAD_NAME, "gpt2") + monkeypatch.setenv(mp_context._ENV_PRELOAD_TRUST, "true") + monkeypatch.setenv(mp_context._ENV_PRELOAD_REVISION, "pinned-rev") + + # Forkserver-helper role: the preload load. + tokenizer_preload._preload() + + # Worker role on a preload MISS: the on-demand fallback in + # content._build_generator. Stub out the corpus generator so only the + # tokenizer-load call shape is exercised. + class _StubGenerator: + def __init__(self, config, tokenizer) -> None: + self.config = config + self.tokenizer = tokenizer + self._hash_id_corpus_rng = None + + monkeypatch.setattr( + "aiperf.dataset.generator.coding_content.CodingContentGenerator", + _StubGenerator, + ) + monkeypatch.setattr( + "aiperf.dataset._tokenizer_preload.get_preloaded", lambda *a, **k: None + ) + CorpusContentSynthesizer._build_generator("gpt2", "coding") + + assert len(calls) == 2, f"expected preload + fallback loads, got {calls}" + preload_call, fallback_call = calls + assert preload_call["name"] == fallback_call["name"] == "gpt2" + for key in ("trust_remote_code", "revision", "resolve_alias"): + assert preload_call[key] == fallback_call[key], ( + f"{key} diverges between preload ({preload_call[key]!r}) and " + f"fallback ({fallback_call[key]!r})" + ) + assert preload_call["trust_remote_code"] is True + assert preload_call["revision"] == "pinned-rev" + assert preload_call["resolve_alias"] is True diff --git a/tests/unit/dataset/test_nonweka_interned_route.py b/tests/unit/dataset/test_nonweka_interned_route.py new file mode 100644 index 0000000000..31ef5d3d44 --- /dev/null +++ b/tests/unit/dataset/test_nonweka_interned_route.py @@ -0,0 +1,334 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Non-weka interned-route contract: every non-weka graph build (dynamo / +native / dag_jsonl) parses ONCE in-process and drains that SAME parse through +the eager interned builder, writing the mandatory graph_meta sidecar DIRECTLY +from the stripped parse; weka alone still takes the worker-pool payload stream + +structural merge. + +These pins cover the route face that the byte-parity oracle +(``test_dag_jsonl_streaming_store_parity`` / ``test_dynamo_streaming_store_parity``) +does not: the returned prefix source identity, the persisted store's per-trace +topology, the sidecar's PARSE-ORDER trace list (parse order is the contract -- +a sorted-by-id list would silently reorder traces), the sidecar's +prefix-cache-count round trip, and the weka-only routing inverse. +""" + +from __future__ import annotations + +from pathlib import Path +from types import MethodType, SimpleNamespace + +import orjson +import pytest + +from aiperf.dataset.graph.adapters.dag_jsonl.trace import from_dag_jsonl +from aiperf.dataset.graph.adapters.dynamo.trace import from_dynamo_trace +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.codecs import decode_graph_meta_sidecar +from aiperf.dataset.graph.graph_meta_sidecar import catalogs_match +from aiperf.dataset.graph.models import LlmNode, resolve_trace_graph +from aiperf.dataset.graph.segment_ir.store_builder import ( + graph_carries_assembly_slots, + iter_trace_segment_payloads, +) +from aiperf.dataset.graph.store_build import GraphStoreBuilder +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedClient, + _encode_inner_key, +) +from tests.unit.dataset.test_dag_jsonl_streaming_store_parity import ( + MULTI_TRACE_DAG_LINES, +) + +WEKA_FIXTURE = Path(__file__).parents[1] / "graph" / "fixtures" / "weka_min.json" + +# Two INDEPENDENT lone-turn traces whose ids sort in the REVERSE of parse order +# ("aa-second" < "zz-first"), so a route that preserves parse order and one that +# sorts by id produce DIFFERENT trace lists -- the pin for the +# parse-order contract. +PARSE_ORDER_DAG_LINES = """\ +{"session_id":"zz-first","turns":[{"model":"Qwen3-0.6B","messages":[{"role":"system","content":"zz-sys"},{"role":"user","content":"zz-u"}],"max_tokens":20}]} +{"session_id":"aa-second","turns":[{"model":"Qwen3-0.6B","messages":[{"role":"system","content":"aa-sys"},{"role":"user","content":"aa-u"}],"max_tokens":20}]} +""" + + +def _dynamo_record(ts: int, sid: str, input_tokens: int, hashes: list[int]) -> dict: + """A minimal dynamo request-end record that stamps prefix-cache hash ids.""" + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": {"session_id": sid}, + "request": { + "request_id": f"r{ts}", + "model": "m", + "input_tokens": input_tokens, + "output_tokens": 8, + "cached_tokens": 0, + "replay": { + "trace_block_size": 16, + "input_length": input_tokens, + "input_sequence_hashes": hashes, + }, + }, + } + + +def _make_interned_route_stub(benchmark_id: str) -> SimpleNamespace: + """Stub carrying only what the in-process interned branch reads from self. + + Binds the REAL interned-drain/sidecar helpers; the weka trie drain and + structural merge fail loudly if reached (no non-weka format may take them). + """ + stub = SimpleNamespace( + run=SimpleNamespace(benchmark_id=benchmark_id), + info=lambda *a, **k: None, + warning=lambda *a, **k: None, + _sidecar_path=None, + ) + for name in ("_write_graph_sidecar", "_build_interned_unified_store"): + setattr(stub, name, MethodType(getattr(GraphStoreBuilder, name), stub)) + + async def _fail_trie(payloads, base_path): # noqa: ANN001, ARG001 + raise AssertionError( + "non-weka format must not take the weka trie payload drain" + ) + + def _fail_merge(structural_sink): # noqa: ANN001, ARG001 + raise AssertionError("non-weka format must not merge a structural stream") + + stub._build_graph_store_streaming_trie = _fail_trie + stub._merge_structural_graphs = _fail_merge + return stub + + +@pytest.mark.asyncio +async def test_multi_trace_interned_route_persists_both_topologies( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A multi-trace slot-free dag_jsonl file takes the interned route and the + persisted store serves each trace's OWN topology. + + ``_build_graph_store_streaming`` returns the FULL parse as the prefix source + (identity), the catalog covers both traces, and -- read back from the + unified store index -- each trace's node ordinals match its own catalog + (a merge-collapse bug would give both traces one tree's single-node graph). + """ + from aiperf.dataset.graph import workload_detect + + src = tmp_path / "multi.dag.jsonl" + src.write_text(MULTI_TRACE_DAG_LINES) + parsed = from_dag_jsonl(str(src)) + assert len(parsed.traces) == 2 + assert not graph_carries_assembly_slots(parsed) + # Premise: the two traces genuinely carry DIFFERENT topologies. + node_sets = {t.id: set(resolve_trace_graph(parsed, t).nodes) for t in parsed.traces} + assert node_sets["conv-a"] != node_sets["conv-b"] + monkeypatch.setattr( + workload_detect, "parse_graph_workload", lambda run, path: parsed + ) + + stub = _make_interned_route_stub("bench-multi-topo") + catalog, returned = await GraphStoreBuilder._build_graph_store_streaming( + stub, src, tmp_path, "dag_jsonl" + ) + + assert returned is parsed + assert set(catalog) == {"conv-a", "conv-b"} and len(catalog) == 2 + # Store-side topology check: the persisted index carries each trace's own + # node ordinals and its envelopes materialize. + with GraphSegmentUnifiedClient(tmp_path, "bench-multi-topo").open() as client: + for trace_id, ordinals in catalog.items(): + store_keys = set(client._node_offsets.get(trace_id, {})) + expected_keys = { + _encode_inner_key(ordinal, "profiling") for ordinal in ordinals.values() + } + assert store_keys == expected_keys and store_keys + for ordinal in ordinals.values(): + assert client.get_node_envelope(trace_id, ordinal, "profiling") + # The two traces did NOT collapse to a single shared topology: distinct + # per-trace node-id sets survived into the store's build catalog. + assert set(catalog["conv-a"]) != set(catalog["conv-b"]) + + +@pytest.mark.asyncio +async def test_interned_route_sidecar_preserves_parse_order( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The graph_meta sidecar ships traces in PARSE order, content-free. + + The contract: the in-process interned route writes the sidecar + DIRECTLY from the stripped parse, so the loaded trace list is the parse + order ("zz-first", "aa-second"), NOT a + sorted-by-id order. The loaded graph is content-free (empty pool, empty node + prompts) and its catalog still matches the build catalog. + """ + from aiperf.dataset.graph import workload_detect + + src = tmp_path / "order.dag.jsonl" + src.write_text(PARSE_ORDER_DAG_LINES) + parsed = from_dag_jsonl(str(src)) + assert [t.id for t in parsed.traces] == ["zz-first", "aa-second"] + assert not graph_carries_assembly_slots(parsed) + monkeypatch.setattr( + workload_detect, "parse_graph_workload", lambda run, path: parsed + ) + + stub = _make_interned_route_stub("bench-order") + catalog, returned = await GraphStoreBuilder._build_graph_store_streaming( + stub, src, tmp_path, "dag_jsonl" + ) + assert returned is parsed + assert stub._sidecar_path is not None + + loaded, _fingerprint, _version = decode_graph_meta_sidecar( + stub._sidecar_path.read_bytes() + ) + + # PARSE order, not sorted-by-id (sorted would be ["aa-second", "zz-first"]). + assert [t.id for t in loaded.traces] == ["zz-first", "aa-second"] + # Content-free: the pool is emptied and every LlmNode prompt is stripped. + assert not loaded.segment_pool.by_id + for trace in loaded.traces: + for node in resolve_trace_graph(loaded, trace).nodes.values(): + if isinstance(node, LlmNode): + assert node.prompt == [] + # The loaded structural graph still describes the build's topology. + assert catalogs_match(loaded, catalog) + + +@pytest.mark.asyncio +async def test_interned_route_sidecar_roundtrips_prefix_cache_counts( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The sidecar preserves the per-node prefix-cache counts the prefix map reads. + + A dynamo parse stamps ``metadata["dispatch"]`` prefix-cache counts. The + prefix map computed from the LOADED sidecar graph must equal the map from + the live parse -- proving ``strip_replay_text`` + encode/decode preserve the + dispatch metadata the prefix-cache metric consumes. + """ + from aiperf.dataset.graph import workload_detect + + dyn = tmp_path / "dyn.jsonl" + dyn.write_bytes( + b"\n".join( + orjson.dumps(r) + for r in ( + _dynamo_record(1000, "s1", 32, [111, 222]), + _dynamo_record(2000, "s1", 64, [111, 222, 333, 444]), + _dynamo_record(3000, "s2", 48, [555, 666, 777]), + ) + ) + ) + parsed = from_dynamo_trace(dyn, content_root_seed=42, content_tokenizer="builtin") + eager_map = GraphStoreBuilder._build_graph_prefix_cache_by_trace(parsed) + assert eager_map, "dynamo fixture must stamp a non-empty prefix-cache map" + monkeypatch.setattr( + workload_detect, "parse_graph_workload", lambda run, path: parsed + ) + + stub = _make_interned_route_stub("bench-prefix") + _catalog, returned = await GraphStoreBuilder._build_graph_store_streaming( + stub, dyn, tmp_path, "dynamo" + ) + assert returned is parsed + assert stub._sidecar_path is not None + + loaded, _fingerprint, _version = decode_graph_meta_sidecar( + stub._sidecar_path.read_bytes() + ) + loaded_map = GraphStoreBuilder._build_graph_prefix_cache_by_trace(loaded) + assert loaded_map == eager_map + + +@pytest.mark.asyncio +async def test_weka_route_still_takes_payload_stream_and_merge( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The weka-only inverse: ``fmt="weka_trace"`` takes the worker-pool payload + stream + structural merge, NOT the in-process interned drain. + + Canned ``TraceSegmentPayload``s (captured from a real weka parse so the + structural graph and catalog stay consistent) replace the streaming source; + the route must drain them through ``_build_graph_store_streaming_trie``, so + the returned prefix source IS the MERGED structural graph (content-free, + identity-checked), the sidecar is written, and the interned in-process drain + is never reached. + """ + from aiperf.dataset.graph import workload_detect + from aiperf.dataset.graph.adapters.weka import trace as weka_trace + from aiperf.dataset.graph.parse_context import GraphParseContext + + parsed_weka = from_weka_trace(str(WEKA_FIXTURE), content_root_seed=0) + canned_payloads = list(iter_trace_segment_payloads(parsed_weka)) + assert canned_payloads and any(p.structural_graph for p in canned_payloads) + + monkeypatch.setattr( + weka_trace, + "stream_weka_trace_segment_payloads", + lambda source, **kwargs: iter(canned_payloads), + ) + # The route resolves the run-derived knob bundle (the ONE + # ``resolve_graph_parse_context`` resolution) before calling the (patched) + # stream; neutralize it so the bare stub run suffices. + monkeypatch.setattr( + workload_detect, + "resolve_graph_parse_context", + lambda run: GraphParseContext(content_root_seed=0, idle_gap_cap_seconds=None), + ) + + stub = SimpleNamespace( + run=SimpleNamespace(benchmark_id="bench-weka-route"), + info=lambda *a, **k: None, + warning=lambda *a, **k: None, + _sidecar_path=None, + ) + for name in ( + "_build_graph_store_streaming_trie", + "_write_graph_sidecar", + ): + setattr(stub, name, MethodType(getattr(GraphStoreBuilder, name), stub)) + + captured: dict[str, object] = {} + real_merge = MethodType(GraphStoreBuilder._merge_structural_graphs, stub) + + def _capturing_merge(structural_sink): # noqa: ANN001 + merged = real_merge(structural_sink) + captured["merged"] = merged + return merged + + stub._merge_structural_graphs = _capturing_merge + + async def _fail_interned(parsed, unified): # noqa: ANN001, ARG001 + raise AssertionError("weka_trace must not take the in-process interned drain") + + stub._build_interned_unified_store = _fail_interned + monkeypatch.setattr( + workload_detect, + "parse_graph_workload", + lambda run, path: (_ for _ in ()).throw( + AssertionError("weka_trace must not re-parse in-process") + ), + ) + + catalog, returned = await GraphStoreBuilder._build_graph_store_streaming( + stub, WEKA_FIXTURE, tmp_path, "weka_trace" + ) + + # Routed through the trie drain: the returned prefix source IS the merged + # structural graph, not the parse. + assert returned is captured["merged"] + assert returned is not parsed_weka + # The merged structural graph is content-free (empty pool, stripped prompts). + assert not returned.segment_pool.by_id + assert set(catalog) == {t.id for t in parsed_weka.traces} and catalog + assert stub._sidecar_path is not None + assert stub._sidecar_path.exists() + # The store serves the drained trace's node envelopes. + with GraphSegmentUnifiedClient(tmp_path, "bench-weka-route").open() as client: + for trace_id, ordinals in catalog.items(): + for ordinal in ordinals.values(): + assert client.get_node_envelope(trace_id, ordinal, "profiling") diff --git a/tests/unit/dataset/test_streaming_trie_sidecar.py b/tests/unit/dataset/test_streaming_trie_sidecar.py new file mode 100644 index 0000000000..caa7a8a3e4 --- /dev/null +++ b/tests/unit/dataset/test_streaming_trie_sidecar.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``GraphStoreBuilder._write_graph_sidecar`` mandatory-write contract. + +The TimingManager only ingests the graph_meta sidecar from the path the +graph-typed dataset broadcast advertises (no re-parse fallback), so every +graph build route MUST land the file, record its path, or fail the run. +""" + +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from aiperf.common.exceptions import DatasetError +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.codecs import encode_parsed_graph_msgpack +from aiperf.dataset.graph.graph_meta_sidecar import ( + sidecar_path_for, + strip_replay_text, +) +from aiperf.dataset.graph.graph_path_catalog import build_graph_path_catalog +from aiperf.dataset.graph.merge import merge_parsed_graphs +from aiperf.dataset.graph.store_build import GraphStoreBuilder + +WEKA_FIXTURE = Path(__file__).parents[1] / "graph" / "fixtures" / "weka_min.json" + + +class _StubManager: + """Just the attributes ``_write_graph_sidecar`` reads/writes on self.""" + + def __init__(self) -> None: + self.run = SimpleNamespace(benchmark_id="bench-sidecar-test") + self._sidecar_path: Path | None = None + self.infos: list[object] = [] + + def info(self, msg: object) -> None: + self.infos.append(msg() if callable(msg) else msg) + + +@pytest.mark.skipif(not WEKA_FIXTURE.exists(), reason="weka fixture missing") +def test_write_graph_sidecar_writes_and_records_path(tmp_path): + parsed = from_weka_trace(WEKA_FIXTURE, content_root_seed=0) + structural = strip_replay_text(parsed) + sink = [encode_parsed_graph_msgpack(structural)] + catalog = build_graph_path_catalog(merge_parsed_graphs([structural])) + manager = _StubManager() + + merged = GraphStoreBuilder._merge_structural_graphs(manager, sink) + GraphStoreBuilder._write_graph_sidecar(manager, merged, catalog, tmp_path) + + expected = sidecar_path_for(tmp_path, "bench-sidecar-test") + assert expected.exists() + assert manager._sidecar_path == expected + + +@pytest.mark.skipif(not WEKA_FIXTURE.exists(), reason="weka fixture missing") +def test_write_graph_sidecar_catalog_mismatch_raises(tmp_path): + parsed = from_weka_trace(WEKA_FIXTURE, content_root_seed=0) + manager = _StubManager() + with pytest.raises(DatasetError, match="catalog"): + GraphStoreBuilder._write_graph_sidecar( + manager, parsed, {"ghost-trace": {"n": 0}}, tmp_path + ) + assert not sidecar_path_for(tmp_path, "bench-sidecar-test").exists() + assert manager._sidecar_path is None + + +def test_merge_structural_graphs_empty_sink_raises(): + manager = _StubManager() + with pytest.raises(DatasetError, match="no structural graphs"): + GraphStoreBuilder._merge_structural_graphs(manager, []) + + +def test_merge_structural_graphs_undecodable_blob_raises(): + manager = _StubManager() + with pytest.raises(DatasetError, match="failed to merge"): + GraphStoreBuilder._merge_structural_graphs(manager, [b"not-msgpack"]) diff --git a/tests/unit/dataset/test_trie_unified_build.py b/tests/unit/dataset/test_trie_unified_build.py new file mode 100644 index 0000000000..b1f2be3ec5 --- /dev/null +++ b/tests/unit/dataset/test_trie_unified_build.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +from pathlib import Path + +import orjson +import pytest + +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.segment_ir.store_builder import ( + _prompt_segment_ids, + _trie_llm_nodes, + build_unified_trie_store_interned, + trie_node_ordinals, +) +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) + +FIXTURES = Path(__file__).parent.parent / "graph" / "fixtures" + + +@pytest.fixture +def one_trie_trace(monkeypatch): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + parsed = from_weka_trace(FIXTURES / "weka_subagent.json") + assert parsed.segment_pool is not None + known_path = None + known_ordinal = None + known_trace = None + for trace in parsed.traces: + nodes = _trie_llm_nodes(parsed, trace) + # Assign ordinals via the SAME builder helper (sort key + # (arrival_offset_us or 0, node_id)) so known_ordinal matches the + # ordinal build_unified_trie_store_interned stores; a hand-rolled + # enumerate(sorted(nodes)) would query a different node's manifest + # whenever arrival order differs from lexical node-id order. + ordinals = trie_node_ordinals(nodes) + for nid, node in nodes.items(): + path = _prompt_segment_ids(node) + if path and (known_path is None or len(path) > len(known_path)): + known_path, known_ordinal, known_trace = path, ordinals[nid], trace.id + assert known_path + return parsed, known_trace, known_ordinal, known_path + + +def _dyn_subagent_fixture(tmp_path: Path) -> Path: + """Current-schema dynamo trace: one root session spawning one subagent. + + ``s_root`` makes several LLM calls; ``s_child`` (``parent_session_id=s_root``) + is a subagent session. The trie adapter flattens BOTH into one graph of + ``{agent_id}_a{k}`` LlmNodes (no ``parsed.subgraphs``), so the built store + must carry a bare-id manifest for every node -- child-session nodes included. + """ + + def _rec(ts: int, sid: str, parent: str | None = None) -> dict: + ctx: dict = {"session_id": sid} + if parent is not None: + ctx["parent_session_id"] = parent + return { + "schema": "dynamo.request.trace.v1", + "event_type": "request_end", + "event_time_unix_ms": ts, + "event_source": "dynamo", + "agent_context": ctx, + "request": { + "request_id": f"{sid}-r{ts}", + "model": "m", + "input_tokens": 32, + "output_tokens": 16, + "cached_tokens": 0, + }, + } + + records = [ + _rec(1000, "s_root"), + _rec(1100, "s_root"), + _rec(1150, "s_child", "s_root"), + _rec(1170, "s_child", "s_root"), + _rec(1200, "s_root"), + ] + p = tmp_path / "dyn_subagent.jsonl" + p.write_bytes(b"\n".join(orjson.dumps(r) for r in records)) + return p + + +@pytest.mark.asyncio +async def test_unified_store_includes_child_session_manifests(tmp_path, monkeypatch): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + from aiperf.dataset.graph.adapters.dynamo.trace import from_dynamo_trace + + p = _dyn_subagent_fixture(tmp_path) + parsed = from_dynamo_trace(p, content_root_seed=1234, content_tokenizer="builtin") + assert parsed.segment_pool is not None + + store = GraphSegmentUnifiedBackingStore(tmp_path, "bench") + catalog = await build_unified_trie_store_interned(parsed, store) + + # Child-session nodes carry bare-id catalog keys next to the root's. + keys = {k for trace_keys in catalog.values() for k in trace_keys} + child_keys = {k for k in keys if k.startswith("s_child:")} + assert len(child_keys) == 2, f"missing child-session keys in {sorted(keys)}" + + # Every node's manifest round-trips out of the store at its ordinal. + with GraphSegmentUnifiedClient(tmp_path, "bench").open() as c: + for trace_id, trace_keys in catalog.items(): + for _key, ordinal in trace_keys.items(): + env = orjson.loads(c.get_node_envelope(trace_id, ordinal, "profiling")) + handles = env["handles"] + msgs = c.materialize_handles(handles) + assert msgs and all(set(m) == {"role", "content"} for m in msgs) + + +@pytest.mark.asyncio +async def test_build_unified_trie_store_interned_round_trips(tmp_path, one_trie_trace): + parsed, trace_id, ordinal, path = one_trie_trace + store = GraphSegmentUnifiedBackingStore(tmp_path, "bench") + catalog = await build_unified_trie_store_interned(parsed, store) + assert trace_id in catalog + + with GraphSegmentUnifiedClient(tmp_path, "bench").open() as c: + env = orjson.loads(c.get_node_envelope(trace_id, ordinal, "profiling")) + handles = env["handles"] + assert len(handles) == len(path) + msgs = c.materialize_handles(handles) + assert all(set(m) == {"role", "content"} for m in msgs) diff --git a/tests/unit/dataset/test_unified_flag_wiring.py b/tests/unit/dataset/test_unified_flag_wiring.py new file mode 100644 index 0000000000..c0d80b5c2e --- /dev/null +++ b/tests/unit/dataset/test_unified_flag_wiring.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unified-store wiring through the GraphStoreBuilder build path. + +Behavioral wiring proofs: the assertions monkeypatch the real builders and +drive the real methods, so a docstring mention of a symbol cannot satisfy them +(the trap the previous ``inspect.getsource`` substring checks fell into). +""" + +from __future__ import annotations + +import inspect +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from aiperf.config.resolution.plan import GraphWorkloadRef, ResolvedConfig +from aiperf.dataset import dataset_manager +from aiperf.dataset.graph import store_build +from aiperf.dataset.graph.models import GraphRecord, ParsedGraph, TraceRecord +from aiperf.dataset.graph.segment_ir import store_builder + +# Recognized weka HF dataset id (marker + nonexistent org/name path); mirrors +# the id the real-ingest suite pins. +_HF_WEKA_ID = "semianalysisai/cc-traces-weka-062126" + + +def _bare_builder( + graph_ref: GraphWorkloadRef | None = None, +) -> store_build.GraphStoreBuilder: + """A GraphStoreBuilder on a stub run for direct method-wiring tests. + + ``build`` reads the run's memoized graph resolution + (``resolve_graph_workload``), so the stub carries a real ``ResolvedConfig``; + pass ``graph_ref`` to preset the memo the production resolver chain (or the + DatasetManager's accessor call) populates before any build. + """ + resolved = ResolvedConfig( + graph_workload=graph_ref, graph_workload_resolved=graph_ref is not None + ) + return store_build.GraphStoreBuilder( + run=SimpleNamespace(benchmark_id="bid", resolved=resolved) + ) + + +@pytest.mark.asyncio +async def test_build_interned_unified_store_calls_interned_builder( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The interned unified build routes through ``build_unified_trie_store_interned``. + + Monkeypatches the node-typed interned builder at its definition site (the + method imports it locally) and asserts the flag path actually CALLS it with + the parse + store it was handed, returning the builder's catalog verbatim. + """ + calls: dict[str, tuple[object, object]] = {} + + async def _fake_builder(parsed: object, unified: object) -> dict: + calls["args"] = (parsed, unified) + return {"trace_a": {"n0": 0}} + + monkeypatch.setattr( + store_builder, "build_unified_trie_store_interned", _fake_builder + ) + + sentinel_parsed, sentinel_store = object(), object() + catalog = await _bare_builder()._build_interned_unified_store( + sentinel_parsed, sentinel_store + ) + + assert calls["args"] == (sentinel_parsed, sentinel_store) + assert catalog == {"trace_a": {"n0": 0}} + + +@pytest.mark.asyncio +async def test_build_routes_hf_id_to_streaming_build( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """An HF weka dataset id reaches the ONE streaming store build as weka_trace. + + ``GraphStoreBuilder.build`` always calls ``_build_graph_store_streaming``; + the wiring under test is the format threading -- the run's memoized + ``weka_trace`` ref (populated from a recognized weka HF repo id at config + resolution; see ``test_graph_workload_resolution``) must arrive as + ``fmt == "weka_trace"`` so the builder takes the worker-pool payload + branch. + """ + builder = _bare_builder( + GraphWorkloadRef(path=Path(_HF_WEKA_ID), format="weka_trace") + ) + called: dict[str, object] = {} + + async def _fake_streaming(path: Path, base_path: Path, fmt: str | None) -> tuple: + called["path"] = path + called["fmt"] = fmt + return {"trace_a": {"n0": 0}}, ParsedGraph( + graph=GraphRecord(), traces=[TraceRecord(id="trace_a", tags=["x"])] + ) + + # The stub run carries only benchmark_id, so the run-derived env/forkserver + # hooks must be bypassed: the module-level tokenizer-env publisher (which + # now takes the run) is monkeypatched at store_build's call site, and the + # forkserver pre-start / drain are instance-shadowed. The sidecar path is + # preset because the shadowed drain never writes one and build() hard-fails + # a sidecar-less build. + monkeypatch.setattr( + store_build, "publish_graph_loader_tokenizer_env", lambda run: None + ) + builder._prestart_loader_forkserver = lambda: None + builder._build_graph_store_streaming = _fake_streaming + builder._sidecar_path = tmp_path / "graph_meta.sidecar" + + result = await builder.build(Path(_HF_WEKA_ID)) + + assert called["path"] == Path(_HF_WEKA_ID) + assert called["fmt"] == "weka_trace" + assert result.facet.trace_ids == ["trace_a"] + assert result.sidecar_path == tmp_path / "graph_meta.sidecar" + + +@pytest.mark.asyncio +async def test_build_prestarts_forkserver_before_build( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The forkserver pre-start must run BEFORE the offloaded store build. + + ``_eagerly_start_forkserver`` dup2-swaps process-wide stdio; if the helper + were first started lazily inside the ``asyncio.to_thread`` build, the swap + would race concurrent event-loop logging. The env publish must come first + so the helper snapshots the run's trust/revision triple. + """ + builder = _bare_builder() + order: list[str] = [] + + async def _fake_streaming(path: Path, base_path: Path, fmt: str | None) -> tuple: + order.append("store-build") + return {"trace_a": {"n0": 0}}, ParsedGraph( + graph=GraphRecord(), traces=[TraceRecord(id="trace_a", tags=["x"])] + ) + + monkeypatch.setattr( + store_build, + "publish_graph_loader_tokenizer_env", + lambda run: order.append("publish-env"), + ) + builder._prestart_loader_forkserver = lambda: order.append("prestart-forkserver") + builder._build_graph_store_streaming = _fake_streaming + builder._sidecar_path = tmp_path / "graph_meta.sidecar" + + local_trace = tmp_path / "trace.json" + local_trace.write_text("{}") + await builder.build(local_trace) + + assert order == ["publish-env", "prestart-forkserver", "store-build"] + + +@pytest.mark.asyncio +async def test_configure_graph_workload_returns_builder_facet( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """``DatasetManager._configure_graph_workload`` hands back the builder's facet. + + The DatasetManager builds nothing itself: the wrapper must run + the endpoint gate, delegate to ``GraphStoreBuilder.build``, and return the + result's facet verbatim (component-integration callers consume + ``.trace_ids`` off it). + """ + sentinel_facet = object() + built: dict[str, object] = {} + + async def _fake_build(self: store_build.GraphStoreBuilder, graph_path: Path): + built["run"] = self.run + built["path"] = graph_path + return SimpleNamespace( + facet=sentinel_facet, + sidecar_path=tmp_path / "sidecar", + base_path=tmp_path, + ) + + monkeypatch.setattr(store_build.GraphStoreBuilder, "build", _fake_build) + + gated: list[object] = [] + from aiperf.dataset.graph import workload_detect + + monkeypatch.setattr(workload_detect, "validate_graph_endpoint_type", gated.append) + + dm = dataset_manager.DatasetManager.__new__(dataset_manager.DatasetManager) + dm.run = SimpleNamespace(benchmark_id="bid") + + facet = await dm._configure_graph_workload(tmp_path / "trace.json") + + assert facet is sentinel_facet + assert built["run"] is dm.run + assert built["path"] == tmp_path / "trace.json" + assert gated == [dm.run], "the endpoint gate must run before the build" + + +def test_store_build_has_no_unsupported_combo_error() -> None: + """Tombstone: the old unsupported-combo error must never come back. + + An ABSENCE can only be asserted against the source text; the positive + routing halves are covered behaviorally above. Checked against BOTH homes + the build logic has lived in. + """ + src = inspect.getsource(dataset_manager.DatasetManager._configure_graph_workload) + assert "WekaUnifiedStoreUnsupportedError" not in src + assert "WekaUnifiedStoreUnsupportedError" not in inspect.getsource(store_build) diff --git a/tests/unit/dataset/test_weka_streaming_gate.py b/tests/unit/dataset/test_weka_streaming_gate.py new file mode 100644 index 0000000000..0f6e2dbd49 --- /dev/null +++ b/tests/unit/dataset/test_weka_streaming_gate.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The streaming store build must recover the per-node prefix-cache map from +the merged structural graphs, and the weka streaming entry point must drain +local directory sources. (The former eager-vs-streaming store-route pins are +superseded by the streaming builder's own branch tests: +``test_weka_content_knob_wiring``, ``test_slot_graph_eager_drain``, and +``test_dag_jsonl_streaming_store_parity`` -- every graph workload now takes +the one streaming store build.)""" + +from __future__ import annotations + +from pathlib import Path + +from aiperf.dataset.graph.adapters.weka.trace import ( + from_weka_trace, + stream_weka_trace_segment_payloads, +) +from aiperf.dataset.graph.codecs import decode_parsed_graph_msgpack +from aiperf.dataset.graph.merge import merge_parsed_graphs +from aiperf.dataset.graph.segment_ir.store_builder import ( + iter_trace_segment_payloads, +) +from aiperf.dataset.graph.store_build import GraphStoreBuilder + +FIX_MIN = Path(__file__).parents[1] / "graph" / "fixtures" / "weka_min.json" +FIX_SUB = Path(__file__).parents[1] / "graph" / "fixtures" / "weka_subagent.json" + + +def test_streaming_structural_prefix_cache_matches_eager_parse() -> None: + parsed = from_weka_trace(str(FIX_MIN), content_root_seed=42) + eager_map = GraphStoreBuilder._build_graph_prefix_cache_by_trace(parsed) + assert eager_map, "fixture must produce a non-empty prefix-cache map" + + structural_blobs = [ + p.structural_graph + for p in iter_trace_segment_payloads(parsed) + if p.structural_graph + ] + merged = merge_parsed_graphs( + decode_parsed_graph_msgpack(b) for b in structural_blobs + ) + streaming_map = GraphStoreBuilder._build_graph_prefix_cache_by_trace(merged) + assert streaming_map == eager_map + + +def test_stream_segment_payloads_local_directory_yields_all_traces(tmp_path) -> None: + """The streaming entry point drains a LOCAL directory of weka .json files. + + Pins the local-source branch the weka streaming store build depends on: a + directory source must stream one payload per file, keyed by each file's + own trace id. + """ + weka_dir = tmp_path / "corpus" + weka_dir.mkdir() + (weka_dir / "a.json").write_bytes(FIX_MIN.read_bytes()) + (weka_dir / "b.json").write_bytes(FIX_SUB.read_bytes()) + + payloads = list( + stream_weka_trace_segment_payloads(str(weka_dir), content_root_seed=42) + ) + + assert sorted(p.trace_id for p in payloads) == [ + "trace_03_n3", + "trace_sub_n2s1", + ] diff --git a/tests/unit/exporters/aggregate/test_aggregate_submission_outcome.py b/tests/unit/exporters/aggregate/test_aggregate_submission_outcome.py new file mode 100644 index 0000000000..8b7ac45d87 --- /dev/null +++ b/tests/unit/exporters/aggregate/test_aggregate_submission_outcome.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Aggregate (multi-run) scenario-submission verdict on the confidence export. + +Covers the cross-run fold ported from AgentX +(``aggregate_confidence_json_exporter.py`` + ``aggregate_base_exporter.py``): +the static scenario-lock outcome combined ACROSS RUNS with the cross-run +context-overflow rate and cancellation, via ``compute_submission_outcome``. + +Uses real ``AggregateResult`` / ``ConfidenceMetric`` / ``ScenarioOutcome`` +objects (no MagicMock) so the metric-key plumbing and the rate-equivalence +between summed totals and per-run means are exercised end-to-end. +""" + +from __future__ import annotations + +import orjson +import pytest + +from aiperf.common.scenario.base import ScenarioOutcome +from aiperf.exporters.aggregate import ( + AggregateConfidenceJsonExporter, + AggregateExporterConfig, +) +from aiperf.orchestrator.aggregation.base import AggregateResult +from aiperf.orchestrator.aggregation.confidence import ConfidenceMetric + + +def _count_metric(mean: float) -> ConfidenceMetric: + """A degenerate count ``ConfidenceMetric`` whose cross-run mean is ``mean``.""" + return ConfidenceMetric( + mean=mean, + std=0.0, + min=mean, + max=mean, + cv=0.0, + se=0.0, + ci_low=mean, + ci_high=mean, + t_critical=float("nan"), + unit="requests", + ) + + +def _make_result(*, metrics: dict, metadata: dict) -> AggregateResult: + """A confidence ``AggregateResult`` over two successful runs.""" + return AggregateResult( + aggregation_type="confidence", + num_runs=2, + num_successful_runs=2, + failed_runs=[], + metrics=dict(metrics), + metadata=dict(metadata), + ) + + +def _export_metadata(tmp_path, result: AggregateResult) -> dict: + """Run the exporter and return the parsed top-level ``metadata`` dict.""" + config = AggregateExporterConfig(result=result, output_dir=tmp_path) + exporter = AggregateConfidenceJsonExporter(config=config) + payload = orjson.loads(exporter._generate_content()) + return payload["metadata"] + + +def test_cross_run_overflow_rate_flips_submission_invalid_from_means(tmp_path): + """Cross-run overflow rate > limit flips submission_valid False from means. + + Neither single run exceeds 1% in isolation, but the cross-run mean rate + (mean(overflow) / mean(total)) does: 3 / 100 = 3% > 1%. Derived purely + from the confidence ``*_avg`` count metrics (no orchestrator carrier keys), + proving the self-sufficient fallback path. + """ + result = _make_result( + metrics={ + "request_count_avg": _count_metric(97.0), + "error_request_count_avg": _count_metric(0.0), + "context_overflow_count_avg": _count_metric(3.0), + }, + metadata={"scenario": "inferencex-agentx-mvp", "confidence_level": 0.95}, + ) + md = _export_metadata(tmp_path, result) + assert md["scenario"] == "inferencex-agentx-mvp" + assert md["submission_valid"] is False + assert "context_overflow_rate_exceeded" in md["submission_invalid_reasons"] + + +def test_carrier_keys_sum_across_runs_and_are_stripped(tmp_path): + """Orchestrator-stamped cross-run carrier totals win and never leak. + + The summed ``_total_responses`` / ``_context_overflow_count`` carriers + (200 / 6 = 3% > 1%) drive the verdict; all underscore-prefixed carrier + keys are popped from the public metadata. + """ + outcome = ScenarioOutcome(scenario_name="inferencex-agentx-mvp") + result = _make_result( + metrics={}, + metadata={ + "_scenario_name": outcome.scenario_name, + "_validator_submission_valid": True, + "_validator_submission_invalid_reasons": [], + "_total_responses": 200, + "_context_overflow_count": 6, + "_was_cancelled": False, + }, + ) + md = _export_metadata(tmp_path, result) + assert md["scenario"] == "inferencex-agentx-mvp" + assert md["submission_valid"] is False + assert md["submission_invalid_reasons"] == ["context_overflow_rate_exceeded"] + for key in ( + "_scenario_name", + "_validator_submission_valid", + "_validator_submission_invalid_reasons", + "_total_responses", + "_context_overflow_count", + "_was_cancelled", + ): + assert key not in md + + +def test_lock_violation_propagates_into_aggregate(tmp_path): + """A real --unsafe-override ScenarioOutcome lock violation propagates. + + No overflow, no cancellation: the aggregate verdict is exactly the static + lock outcome (False + unsafe_override) carried through the validator + carrier keys. + """ + outcome = ScenarioOutcome( + scenario_name="inferencex-agentx-mvp", + submission_valid=False, + submission_invalid_reasons=["unsafe_override"], + ) + result = _make_result( + metrics={ + "request_count_avg": _count_metric(100.0), + "context_overflow_count_avg": _count_metric(0.0), + }, + metadata={ + "_scenario_name": outcome.scenario_name, + "_validator_submission_valid": outcome.submission_valid, + "_validator_submission_invalid_reasons": outcome.submission_invalid_reasons, + "_total_responses": 100, + "_context_overflow_count": 0, + }, + ) + md = _export_metadata(tmp_path, result) + assert md["submission_valid"] is False + assert md["submission_invalid_reasons"] == ["unsafe_override"] + + +def test_cancellation_flips_aggregate_invalid(tmp_path): + """A cancelled multi-run aggregate is never a valid submission.""" + result = _make_result( + metrics={ + "request_count_avg": _count_metric(50.0), + "context_overflow_count_avg": _count_metric(0.0), + }, + metadata={ + "_scenario_name": "inferencex-agentx-mvp", + "_validator_submission_valid": True, + "_total_responses": 50, + "_context_overflow_count": 0, + "_was_cancelled": True, + }, + ) + md = _export_metadata(tmp_path, result) + assert md["submission_valid"] is False + assert md["submission_invalid_reasons"] == ["run_cancelled"] + + +def test_clean_scenario_run_is_valid(tmp_path): + """A clean scenario run (under threshold, no lock, no cancel) is valid.""" + result = _make_result( + metrics={ + "request_count_avg": _count_metric(1000.0), + "context_overflow_count_avg": _count_metric(1.0), # 0.1% < 1% + }, + metadata={"scenario": "inferencex-agentx-mvp"}, + ) + md = _export_metadata(tmp_path, result) + assert md["scenario"] == "inferencex-agentx-mvp" + assert md["submission_valid"] is True + assert "submission_invalid_reasons" not in md + + +def test_no_scenario_run_omits_submission_fields(tmp_path): + """A non-scenario aggregate carries no submission fields (null-safe omit).""" + result = _make_result( + metrics={ + "request_count_avg": _count_metric(10.0), + "context_overflow_count_avg": _count_metric(9.0), # huge rate, ignored + }, + metadata={"confidence_level": 0.95}, + ) + md = _export_metadata(tmp_path, result) + assert "scenario" not in md + assert "submission_valid" not in md + assert "submission_invalid_reasons" not in md + # Untouched aggregate metadata still flows through. + assert md["aggregation_type"] == "confidence" + assert md["num_profile_runs"] == 2 + + +@pytest.mark.parametrize( + "overflow_mean,total_request_mean,expect_valid", + [ + (1.0, 99.0, True), # 1/100 == 1% boundary, accepted (strictly >) + (2.0, 98.0, False), # 2/100 == 2% > 1% + ], +) # fmt: skip +def test_boundary_rate_from_means( + tmp_path, overflow_mean, total_request_mean, expect_valid +): + """The strictly-greater rate operator holds when derived from means.""" + result = _make_result( + metrics={ + "request_count_avg": _count_metric(total_request_mean), + "context_overflow_count_avg": _count_metric(overflow_mean), + }, + metadata={"scenario": "inferencex-agentx-mvp"}, + ) + md = _export_metadata(tmp_path, result) + assert md["submission_valid"] is expect_valid diff --git a/tests/unit/exporters/aggregate/test_sweep_exporters.py b/tests/unit/exporters/aggregate/test_sweep_exporters.py index 5373d1069d..73f59830ca 100644 --- a/tests/unit/exporters/aggregate/test_sweep_exporters.py +++ b/tests/unit/exporters/aggregate/test_sweep_exporters.py @@ -168,7 +168,7 @@ async def test_export_json_schema_compliant_valid_data( assert data["num_profile_runs"] == 15 assert data["num_successful_runs"] == 15 - # Verify metadata section (Task 9.4) + # Verify metadata section assert "metadata" in data metadata = data["metadata"] assert "sweep_parameters" in metadata @@ -177,7 +177,7 @@ async def test_export_json_schema_compliant_valid_data( assert metadata["sweep_parameters"][0]["values"] == [10, 20, 30] assert metadata["num_combinations"] == 3 - # Verify per-combination metrics section (Task 9.5) + # Verify per-combination metrics section assert "per_combination_metrics" in data per_combination = data["per_combination_metrics"] assert len(per_combination) == 3 @@ -197,7 +197,7 @@ async def test_export_json_schema_compliant_valid_data( assert throughput["std"] == 5.2 assert throughput["unit"] == "requests/sec" - # Verify best configurations section (Task 9.6) + # Verify best configurations section assert "best_configurations" in data best = data["best_configurations"] assert "best_throughput" in best @@ -206,7 +206,7 @@ async def test_export_json_schema_compliant_valid_data( assert "best_latency_p99" in best assert best["best_latency_p99"]["parameters"] == {"concurrency": 10} - # Verify Pareto optimal section (Task 9.7) + # Verify Pareto optimal section assert "pareto_optimal" in data pareto = data["pareto_optimal"] assert len(pareto) == 2 diff --git a/tests/unit/exporters/test_console_exporter.py b/tests/unit/exporters/test_console_exporter.py index 9f333473fb..0a7bca015a 100644 --- a/tests/unit/exporters/test_console_exporter.py +++ b/tests/unit/exporters/test_console_exporter.py @@ -25,6 +25,9 @@ from aiperf.metrics.types.inter_token_latency_metric import InterTokenLatencyMetric from aiperf.metrics.types.output_token_count import OutputTokenCountMetric from aiperf.metrics.types.request_latency_metric import RequestLatencyMetric +from aiperf.metrics.types.theoretical_prefix_cache_metric import ( + TheoreticalPrefixCacheHitMetric, +) from aiperf.metrics.types.ttft_metric import TTFTMetric from aiperf.plugin.enums import EndpointType from tests.unit.exporters.conftest import make_exporter_config @@ -185,6 +188,44 @@ async def test_export_omits_cache_hint_when_cache_reported( await ConsoleMetricsExporter(config).export(Console(width=115)) assert "--enable-prompt-tokens-details" not in capsys.readouterr().out + @pytest.mark.asyncio + async def test_export_renders_theoretical_prefix_cache_hit_in_cache_table( + self, mock_endpoint_config, capsys + ): + """The accumulator-injected tag must not be silently omitted. + + It is registered with console_group=CACHE, so it renders in its own + Cache table rather than the default metrics table. + """ + config = make_exporter_config( + results=ProfileResults( + records=[ + MetricResult( + tag="request_latency", + header="Request Latency", + unit="ms", + avg=1.0, + ), + MetricResult( + tag=TheoreticalPrefixCacheHitMetric.tag, + header=TheoreticalPrefixCacheHitMetric.header, + unit=str(TheoreticalPrefixCacheHitMetric.unit), + avg=42.5, + ), + ], + start_ns=0, + end_ns=0, + completed=0, + ), + cli_config=mock_endpoint_config, + telemetry_results=None, + ) + await ConsoleMetricsExporter(config).export(Console(width=115)) + output = capsys.readouterr().out + assert "NVIDIA AIPerf | LLM Metrics: Cache" in output + assert "Theoretical Prefix Cache Hit" in output + assert "42.50" in output + @pytest.mark.parametrize( "metric_tag, should_show", [ @@ -200,6 +241,8 @@ async def test_export_omits_cache_hint_when_cache_reported( (InterTokenLatencyMetric.tag, True), # Normal metric (RequestLatencyMetric.tag, True), # Normal metric (TTFTMetric.tag, True), # Normal metric + # Externally-injected accumulator metric - shown (CACHE group) + (TheoreticalPrefixCacheHitMetric.tag, True), ], ) # fmt: skip def test_should_show_metrics_based_on_flags( diff --git a/tests/unit/exporters/test_metrics_json_exporter.py b/tests/unit/exporters/test_metrics_json_exporter.py index 05d6cb0f70..7c46632a3e 100644 --- a/tests/unit/exporters/test_metrics_json_exporter.py +++ b/tests/unit/exporters/test_metrics_json_exporter.py @@ -208,6 +208,74 @@ def error_summary(self): assert "count" not in raw["request_count"] assert raw["request_count"]["avg"] == 20.0 + @pytest.mark.asyncio + async def test_run_info_carries_scenario_outcome(self, mock_results, mock_cfg): + """run_info surfaces scenario_name + submission_valid from the resolved + ScenarioOutcome (real BenchmarkRun, real ScenarioOutcome -- not MagicMock).""" + from aiperf.common.scenario import ScenarioOutcome + from aiperf.config.resolution.plan import BenchmarkRun + from tests.unit.conftest import make_cfg_from_v1 + + with tempfile.TemporaryDirectory() as temp_dir: + output_dir = Path(temp_dir) + mock_cfg.artifact_directory = output_dir + + cfg = make_cfg_from_v1(mock_cfg, artifact_directory=output_dir) + run = BenchmarkRun( + benchmark_id="scn123", + cfg=cfg, + artifact_dir=output_dir, + ) + run.resolved.scenario_outcome = ScenarioOutcome( + scenario_name="inferencex-agentx-mvp", + submission_valid=False, + submission_invalid_reasons=["unsafe_override"], + ) + exporter_config = make_exporter_config( + results=mock_results, + cli_config=mock_cfg, + telemetry_results=None, + run=run, + ) + exporter = MetricsJsonExporter(exporter_config) + await exporter.export() + + expected_file = output_dir / OutputDefaults.PROFILE_EXPORT_AIPERF_JSON_FILE + data = JsonExportData.model_validate_json(expected_file.read_text()) + assert data.run_info is not None + assert data.run_info.scenario_name == "inferencex-agentx-mvp" + assert data.run_info.submission_valid is False + assert data.run_info.submission_invalid_reasons == ["unsafe_override"] + + @pytest.mark.asyncio + async def test_run_info_omits_scenario_fields_when_no_scenario( + self, mock_results, mock_cfg + ): + """When no scenario is set, the scenario fields are absent from the JSON + (None + exclude_none drops them).""" + from aiperf.config.resolution.plan import BenchmarkRun + from tests.unit.conftest import make_cfg_from_v1 + + with tempfile.TemporaryDirectory() as temp_dir: + output_dir = Path(temp_dir) + mock_cfg.artifact_directory = output_dir + + cfg = make_cfg_from_v1(mock_cfg, artifact_directory=output_dir) + run = BenchmarkRun(benchmark_id="noscn", cfg=cfg, artifact_dir=output_dir) + exporter_config = make_exporter_config( + results=mock_results, + cli_config=mock_cfg, + telemetry_results=None, + run=run, + ) + exporter = MetricsJsonExporter(exporter_config) + await exporter.export() + + expected_file = output_dir / OutputDefaults.PROFILE_EXPORT_AIPERF_JSON_FILE + raw = json.loads(expected_file.read_text()) + assert "scenario_name" not in raw.get("run_info", {}) + assert "submission_valid" not in raw.get("run_info", {}) + @pytest.mark.asyncio async def test_run_info_populated_when_run_provided(self, mock_results, mock_cfg): """run_info surfaces seed + variation coordinates when ExporterConfig.run is set.""" @@ -934,6 +1002,170 @@ async def test_json_export_with_hostname_metadata(self, mock_results, mock_cfg): assert gpu_summary["hostname"] == "test-hostname" +# --------------------------------------------------------------------------- +# Context-overflow submission-validity fold (InferenceX AgentX RFC §7) +# --------------------------------------------------------------------------- + + +def _make_scenario_run(scenario_outcome): + """Build a real BenchmarkRun carrying a real ScenarioOutcome (no MagicMock).""" + from aiperf.config import BenchmarkConfig + from aiperf.config.resolution.plan import BenchmarkRun + + cfg = BenchmarkConfig( + models=["test-model"], + endpoint={"urls": ["http://localhost:8000/v1/chat/completions"]}, + datasets=[{"name": "profiling", "type": "synthetic"}], + phases=[ + { + "name": "profiling", + "type": "concurrency", + "concurrency": 1, + "sessions": 5, + } + ], + ) + run = BenchmarkRun( + benchmark_id="ctx-overflow-run", + cfg=cfg, + artifact_dir=Path("/tmp/ctx-overflow"), + cli_command=None, + ) + run.resolved.scenario_outcome = scenario_outcome + return run + + +def _count_record(tag: str, value: float): + return MetricResult( + tag=tag, + header=tag, + unit="requests", + avg=value, + ) + + +class _OverflowResults: + def __init__(self, records, was_cancelled=False): + self.metrics = records + self.start_ns = None + self.end_ns = None + self._was_cancelled = was_cancelled + + @property + def records(self): + return self.metrics + + @property + def has_results(self): + return bool(self.metrics) + + @property + def was_cancelled(self): + return self._was_cancelled + + @property + def error_summary(self): + return [] + + +async def _export_run_info(run, results, mock_cfg): + with tempfile.TemporaryDirectory() as temp_dir: + output_dir = Path(temp_dir) + mock_cfg.artifact_directory = output_dir + exporter_config = make_exporter_config( + results=results, cli_config=mock_cfg, run=run + ) + exporter = MetricsJsonExporter(exporter_config) + await exporter.export() + expected_file = output_dir / OutputDefaults.PROFILE_EXPORT_AIPERF_JSON_FILE + with open(expected_file) as f: + data = JsonExportData.model_validate_json(f.read()) + return data.run_info + + +class TestContextOverflowSubmissionFold: + @pytest.mark.asyncio + async def test_overflow_rate_over_threshold_flips_submission_invalid( + self, mock_cfg + ): + """> threshold overflow rate flips run_info.submission_valid to False.""" + from aiperf.common.scenario import ScenarioOutcome + + outcome = ScenarioOutcome(scenario_name="inferencex", submission_valid=True) + run = _make_scenario_run(outcome) + # 5 overflow of 100 total -> 5% > 1% default. + records = [ + _count_record("request_count", 95), + _count_record("error_request_count", 5), + _count_record("context_overflow_count", 5), + ] + results = _OverflowResults(records) + run_info = await _export_run_info(run, results, mock_cfg) + assert run_info is not None + assert run_info.submission_valid is False + assert "context_overflow_rate_exceeded" in run_info.submission_invalid_reasons + + @pytest.mark.asyncio + async def test_overflow_rate_under_threshold_unaffected(self, mock_cfg): + """< threshold overflow rate leaves submission_valid True (lock-only).""" + from aiperf.common.scenario import ScenarioOutcome + + outcome = ScenarioOutcome(scenario_name="inferencex", submission_valid=True) + run = _make_scenario_run(outcome) + # 0 overflow of ~1000 total. + records = [ + _count_record("request_count", 1000), + _count_record("error_request_count", 0), + _count_record("context_overflow_count", 0), + ] + results = _OverflowResults(records) + run_info = await _export_run_info(run, results, mock_cfg) + assert run_info is not None + assert run_info.submission_valid is True + assert run_info.submission_invalid_reasons is None + + @pytest.mark.asyncio + async def test_lock_violation_and_overflow_both_surface(self, mock_cfg): + """unsafe_override lock reason AND overflow reason both appear.""" + from aiperf.common.scenario import ScenarioOutcome + + outcome = ScenarioOutcome( + scenario_name="inferencex", + submission_valid=False, + submission_invalid_reasons=["unsafe_override"], + ) + run = _make_scenario_run(outcome) + records = [ + _count_record("request_count", 90), + _count_record("error_request_count", 10), + _count_record("context_overflow_count", 10), + ] + results = _OverflowResults(records) + run_info = await _export_run_info(run, results, mock_cfg) + assert run_info is not None + assert run_info.submission_valid is False + assert "unsafe_override" in run_info.submission_invalid_reasons + assert "context_overflow_rate_exceeded" in run_info.submission_invalid_reasons + + @pytest.mark.asyncio + async def test_no_scenario_run_info_submission_valid_none(self, mock_cfg): + """A run with no scenario keeps submission_valid None even with overflow.""" + from aiperf.common.scenario import ScenarioOutcome + + outcome = ScenarioOutcome() # scenario_name None + run = _make_scenario_run(outcome) + records = [ + _count_record("request_count", 50), + _count_record("error_request_count", 50), + _count_record("context_overflow_count", 50), + ] + results = _OverflowResults(records) + run_info = await _export_run_info(run, results, mock_cfg) + assert run_info is not None + assert run_info.scenario_name is None + assert run_info.submission_valid is None + + class TestMetricsJsonExporterWarmupMetrics: @pytest.mark.asyncio async def test_json_export_includes_warmup_metrics_separately(self, mock_cfg): diff --git a/tests/unit/graph/fixtures/weka_first_token_anchor.json b/tests/unit/graph/fixtures/weka_first_token_anchor.json new file mode 100644 index 0000000000..9cc1fa0529 --- /dev/null +++ b/tests/unit/graph/fixtures/weka_first_token_anchor.json @@ -0,0 +1,16 @@ +{ + "id": "trace_first_token_anchor", + "models": ["claude-opus-4-5-20251101"], + "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "s", "model": "claude-opus-4-5-20251101", "in": 128, "out": 200, "hash_ids": [1, 2], "input_types": ["text"], "output_types": ["text"], "stop": "tool_use", "api_time": 8.0, "think_time": 0.0, "ttft": 2.0}, + {"t": 0.5, "type": "subagent", "agent_id": "agent_pre", "subagent_type": "Explore", "duration_ms": 1500, "total_tokens": 200, "tool_use_count": 1, "status": "async_launched", "models": ["claude-opus-4-5-20251101"], "tool_tokens": 0, "system_tokens": 0, "requests": [ + {"t": 1.0, "type": "n", "model": "claude-opus-4-5-20251101", "in": 128, "out": 45, "hash_ids": [60, 61], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 1.0, "think_time": 0.0} + ]}, + {"t": 3.5, "type": "subagent", "agent_id": "agent_post", "subagent_type": "Explore", "duration_ms": 1500, "total_tokens": 200, "tool_use_count": 1, "status": "async_launched", "models": ["claude-opus-4-5-20251101"], "tool_tokens": 0, "system_tokens": 0, "requests": [ + {"t": 4.0, "type": "n", "model": "claude-opus-4-5-20251101", "in": 128, "out": 30, "hash_ids": [50, 51], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 1.0, "think_time": 0.0} + ]}, + {"t": 9.5, "type": "n", "model": "claude-opus-4-5-20251101", "in": 192, "out": 60, "hash_ids": [1, 2, 3], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 0.5, "think_time": 0.0} + ] +} diff --git a/tests/unit/graph/fixtures/weka_min.json b/tests/unit/graph/fixtures/weka_min.json new file mode 100644 index 0000000000..5546c10fd5 --- /dev/null +++ b/tests/unit/graph/fixtures/weka_min.json @@ -0,0 +1,11 @@ +{ + "id": "trace_03_n3", + "models": ["claude-opus-4-5-20251101"], + "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "n", "model": "claude-opus-4-5-20251101", "in": 180, "out": 25, "hash_ids": [1, 2], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 0.8, "think_time": 0.0}, + {"t": 1.5, "type": "n", "model": "claude-opus-4-5-20251101", "in": 230, "out": 35, "hash_ids": [1, 2, 3], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 1.0, "think_time": 0.5}, + {"t": 3.0, "type": "n", "model": "claude-opus-4-5-20251101", "in": 280, "out": 45, "hash_ids": [1, 2, 3, 4], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 1.3, "think_time": 0.5} + ] +} diff --git a/tests/unit/graph/fixtures/weka_start_anchor.json b/tests/unit/graph/fixtures/weka_start_anchor.json new file mode 100644 index 0000000000..8535ef3da8 --- /dev/null +++ b/tests/unit/graph/fixtures/weka_start_anchor.json @@ -0,0 +1,16 @@ +{ + "id": "trace_start_anchor", + "models": ["claude-opus-4-5-20251101"], + "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "n", "model": "claude-opus-4-5-20251101", "in": 128, "out": 200, "hash_ids": [1, 2], "input_types": ["text"], "output_types": ["text"], "stop": "tool_use", "api_time": 8.0, "think_time": 0.0}, + {"t": 2.0, "type": "subagent", "agent_id": "agent_spawn", "subagent_type": "Explore", "duration_ms": 1500, "total_tokens": 300, "tool_use_count": 1, "status": "async_launched", "models": ["claude-opus-4-5-20251101"], "tool_tokens": 0, "system_tokens": 0, "requests": [ + {"t": 2.5, "type": "n", "model": "claude-opus-4-5-20251101", "in": 128, "out": 30, "hash_ids": [50, 51], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 1.0, "think_time": 0.0} + ]}, + {"t": 4.5, "type": "subagent", "agent_id": "agent_chain", "subagent_type": "Task", "duration_ms": 1500, "total_tokens": 300, "tool_use_count": 1, "status": "async_launched", "models": ["claude-opus-4-5-20251101"], "tool_tokens": 0, "system_tokens": 0, "requests": [ + {"t": 5.0, "type": "n", "model": "claude-opus-4-5-20251101", "in": 192, "out": 45, "hash_ids": [60, 61, 62], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 1.0, "think_time": 0.0} + ]}, + {"t": 9.0, "type": "n", "model": "claude-opus-4-5-20251101", "in": 256, "out": 60, "hash_ids": [1, 2, 3, 4], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 0.5, "think_time": 0.0} + ] +} diff --git a/tests/unit/graph/fixtures/weka_subagent.json b/tests/unit/graph/fixtures/weka_subagent.json new file mode 100644 index 0000000000..175f64e87a --- /dev/null +++ b/tests/unit/graph/fixtures/weka_subagent.json @@ -0,0 +1,27 @@ +{ + "id": "trace_sub_n2s1", + "models": ["claude-opus-4-5-20251101"], + "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "n", "model": "claude-opus-4-5-20251101", "in": 180, "out": 25, "hash_ids": [1, 2], "input_types": ["text"], "output_types": ["text"], "stop": "tool_use", "api_time": 0.8, "think_time": 0.0}, + { + "t": 1.0, + "type": "subagent", + "agent_id": "agent_001", + "subagent_type": "Explore", + "duration_ms": 4000, + "total_tokens": 600, + "tool_use_count": 1, + "status": "completed", + "models": ["claude-opus-4-5-20251101"], + "tool_tokens": 0, + "system_tokens": 0, + "requests": [ + {"t": 1.2, "type": "n", "model": "claude-opus-4-5-20251101", "in": 200, "out": 30, "hash_ids": [10, 11], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 0.9, "think_time": 0.0}, + {"t": 2.4, "type": "n", "model": "claude-opus-4-5-20251101", "in": 260, "out": 40, "hash_ids": [10, 11, 12], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 1.1, "think_time": 0.2} + ] + }, + {"t": 6.0, "type": "n", "model": "claude-opus-4-5-20251101", "in": 280, "out": 45, "hash_ids": [1, 2, 3, 4], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 1.3, "think_time": 0.5} + ] +} diff --git a/tests/unit/graph/test_accelerated_warmup.py b/tests/unit/graph/test_accelerated_warmup.py new file mode 100644 index 0000000000..136a985c7e --- /dev/null +++ b/tests/unit/graph/test_accelerated_warmup.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Accelerated cache-pressure warmup -- knob-gated pacing. + +These tests pin the zero-idle warmup replay (live trajectories replay with +ZERO idle delay to drive the server KV cache to pressure): +the WARMUP-phase ``TraceExecutor`` is built with ``compress_edge_delays=True`` +ONLY when a cache-pressure duration is configured, collapsing +every captured inter-node edge delay. Default (knob OFF) honors every edge delay +exactly. The 1-token warmup output +cap (``WARMUP_MAX_OUTPUT_TOKENS``) is independent of pacing and untouched. +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path +from typing import Any + +import pytest + +from aiperf.common.enums import CreditPhase +from aiperf.common.environment import Environment +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.graph.executor import TraceExecutor + +FIX = Path(__file__).parent / "fixtures" / "weka_min.json" + + +class _RecordingCreditIssuer: + """Stub issuer recording dispatch order; returns a placeholder string.""" + + def __init__(self) -> None: + self.dispatched: list[str] = [] + + async def dispatch( + self, + node: Any, + request: Any, + ctx: Any, + **kwargs: Any, + ) -> str: + self.dispatched.append(request.node_id) + return f"placeholder::{request.node_id}" + + +def _count_real_sleeps(monkeypatch) -> list[float]: + """Patch ``asyncio.sleep`` in the executor module to record positive waits. + + ``_apply_firing_delay`` calls ``await asyncio.sleep(wait_us / 1e6)`` only + when a non-zero gate remains after honoring edge delays. Recording the + POSITIVE-duration sleeps it issues is a direct, virtual-time-independent + probe of whether the executor honored (>=1 positive sleep on a delayed + chain) or collapsed (zero positive sleeps) the edge delays. + """ + sleeps: list[float] = [] + import aiperf.graph.executor as executor_mod + + real_sleep = asyncio.sleep + + async def _spy(delay: float, *args, **kwargs): + if delay > 0: + sleeps.append(delay) + return await real_sleep(0) + + monkeypatch.setattr(executor_mod.asyncio, "sleep", _spy) + return sleeps + + +@pytest.mark.asyncio +async def test_executor_honors_edge_delays_by_default(monkeypatch): + """Default executor (compress_edge_delays=False) honors the F1 edge delays. + + The trie IR always stamps end-to-start inter-turn delays, so ``weka_min`` + carries them (700 ms, 500 ms) on its main chain and the default executor must + issue at least one positive sleep in ``_apply_firing_delay``. + """ + monkeypatch.setattr(Environment.GRAPH, "IGNORE_EDGE_DELAYS", False) + sleeps = _count_real_sleeps(monkeypatch) + + parsed = from_weka_trace(str(FIX)) + issuer = _RecordingCreditIssuer() + executor = TraceExecutor(parsed, credit_issuer=issuer) + + async with asyncio.TaskGroup(): + for trace in parsed.traces: + await executor.run(trace) + + assert sleeps, ( + "default executor must honor F1 end-to-start edge delays (expected at " + "least one positive firing-delay sleep on weka_min's 700/500 ms chain)" + ) + + +@pytest.mark.asyncio +async def test_executor_collapses_edge_delays_when_compressed(monkeypatch): + """compress_edge_delays=True collapses ALL firing delays (zero-idle / burst). + + The live trajectories replay with zero idle delay. No positive + firing-delay sleep + may be issued, yet every node still dispatches exactly once. + """ + monkeypatch.setattr(Environment.GRAPH, "IGNORE_EDGE_DELAYS", False) + sleeps = _count_real_sleeps(monkeypatch) + + parsed = from_weka_trace(str(FIX)) + from aiperf.dataset.graph.models import LlmNode + + graph = parsed.graph if not parsed.graphs else next(iter(parsed.graphs.values())) + expected = {nid for nid, n in graph.nodes.items() if isinstance(n, LlmNode)} + + issuer = _RecordingCreditIssuer() + executor = TraceExecutor(parsed, credit_issuer=issuer, compress_edge_delays=True) + + async with asyncio.TaskGroup(): + for trace in parsed.traces: + await executor.run(trace) + + assert sleeps == [], ( + f"compressed executor must collapse all firing delays (zero idle), got " + f"positive sleeps {sleeps}" + ) + assert set(issuer.dispatched) == expected, ( + "compressed warmup must still dispatch every node exactly once" + ) + + +class _StubIssuer: + def mark_graph_sending_complete(self): + pass + + def graph_all_returned(self): + return True + + def set_graph_all_returned_event(self): + pass + + +def _strategy( + parsed, *, phase: CreditPhase | None, cache_pressure_duration_s: float | None = None +): + from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy + + class _Cfg: + pass + + cfg = _Cfg() + cfg.phase = phase + return GraphIRReplayStrategy( + config=cfg, + credit_issuer=_StubIssuer(), + parsed_graph=parsed, + register_observer=lambda _obs: None, + start_min_ratio=0.25, + start_max_ratio=0.75, + t_star_random_seed=1234, + cache_pressure_duration_s=cache_pressure_duration_s, + ) + + +@pytest.mark.parametrize( + "phase,duration,expected", + [ + # A WARMUP phase with no pressure duration honors recorded pacing. + (CreditPhase.WARMUP, None, False), + (CreditPhase.PROFILING, None, False), + (None, None, False), + # The configured duration (--agentic-cache-warmup-duration) alone + # activates compressed pacing -- WARMUP phase only. + (CreditPhase.WARMUP, 30.0, True), + (CreditPhase.PROFILING, 30.0, False), + (None, 30.0, False), + ], +) +def test_strategy_accelerated_warmup_gate(phase, duration, expected): + """Accelerated warmup is active iff WARMUP phase + a pressure duration.""" + parsed = from_weka_trace(str(FIX)) + strategy = _strategy(parsed, phase=phase, cache_pressure_duration_s=duration) + assert strategy.accelerated_warmup is expected diff --git a/tests/unit/graph/test_analysis_core.py b/tests/unit/graph/test_analysis_core.py new file mode 100644 index 0000000000..f53cf3e921 --- /dev/null +++ b/tests/unit/graph/test_analysis_core.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Direct unit cover for the graph analysis core. + +These functions (`elaborate_trace`, `compute_snapshot`) were previously +exercised only transitively through the strategy/adapter E2E path, so a +structural regression in any of them could slip past the suite. These tests +build a `ParsedGraph` and assert directly on each function's output -- the +parallel-readiness cohort keying and the snapshot t* partition. +""" + +from __future__ import annotations + +import pytest + +from aiperf.dataset.graph.models import ( + GraphRecord, + LlmNode, + ParsedGraph, + StaticEdge, + TraceRecord, +) +from aiperf.graph.analysis.snapshot import compute_snapshot +from aiperf.graph.analysis.timeline import GraphCycleError, elaborate_trace + + +def _llm(output: str, *, offset: int | None = None) -> LlmNode: + return LlmNode(prompt=[f"@{output}"], output=output, arrival_offset_us=offset) + + +def _parsed( + nodes: dict[str, object], + edges: list[StaticEdge], +) -> ParsedGraph: + graph = GraphRecord(nodes=nodes, edges=edges, state={}) + return ParsedGraph( + graph=graph, + traces=[TraceRecord(id="t")], + ) + + +def test_compute_snapshot_inclusive_tstar_boundary_and_warmup_split(): + """arrival < t* -> warmup; arrival >= t* -> profiled (t* itself is kept).""" + parsed = _parsed( + { + "A": _llm("a", offset=0), + "B": _llm("b", offset=100), + "C": _llm("c", offset=200), + }, + [ + StaticEdge(source="START", target="A"), + StaticEdge(source="A", target="B"), + StaticEdge(source="B", target="C"), + ], + ) + snap = compute_snapshot(parsed, parsed.traces[0], t_star_us=100) + + assert {sf.firing.node_id for sf in snap.warmup} == {"A"} + profiled = {sf.firing.node_id: sf for sf in snap.profiled} + assert set(profiled) == {"B", "C"} + # The firing exactly at t* is kept (inclusive) and rebased to dispatch 0. + assert profiled["B"].dispatch_offset_us == 0 + assert profiled["C"].dispatch_offset_us == 100 + + +def test_elaborate_trace_depth_cap_raises_graph_cycle_error(): + """A firing count over depth_cap raises GraphCycleError (validator guard).""" + parsed = _parsed( + {"A": _llm("a"), "B": _llm("b")}, + [ + StaticEdge(source="START", target="A"), + StaticEdge(source="A", target="B"), + ], + ) + with pytest.raises(GraphCycleError): + elaborate_trace(parsed, parsed.traces[0], depth_cap=1) + + +# --------------------------------------------------------------------------- +# R1 -- start-anchored subtrees must elaborate (analysis/runtime agreement) +# --------------------------------------------------------------------------- + + +def _start_anchored_parsed() -> ParsedGraph: + """The R1 review graph: START->a; a->b (completion); a->c (start-anchored); + c->d (completion). The runtime fires all four; the dry-run must too.""" + return _parsed( + { + "a": _llm("a", offset=0), + "b": _llm("b", offset=1_000_000), + "c": _llm("c", offset=2_500_000), + "d": _llm("d", offset=9_000_000), + }, + [ + StaticEdge(source="START", target="a"), + StaticEdge(source="a", target="b", delay_after_predecessor_us=0.0), + StaticEdge( + source="a", target="c", delay_after_predecessor_start_us=2_500_000.0 + ), + StaticEdge(source="c", target="d", delay_after_predecessor_us=0.0), + ], + ) + + +def test_elaborate_trace_follows_start_anchored_edges(): + """The whole subtree under a start-anchored edge fires in the dry-run. + + Before the fix the elaborator consulted only ``successors_after`` and the + timeline stopped at ['a', 'b'] -- the anchored child ``c`` and its + completion successor ``d`` never fired, so ``duration_us`` under-measured + and snapshot planning misclassified the subtree. + """ + parsed = _start_anchored_parsed() + + timeline = elaborate_trace(parsed, parsed.traces[0]) + + assert [f.node_id for f in timeline.firings] == ["a", "b", "c", "d"] + assert timeline.duration_us() == 9_000_000 + + +def test_compute_snapshot_partitions_start_anchored_subtree(): + """Snapshot-at-t* sees anchored-subtree firings as warmup/profiled members.""" + parsed = _start_anchored_parsed() + + snap = compute_snapshot(parsed, parsed.traces[0], t_star_us=2_500_000) + + assert {sf.firing.node_id for sf in snap.warmup} == {"a", "b"} + profiled = {sf.firing.node_id: sf for sf in snap.profiled} + assert set(profiled) == {"c", "d"} + assert profiled["c"].dispatch_offset_us == 0 + assert profiled["d"].dispatch_offset_us == 6_500_000 + + +def test_elaborate_trace_pending_fan_in_waits_for_late_arrival(): + """A scheduled node whose fan-in is unsatisfied stays pending (not dropped). + + Mirrors the runtime's parked ``await_inputs``: ``join`` is scheduled via a + start-anchored edge from ``a`` but requires ``b``'s output channel, so it + fires only after ``b`` does. + """ + from aiperf.dataset.graph.models import ChannelRequirement, ChannelSpec + + join = LlmNode( + prompt=["@j"], + output="j", + inputs=[ChannelRequirement(channel="b", count=1)], + ) + graph = GraphRecord( + nodes={ + "a": _llm("a", offset=0), + "b": _llm("b", offset=100), + "join": join, + }, + edges=[ + StaticEdge(source="START", target="a"), + StaticEdge(source="a", target="b", delay_after_predecessor_us=0.0), + StaticEdge( + source="a", target="join", delay_after_predecessor_start_us=50.0 + ), + ], + state={ + "a": ChannelSpec(), + "b": ChannelSpec(), + "j": ChannelSpec(), + }, + ) + parsed = ParsedGraph(graph=graph, traces=[TraceRecord(id="t")]) + + timeline = elaborate_trace(parsed, parsed.traces[0]) + + order = [f.node_id for f in timeline.firings] + assert set(order) == {"a", "b", "join"} + assert order.index("join") > order.index("b") diff --git a/tests/unit/graph/test_build_strategy_graph_ir.py b/tests/unit/graph/test_build_strategy_graph_ir.py new file mode 100644 index 0000000000..8db369aa2d --- /dev/null +++ b/tests/unit/graph/test_build_strategy_graph_ir.py @@ -0,0 +1,164 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""R3 — PhaseRunner._build_strategy selects + injects GraphIRReplayStrategy. + +Verifies the plugin lookup resolves ``TimingMode.GRAPH_IR`` to +``GraphIRReplayStrategy`` and that the graph-only injection branch supplies the +``ParsedGraph`` (from the conversation source) and the graph-return observer +(from the callback handler) WITHOUT disturbing the other strategies' fixed-kwarg +construction. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from aiperf.plugin import plugins +from aiperf.plugin.enums import PluginType, TimingMode +from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy + +_FIX = Path(__file__).parent / "fixtures" / "weka_min.json" + + +def test_plugin_registry_resolves_graph_ir_to_strategy(): + cls = plugins.get_class(PluginType.TIMING_STRATEGY, TimingMode.GRAPH_IR) + assert cls is GraphIRReplayStrategy + + +def test_build_strategy_injects_parsed_graph_and_observer(): + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + + parsed = from_weka_trace(str(_FIX)) + + # Minimal stand-ins for the runner collaborators the graph branch reads. + class _Source: + parsed_graph = parsed + + class _Handler: + def __init__(self) -> None: + self.installed = None + + def set_graph_return_observer(self, obs) -> None: + self.installed = obs + + class _Config: + timing_mode = TimingMode.GRAPH_IR + phase = None + concurrency = None + + # Build via the same code path PhaseRunner._build_strategy uses, without + # standing up a full PhaseRunner (its __init__ needs the whole pipeline). + handler = _Handler() + StrategyClass = plugins.get_class(PluginType.TIMING_STRATEGY, TimingMode.GRAPH_IR) + strategy = StrategyClass( + config=_Config(), + conversation_source=_Source(), + scheduler=None, + stop_checker=None, + credit_issuer=object(), + lifecycle=None, + parsed_graph=_Source.parsed_graph, + register_observer=handler.set_graph_return_observer, + ) + assert isinstance(strategy, GraphIRReplayStrategy) + + +@pytest.mark.asyncio +async def test_setup_phase_installs_single_observer(): + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + + parsed = from_weka_trace(str(_FIX)) + installed = [] + strategy = GraphIRReplayStrategy( + credit_issuer=object(), + parsed_graph=parsed, + register_observer=installed.append, + ) + await strategy.setup_phase() + assert len(installed) == 1 + assert callable(installed[0]) + + +@pytest.fixture +def graph_runner(): + """Runner + graph channel wired the way ``PhaseRunner`` does. + + Mirrors ``test_tstar_activation``'s ``PhaseRunner.__new__`` harness so these + tests exercise the real ``_build_graph_ir_strategy`` seam (kwarg threading + and the consume-once handoff pop) without standing up the full pipeline. + Returns ``(runner, channel, captured, StrategyClass)`` where ``captured`` is + updated with the strategy kwargs on every build. + """ + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + from aiperf.timing.graph_channel import GraphPhaseChannel + from aiperf.timing.phase.runner import PhaseRunner + + parsed = from_weka_trace(str(_FIX)) + captured: dict = {} + + class _CapturingStrategy(GraphIRReplayStrategy): + def __init__(self, **kw): + captured.update(kw) + super().__init__(**kw) + + class _Handler: + def set_graph_return_observer(self, obs) -> None: + self._obs = obs + + def set_graph_first_token_observer(self, obs) -> None: + self._ft_obs = obs + + class _Config: + timing_mode = TimingMode.GRAPH_IR + phase = None + concurrency = None + expected_num_sessions = None + + channel = GraphPhaseChannel(parsed_graph=parsed) + runner = PhaseRunner.__new__(PhaseRunner) + runner._config = _Config() + runner._conversation_source = None + runner._graph_channel = channel + runner._scheduler = None + runner._stop_checker = None + runner._credit_issuer = object() + runner._lifecycle = None + runner._callback_handler = _Handler() + + return runner, channel, captured, _CapturingStrategy + + +def test_build_graph_ir_strategy_threads_and_consumes_warmup_handoff(graph_runner): + """The runner threads the stashed handoff into the strategy and clears it. + + Consume-once: a second phase built from the same graph channel must + NOT see a stale handoff (multi-phase profiling configs). + """ + from aiperf.timing.graph_warmup_handoff import GraphWarmupHandoff + + runner, channel, captured, StrategyClass = graph_runner + + handoff = GraphWarmupHandoff( + lanes={}, drain_end_wall_us=0.0, corpus_cursor=0, pressure_lane_count=0 + ) + channel.warmup_handoff = handoff + + runner._build_graph_ir_strategy(StrategyClass) + + assert captured["warmup_handoff"] is handoff + assert channel.warmup_handoff is None + + runner._build_graph_ir_strategy(StrategyClass) + assert captured["warmup_handoff"] is None + + +def test_build_graph_ir_strategy_threads_pressure_duration(graph_runner): + """The runner sources the cache-pressure duration from the phase config.""" + runner, _channel, captured, StrategyClass = graph_runner + runner._config.cache_pressure_duration = 30.0 + + runner._build_graph_ir_strategy(StrategyClass) + + assert captured["cache_pressure_duration_s"] == 30.0 diff --git a/tests/unit/graph/test_cache_bust_marker.py b/tests/unit/graph/test_cache_bust_marker.py new file mode 100644 index 0000000000..b3a9ad7eab --- /dev/null +++ b/tests/unit/graph/test_cache_bust_marker.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""AgentX FIRST_TURN_PREFIX cache-bust marker on the graph-IR dispatch path. + +The marker is PER-TRACE-INSTANCE: minted from ``(benchmark_id, credit.trace_id)`` +and stamped onto the first user turn of the wire ``messages`` at worker +materialize time. Every dispatch of one trace instance shares the SAME marker +(its own prefix stays consistent); distinct instances get distinct markers +(cross-instance bust); a recycled template (a fresh instance id, e.g. ``t-1#1``) +mints a fresh marker. The default (``CacheBustTarget.NONE``) stamps nothing, so +the verbatim replay path is byte-identical to today unless the run passes +``--cache-bust``. Mirrors agentx's per-trajectory-tree scoping +(``cache_bust.py::resolve_tree_marker``, reset on recycle). +""" + +from __future__ import annotations + +import copy + +import pytest + +from aiperf.common.enums import CacheBustTarget +from aiperf.graph.worker_materialize import stamp_cache_bust_marker +from aiperf.timing.strategies.cache_bust import ( + build_trace_instance_marker, + inject_marker_at_first_user_message, +) + +_BENCH = "bench-42" +_FTP = CacheBustTarget.FIRST_TURN_PREFIX + + +def _payload() -> dict: + """A representative materialized graph payload (system + two user turns).""" + return { + "messages": [ + {"role": "system", "content": "you are a helpful assistant"}, + {"role": "user", "content": "first user turn"}, + {"role": "assistant", "content": "ok"}, + {"role": "user", "content": "second user turn"}, + ], + "model": "m", + "stream": True, + } + + +def _stamp(payload: dict, trace_instance_id: str, target: CacheBustTarget = _FTP): + stamp_cache_bust_marker( + payload, + benchmark_id=_BENCH, + trace_instance_id=trace_instance_id, + target=target, + ) + + +def _first_user_content(payload: dict) -> str: + return payload["messages"][1]["content"] + + +def test_marker_format_is_rid_prefix_with_blank_line(): + """The marker is ``[rid:<12hex>]\\n\\n`` -- agentx FIRST_TURN_PREFIX shape.""" + marker = build_trace_instance_marker(_BENCH, "t-1#0", target=_FTP) + assert marker is not None + assert marker.startswith("[rid:") + assert marker.endswith("]\n\n") + digest = marker[len("[rid:") : -len("]\n\n")] + assert len(digest) == 12 + assert all(c in "0123456789abcdef" for c in digest) + + +def test_none_target_mints_no_marker(): + """``NONE`` mints ``None`` so callers can pass it through unconditionally.""" + assert ( + build_trace_instance_marker(_BENCH, "t-1#0", target=CacheBustTarget.NONE) + is None + ) + + +def test_stamp_none_is_byte_identical_noop(): + """With cache-bust NONE the materialized payload is unchanged (today's behavior).""" + payload = _payload() + original = copy.deepcopy(payload) + _stamp(payload, "t-1#0", target=CacheBustTarget.NONE) + assert payload == original + + +def test_stamp_first_turn_prefix_marks_only_first_user_turn(): + """The marker prepends to the FIRST user turn only; all other turns untouched.""" + payload = _payload() + _stamp(payload, "t-1#0") + msgs = payload["messages"] + assert msgs[0] == {"role": "system", "content": "you are a helpful assistant"} + assert msgs[1]["role"] == "user" + assert msgs[1]["content"].startswith("[rid:") + assert msgs[1]["content"].endswith("first user turn") + assert msgs[2] == {"role": "assistant", "content": "ok"} + # SECOND user turn is NOT stamped. + assert msgs[3] == {"role": "user", "content": "second user turn"} + + +def test_marker_is_shared_across_all_dispatches_of_one_trace_instance(): + """Every dispatch of ONE trace instance carries the IDENTICAL marker. + + Two distinct dispatches (different turns/nodes of the same trace instance + ``t-1#0``) must produce the same first-user marker, so the instance's own + conversation prefix stays consistent and prefix-caches WITHIN the instance. + """ + p_turn0 = _payload() + p_turn1 = _payload() + _stamp(p_turn0, "t-1#0") + _stamp(p_turn1, "t-1#0") + m0 = _first_user_content(p_turn0)[: len("[rid:000000000000]\n\n")] + m1 = _first_user_content(p_turn1)[: len("[rid:000000000000]\n\n")] + assert m0 == m1 + + +def test_marker_differs_across_distinct_trace_instances(): + """Distinct trace instances mint distinct markers (cross-instance bust).""" + a = build_trace_instance_marker(_BENCH, "t-1#0", target=_FTP) + b = build_trace_instance_marker(_BENCH, "t-2#0", target=_FTP) + assert a != b + + +def test_marker_resets_on_recycle_of_same_template(): + """A recycled template (fresh instance id ``t-1#1``) mints a FRESH marker. + + Recycling reuses the dataset trace TEMPLATE ``t-1`` in a new session slot, + minting a new instance id (``#1`` vs ``#0``). The marker must reset so the + recycled instance does not warm the server's cache on the prior instance's + prefix -- mirroring agentx's per-recycle ``recycle_pass`` bump. + """ + instance0 = build_trace_instance_marker(_BENCH, "t-1#0", target=_FTP) + recycled = build_trace_instance_marker(_BENCH, "t-1#1", target=_FTP) + assert instance0 != recycled + + +def test_subagent_descendant_shares_root_instance_marker(): + """A nested/subagent dispatch keyed on the SAME root instance id reuses the marker. + + The adapter pins ``credit.trace_id`` to the root instance for nested/subagent + dispatches too (only the runtime ``parent_trace_id`` carries ``::sa:`` / + ``::loop#N`` suffixes), so a subagent turn keyed on ``t-1#0`` gets the same + marker as the main session's turns -- matching agentx, where the whole + trajectory TREE shares one value. + """ + main = build_trace_instance_marker(_BENCH, "t-1#0", target=_FTP) + subagent = build_trace_instance_marker(_BENCH, "t-1#0", target=_FTP) + assert main == subagent + + +def test_marker_is_deterministic_for_same_inputs(): + """Same (benchmark_id, trace_instance_id) -> same marker (reproducible reruns).""" + a = build_trace_instance_marker(_BENCH, "t-9#0", target=_FTP) + b = build_trace_instance_marker(_BENCH, "t-9#0", target=_FTP) + assert a == b + + +def test_marker_varies_per_benchmark_id(): + """Different run salts mint different markers for the same trace instance.""" + a = build_trace_instance_marker("bench-A", "t-1#0", target=_FTP) + b = build_trace_instance_marker("bench-B", "t-1#0", target=_FTP) + assert a != b + + +def test_stamp_is_idempotent(): + """Re-stamping the same instance's marker does not stack it (agentx idempotency).""" + payload = _payload() + _stamp(payload, "t-1#0") + once = copy.deepcopy(payload) + _stamp(payload, "t-1#0") + assert payload == once + + +def test_inject_into_multimodal_first_user_content(): + """Multimodal list content gets a leading text marker part (agentx parity).""" + messages = [ + {"role": "user", "content": [{"type": "text", "text": "hi"}]}, + ] + marker = build_trace_instance_marker(_BENCH, "t-1#0", target=_FTP) + inject_marker_at_first_user_message(messages, marker) + content = messages[0]["content"] + assert content[0] == {"type": "text", "text": marker.strip()} + assert content[1] == {"type": "text", "text": "hi"} + + +def test_inject_no_user_turn_is_noop(): + """No user-role message -> nothing is stamped (no crash).""" + messages = [{"role": "system", "content": "sys"}] + original = copy.deepcopy(messages) + marker = build_trace_instance_marker(_BENCH, "t-1#0", target=_FTP) + inject_marker_at_first_user_message(messages, marker) + assert messages == original + + +def test_inject_none_marker_is_noop(): + """A ``None`` marker (NONE target) stamps nothing.""" + messages = [{"role": "user", "content": "hi"}] + original = copy.deepcopy(messages) + inject_marker_at_first_user_message(messages, None) + assert messages == original + + +@pytest.mark.parametrize( + "missing_messages", + [{}, {"messages": None}, {"messages": "not-a-list"}], +) +def test_stamp_tolerates_payload_without_messages_list(missing_messages): + """A payload lacking a ``messages`` list is left untouched (graceful).""" + payload = dict(missing_messages) + original = copy.deepcopy(payload) + _stamp(payload, "t-1#0") + assert payload == original diff --git a/tests/unit/graph/test_channel_store_reducer_enum.py b/tests/unit/graph/test_channel_store_reducer_enum.py new file mode 100644 index 0000000000..84f7fa1e7a --- /dev/null +++ b/tests/unit/graph/test_channel_store_reducer_enum.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""R6 -- ``VersionedChannelStore`` uses ``ReducerName`` members, never ``.value``. + +Repo contract: string enums are compared/looked-up with the member directly. +The functional tests pin that the enum-member registry lookup and the +overwrite-conflict compare keep working; the mechanical test rejects a +re-introduction of ``.value`` access on the reducer field. +""" + +from __future__ import annotations + +import inspect + +import pytest + +from aiperf.dataset.graph.models import ChannelSpec, ReducerName +from aiperf.graph import channel_store as channel_store_module +from aiperf.graph.channel_store import VersionedChannelStore +from aiperf.graph.reducers import OverwriteConflictError + + +def _store(specs: dict[str, ChannelSpec], producers: dict[str, int]): + return VersionedChannelStore( + initial={}, + channel_specs=specs, + producers_per_channel=producers, + ) + + +def test_no_reducer_value_access_in_channel_store_source(): + """String-enum contract: no ``.value`` on the ``ReducerName`` field.""" + src = inspect.getsource(channel_store_module) + assert "reducer.value" not in src + + +def test_add_messages_reducer_resolved_from_enum_member(): + """The registry lookup accepts the ``ReducerName`` member directly.""" + store = _store({"msgs": ChannelSpec(reducer=ReducerName.ADD_MESSAGES)}, {"msgs": 2}) + store.write(["msgs"], [{"role": "user", "content": "a"}], writer_node_id="n1") + store.write(["msgs"], [{"role": "user", "content": "b"}], writer_node_id="n2") + + snapshot = store.snapshot() + + assert [m["content"] for m in snapshot["msgs"]] == ["a", "b"] + + +def test_overwrite_conflict_detected_via_enum_member_compare(): + """The overwrite second-writer rejection keys on the enum member.""" + store = _store({"out": ChannelSpec(reducer=ReducerName.OVERWRITE)}, {"out": 2}) + store.write(["out"], "first", writer_node_id="n1") + + with pytest.raises(OverwriteConflictError): + store.write(["out"], "second", writer_node_id="n2") + + assert store.snapshot()["out"] == "first" diff --git a/tests/unit/graph/test_collapse_leading_start_offsets.py b/tests/unit/graph/test_collapse_leading_start_offsets.py new file mode 100644 index 0000000000..1871888459 --- /dev/null +++ b/tests/unit/graph/test_collapse_leading_start_offsets.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""R3 -- ``collapse_leading_start_offsets`` (the `--burst-phase-starts` collapse). + +The leading t*-relative phase-start offset lives on a firing's START in-edge +``min_start_delay_us`` (stamped by ``interval_order.build_interval_edges`` and +the ``snapshot_chop`` frontier re-root); NO producer stamps the node-level +field. Burst must zero exactly those leading offsets while every non-START +edge keeps its recorded inter-turn pacing. The timing strategy's +``_burst_collapse_leading_offsets`` is expected to delegate here. +""" + +from __future__ import annotations + +from aiperf.dataset.graph.models import GraphRecord, LlmNode, StaticEdge +from aiperf.graph.scheduler import collapse_leading_start_offsets + + +def _llm(output: str, *, min_start: float | None = None) -> LlmNode: + return LlmNode(prompt=[f"@{output}"], output=output, min_start_delay_us=min_start) + + +def _graph() -> GraphRecord: + return GraphRecord( + nodes={ + "a": _llm("a"), + "b": _llm("b", min_start=7_000_000.0), + "c": _llm("c"), + }, + edges=[ + StaticEdge(source="START", target="a", min_start_delay_us=5_000_000.0), + StaticEdge( + source="a", + target="b", + delay_after_predecessor_us=2_000_000.0, + min_start_delay_us=1_000_000.0, + ), + StaticEdge( + source="a", target="c", delay_after_predecessor_start_us=3_000_000.0 + ), + ], + state={}, + ) + + +def test_collapse_zeroes_start_edge_min_start_delay(): + """The START-edge leading offset (the t* anchor carrier) is zeroed.""" + collapsed = collapse_leading_start_offsets(_graph()) + + start_edges = [e for e in collapsed.edges if e.source == "START"] + assert len(start_edges) == 1 + assert start_edges[0].min_start_delay_us == 0.0 + + +def test_collapse_preserves_inter_turn_edge_pacing(): + """Non-START edges keep every recorded delay (burst governs starts only).""" + collapsed = collapse_leading_start_offsets(_graph()) + + by_target = {e.target: e for e in collapsed.edges if e.source != "START"} + assert by_target["b"].delay_after_predecessor_us == 2_000_000.0 + assert by_target["b"].min_start_delay_us == 1_000_000.0 + assert by_target["c"].delay_after_predecessor_start_us == 3_000_000.0 + + +def test_collapse_node_level_offsets_zero_leading_preserve_mid_graph(): + """Node-level ``min_start_delay_us`` is a leading offset ONLY on a node with + no real predecessor. + + Burst zeroes those (a hand-authored leading anchor) and leaves mid-graph + node-level pacing intact -- the AND-fan-in residual fold lands node-level on + a node that KEEPS a surviving-pred edge, so zeroing it would silently revert + the fold. + """ + graph = GraphRecord( + nodes={ + "x": _llm("x", min_start=4_000_000.0), + "y": _llm("y", min_start=2_000_000.0), + }, + edges=[ + StaticEdge(source="START", target="x", min_start_delay_us=6_000_000.0), + StaticEdge(source="x", target="y", delay_after_predecessor_us=1_000_000.0), + ], + state={}, + ) + + collapsed = collapse_leading_start_offsets(graph) + + # x roots at START only -> its node-level delay is a leading offset, zeroed. + assert collapsed.nodes["x"].min_start_delay_us == 0.0 + # y has a real predecessor -> its node-level delay is mid-graph pacing, kept. + assert collapsed.nodes["y"].min_start_delay_us == 2_000_000.0 + + +def test_collapse_is_pure_and_identity_preserving(): + """The input graph is untouched; untouched nodes/edges keep identity.""" + graph = _graph() + + collapsed = collapse_leading_start_offsets(graph) + + # Purity: the original still carries its offsets. + assert graph.edges[0].min_start_delay_us == 5_000_000.0 + assert graph.nodes["b"].min_start_delay_us == 7_000_000.0 + # Identity for untouched members (no needless struct rebuilds). + assert collapsed.nodes["a"] is graph.nodes["a"] + assert collapsed.edges[1] is graph.edges[1] + + +def test_collapse_noop_graph_round_trips(): + """A graph with no leading offsets comes back structurally identical.""" + graph = GraphRecord( + nodes={"a": _llm("a")}, + edges=[StaticEdge(source="START", target="a")], + state={}, + ) + + collapsed = collapse_leading_start_offsets(graph) + + assert collapsed.nodes["a"] is graph.nodes["a"] + assert collapsed.edges[0] is graph.edges[0] + + +def test_collapse_leading_offsets_preserves_folded_join_residual(): + """Burst zeroes true leading offsets, never mid-graph folded residuals. + + collapse_leading_start_offsets' contract: burst governs only the phase + START. A folded join residual sits on a node with a real (non-START) + in-edge -- zeroing it would silently revert the AND-fan-in fold on + --burst-phase-starts runs. + """ + graph = GraphRecord( + nodes={ + "b": _llm("b"), + "j": _llm("j", min_start=200_000.0), + }, + edges=[ + StaticEdge(source="START", target="b", min_start_delay_us=3_000_000.0), + StaticEdge(source="b", target="j", delay_after_predecessor_us=100_000.0), + ], + state={}, + ) + + out = collapse_leading_start_offsets(graph) + + # leading START-edge offset zeroed; the folded node-level residual survives + start_edge = next(e for e in out.edges if e.source == "START") + assert not start_edge.min_start_delay_us + assert out.nodes["j"].min_start_delay_us == 200_000.0 diff --git a/tests/unit/graph/test_content_seed_determinism.py b/tests/unit/graph/test_content_seed_determinism.py new file mode 100644 index 0000000000..87b6ac1e1c --- /dev/null +++ b/tests/unit/graph/test_content_seed_determinism.py @@ -0,0 +1,206 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Content-seed determinism through ``parse_graph_workload``. + +Content-seed determinism across in-process and spawn-started pool-worker parses +of the same run config. The weka content seed is +``resolve_graph_content_seed(run)`` -- the run ``--random-seed`` verbatim, with +NO weka-specific fallback. With an explicit seed every parse synthesizes +byte-identical content; with the seed unset (``None``) synthesis defers to the +process's ambient global RNG manager, and only the seed-independent +catalog/ordinal invariants are guaranteed. The TimingManager never parses -- it +ingests the graph_meta sidecar from the graph-typed dataset broadcast. + +The byte-identity assertions here compare the REAL synthesized bytes -- the +segment pool's materialized ``(role, content, parent_id)`` per +content-addressed segment id -- on FULL parses. ``TraceRecord.replay_outputs`` +is always empty on the weka path and proves nothing. +""" + +from __future__ import annotations + +from pathlib import Path + +import orjson + +from aiperf.config.flags.cli_config import CLIConfig +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.graph_path_catalog import build_graph_path_catalog +from aiperf.dataset.graph.models import ParsedGraph +from aiperf.dataset.graph.workload_detect import parse_graph_workload +from aiperf.timing.config import resolve_graph_content_seed +from tests.unit.conftest import make_run_from_cli + +WEKA_MIN = Path(__file__).parent / "fixtures" / "weka_min.json" + + +def _graph_run(input_file: Path = WEKA_MIN, **cli_overrides): + cfg = CLIConfig( + model_names=["test-model"], + input_file=str(input_file), + # Content tokenizer now comes from the run config (no env override): pin + # the offline builtin so the fake "test-model" doesn't trigger a HF load. + # Mirrors the CLI's fake-model -> builtin substitution the full pipeline + # applies before parsing. + tokenizer_name="builtin", + **cli_overrides, + ) + return make_run_from_cli(cfg) + + +def _pool_contents(parsed: ParsedGraph) -> dict[str, tuple[str, str, str | None]]: + """Materialized segment content keyed by content-addressed segment id. + + This is the real synthesized-content image of a full parse: comparing two + of these compares the actual bytes workers materialize, not the always-empty + ``TraceRecord.replay_outputs``. + """ + pool = parsed.segment_pool + assert pool is not None, "weka parse must surface the segment pool" + assert pool._by_id, "weka parse must carry real content segments" + return {sid: (s.role, s.content, s.parent_id) for sid, s in pool._by_id.items()} + + +def test_prompt_corpus_flows_to_synthesis_and_resolver() -> None: + """``--prompt-corpus`` lands on ``synthesis.corpus`` and the graph resolver + reads it; default (unset) resolves to ``coding``.""" + from aiperf.dataset.graph.workload_detect import _resolve_graph_corpus + + default_run = _graph_run() + assert _resolve_graph_corpus(default_run) == "coding" + + sonnet_run = _graph_run(prompt_corpus="sonnet") + assert sonnet_run.cfg.get_default_dataset().synthesis.corpus == "sonnet" + assert _resolve_graph_corpus(sonnet_run) == "sonnet" + + +def test_idle_gap_cap_flows_to_synthesis_and_resolver() -> None: + """``--synthesis-idle-gap-cap`` lands on ``synthesis.idle_gap_cap_seconds`` and + the graph resolver reads it; default (unset) resolves to 60.0.""" + from aiperf.dataset.graph.workload_detect import _resolve_graph_idle_gap_cap + + assert _resolve_graph_idle_gap_cap(_graph_run()) == 60.0 + + capped = _graph_run(synthesis_idle_gap_cap=30.0) + assert capped.cfg.get_default_dataset().synthesis.idle_gap_cap_seconds == 30.0 + assert _resolve_graph_idle_gap_cap(capped) == 30.0 + + +def test_parse_graph_workload_explicit_null_idle_gap_replays_raw_offsets( + tmp_path: Path, +) -> None: + """Explicit ``synthesis.idle_gap_cap_seconds: null`` replays RAW offsets. + + The regression both adversarial reviews of the GraphParseContext plan + proved the original plan would ship silently: the run's explicit-null + idle-gap disable collapsing into the 60s adapter default somewhere along + resolver -> ctx -> adapter, silently warping the replay geometry. The + fixture's second request starts at raw ``t=137124`` after the first ends at + ``t=2`` (a 137122s idle gap >> any cap), so the two outcomes are far apart: + unwarped replay keeps ``arrival_offset_us == 137_124_000_000``; a collapse + into the 60s cap yields ``62_000_000``. The default-cap arm pins the warped + value too, proving the fixture actually engages the cap (falsifiability). + + Construction note: the CLI converter DROPS ``None`` values + (``_converter_dataset.py``: ``if value is not None``), so + ``make_run_from_cli(synthesis_idle_gap_cap=None)`` yields the DEFAULT 60s + cap, not explicit-null. Build the run, then set + ``synthesis.idle_gap_cap_seconds = None`` on the default dataset directly + (attaching a synthesis object if absent). Explicit null IS end-user + reachable via a YAML config, so this pin covers a real configuration. + """ + from aiperf.config.dataset.trace import SynthesisConfig + + trace_file = tmp_path / "idle_gap.json" + trace_file.write_bytes( + orjson.dumps( + { + "id": "trace_idle_gap", + "models": ["M"], + "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "n", "model": "M", "in": 128, "out": 64, + "hash_ids": [1, 2], "api_time": 2.0}, + {"t": 137124.0, "type": "n", "model": "M", "in": 256, "out": 64, + "hash_ids": [1, 2, 3, 4], "api_time": 1.0}, + ], + } + ) + ) # fmt: skip + + # Falsifiability arm: the default (unset) cap DOES warp this fixture -- + # turn1 lands at end(2s) + capped idle (60s) = 62s. + warped = parse_graph_workload(_graph_run(input_file=trace_file), trace_file) + assert { + nid: n.arrival_offset_us for nid, n in warped.graph.nodes.items() + } == {"trace_idle_gap:0": 0, "trace_idle_gap:1": 62_000_000} # fmt: skip + + null_run = _graph_run(input_file=trace_file) + dataset = null_run.cfg.get_default_dataset() + if dataset.synthesis is None: + dataset.synthesis = SynthesisConfig() + dataset.synthesis.idle_gap_cap_seconds = None + + unwarped = parse_graph_workload(null_run, trace_file) + assert { + nid: n.arrival_offset_us for nid, n in unwarped.graph.nodes.items() + } == {"trace_idle_gap:0": 0, "trace_idle_gap:1": 137_124_000_000} # fmt: skip + + +def test_parse_graph_workload_random_seed_unset_two_parses_identical() -> None: + """Two full parses at random_seed=None match real content bytes + ordinals.""" + run = _graph_run() + assert run.random_seed is None, "default single run must leave --random-seed unset" + + build_plane = parse_graph_workload(run, WEKA_MIN) + timing_plane = parse_graph_workload(run, WEKA_MIN) + + # Content: the materialized segment bytes must be identical in-process + # (both parses defer to the same ambient RNG manager). + assert _pool_contents(build_plane) == _pool_contents(timing_plane) + + # Ordinals: the build/runtime addressing catalog must match too. + assert build_graph_path_catalog(build_plane) == build_graph_path_catalog( + timing_plane + ) + + +def test_parse_graph_workload_explicit_seed_threads_to_content_synthesis() -> None: + """``parse_graph_workload`` threads exactly ``resolve_graph_content_seed(run)``. + + If the seed were dropped (or a different value substituted), the direct + same-seed ``from_weka_trace`` parse would synthesize different bytes and + this comparison would fail -- seed-dependence is proven by the divergence + test below. + """ + run = _graph_run(random_seed=7) + assert resolve_graph_content_seed(run) == 7 + + via_parse = parse_graph_workload(run, WEKA_MIN) + direct_same_seed = from_weka_trace(str(WEKA_MIN), content_root_seed=7) + assert _pool_contents(via_parse) == _pool_contents(direct_same_seed) + + +def test_parse_graph_workload_different_seeds_diverge_content_not_catalog() -> None: + """Different seeds synthesize different bytes; catalog ordinals are seed-free. + + The divergence arm makes the byte-identity assertions falsifiable: if + content were seed-independent (or the pools compared vacuously), this test + would fail. + """ + parsed_seed_7 = parse_graph_workload(_graph_run(random_seed=7), WEKA_MIN) + parsed_seed_8 = parse_graph_workload(_graph_run(random_seed=8), WEKA_MIN) + + contents_7 = _pool_contents(parsed_seed_7) + contents_8 = _pool_contents(parsed_seed_8) + assert contents_7 != contents_8, "distinct seeds must synthesize distinct bytes" + # The synthesized text itself differs, not just the token-derived ids. + assert sorted(c for _, c, _ in contents_7.values()) != sorted( + c for _, c, _ in contents_8.values() + ) + + # Addressing is seed-independent: same catalog regardless of content seed. + assert build_graph_path_catalog(parsed_seed_7) == build_graph_path_catalog( + parsed_seed_8 + ) diff --git a/tests/unit/graph/test_corpus_content_synthesizer.py b/tests/unit/graph/test_corpus_content_synthesizer.py new file mode 100644 index 0000000000..81eafa31c9 --- /dev/null +++ b/tests/unit/graph/test_corpus_content_synthesizer.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CorpusContentSynthesizer tokenizer-loading behavior. + +Recovered from the deleted cross-run cache test module (the cache is gone, but +the fail-loud tokenizer contract it guarded lives on in +:meth:`CorpusContentSynthesizer._build_generator`). +""" + +import pytest + +from aiperf.common.tokenizer import BUILTIN_TOKENIZER_NAME, Tokenizer +from aiperf.dataset import _tokenizer_preload +from aiperf.dataset.graph.adapters.shared import content as _shared_content + + +def test_tokenizer_load_failure_fails_loud_no_builtin_substitution( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A transient tokenizer load failure raises -- it never silently uses builtin. + + Silently substituting the builtin tokenizer would produce builtin-sized + content while claiming the requested tokenizer. The synthesizer must + propagate the failure so mislabeled content is never produced. + """ + _tokenizer_preload.clear_preloaded() + _shared_content.CorpusContentSynthesizer.reset_worker_cache() + + requested = "some-org/some-private-tokenizer" + boom = RuntimeError("HF unreachable") + + def _always_fail(name: str, **kwargs: object) -> Tokenizer: + raise boom + + monkeypatch.setattr(Tokenizer, "from_pretrained", staticmethod(_always_fail)) + + with pytest.raises(RuntimeError, match="HF unreachable"): + _shared_content.CorpusContentSynthesizer._build_generator( + requested, _shared_content.PROMPT_CORPUS_CODING + ) + + # Sanity: the requested name is not the builtin name, so the old code path + # would have substituted builtin instead of raising. + assert requested != BUILTIN_TOKENIZER_NAME diff --git a/tests/unit/graph/test_credit_counter_graph_bypass.py b/tests/unit/graph/test_credit_counter_graph_bypass.py new file mode 100644 index 0000000000..8a4e6ad7f2 --- /dev/null +++ b/tests/unit/graph/test_credit_counter_graph_bypass.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fix C (F2/A2) — graph credits bypass CreditCounter session arithmetic. + +The proven bug: graph nodes flow through ``CreditCounter.increment_sent`` with +``agent_depth==0`` and per-node ``turn_index==0``, so each NODE bumped +``_sent_sessions`` / ``_total_session_turns`` and tripped ``is_final_credit`` +after the Nth node (freezing the sent-count mid-trace -> lost records). + +The fix gates on ``turn.trace_id is not None``: graph credits bump ONLY +``_requests_sent`` (the per-node record count) and NEVER flip ``is_final_credit``. +These tests pin both the graph bypass and the unchanged non-graph path. +""" + +from __future__ import annotations + +from aiperf.common.enums import CreditPhase +from aiperf.credit.structs import TurnToSend +from aiperf.plugin.enums import TimingMode +from aiperf.timing.config import CreditPhaseConfig +from aiperf.timing.phase.credit_counter import CreditCounter + + +def _graph_turn(trace_id: str = "t-1#0") -> TurnToSend: + """A graph node dispatch: distinct nodes all fire with turn_index==0.""" + return TurnToSend( + conversation_id=trace_id, + x_correlation_id=f"x|{trace_id}", + turn_index=0, + num_turns=1, + trace_id=trace_id, + ) + + +def _linear_turn(turn_index: int, num_turns: int) -> TurnToSend: + return TurnToSend( + conversation_id="c0", + x_correlation_id="x0", + turn_index=turn_index, + num_turns=num_turns, + ) + + +def _counter(**caps) -> CreditCounter: + config = CreditPhaseConfig( + phase=CreditPhase.PROFILING, + timing_mode=TimingMode.GRAPH_IR, + **caps, + ) + return CreditCounter(config) + + +def test_graph_credits_do_not_increment_sessions(): + """Three distinct-node graph credits bump requests_sent only, never sessions.""" + counter = _counter() + for _ in range(3): + _, is_final = counter.increment_sent(_graph_turn()) + assert is_final is False + + assert counter.requests_sent == 3 + assert counter.sent_sessions == 0 + assert counter.total_session_turns == 0 + assert counter.root_requests_sent == 0 + + +def test_graph_credits_never_trip_session_cap(): + """The proven bug: with expected_num_sessions=1, is_final_credit fired on + the FIRST node. After the fix it must NEVER fire for a graph credit.""" + counter = _counter(expected_num_sessions=1) + for _ in range(5): + _, is_final = counter.increment_sent(_graph_turn()) + assert is_final is False + assert counter.sent_sessions == 0 + + +def test_graph_credits_never_trip_request_cap_via_counter(): + """Even with a request-count cap, the COUNTER must not auto-flip + is_final_credit for graph credits -- the strategy owns the cap (the + RequestCountStopCondition gate enforces it at issue time).""" + counter = _counter(total_expected_requests=2) + for _ in range(3): + _, is_final = counter.increment_sent(_graph_turn()) + assert is_final is False + assert counter.requests_sent == 3 + + +def test_non_graph_session_accounting_unchanged(): + """A 2-turn linear session counts as ONE session with 2 turns, and + is_final_credit fires only once both clauses are satisfied -- the graph + trace_id gate must not perturb this.""" + counter = _counter(expected_num_sessions=1) + + _, final0 = counter.increment_sent(_linear_turn(0, 2)) + assert counter.sent_sessions == 1 + assert counter.total_session_turns == 2 + assert counter.requests_sent == 1 + assert final0 is False + + _, final1 = counter.increment_sent(_linear_turn(1, 2)) + assert counter.sent_sessions == 1 + assert counter.requests_sent == 2 + assert final1 is True + + +def test_graph_returns_do_not_touch_session_counters(): + """TC3: graph returns bypass session counters like the sent side. + + Every graph credit is minted ``turn_index=0, num_turns=1`` so + ``is_final_turn`` is always True at return time; without the bypass each + NODE return bumped ``completed_sessions`` while ``sent_sessions`` stayed 0 + (negative ``in_flight_sessions``, bogus progress %). Session stats must + stay 0 on BOTH sides for a graph phase; progress is request-count driven. + """ + counter = _counter() + for _ in range(4): + counter.increment_sent(_graph_turn()) + for _ in range(2): + counter.increment_returned(is_final_turn=True, cancelled=False) + counter.increment_returned(is_final_turn=True, cancelled=False, errored=True) + counter.increment_returned(is_final_turn=True, cancelled=True) + + assert counter.completed_sessions == 0 + assert counter.cancelled_sessions == 0 + assert counter.sent_sessions == 0 + assert counter.in_flight_sessions == 0 # never negative + # Request-level counters still drive progress + the all-returned invariant. + assert counter.requests_completed == 3 + assert counter.requests_cancelled == 1 + assert counter.request_errors == 1 + assert counter.in_flight == 0 + + +def test_graph_returns_still_drive_all_returned_invariant(): + """The completion barrier stays request-count driven for graph phases.""" + counter = _counter() + for _ in range(2): + counter.increment_sent(_graph_turn()) + counter.freeze_sent_counts() + + assert counter.increment_returned(is_final_turn=True, cancelled=False) is False + assert counter.increment_returned(is_final_turn=True, cancelled=False) is True + assert counter.check_all_returned_or_cancelled() is True + + +def test_non_graph_phase_returned_session_accounting_unchanged(): + """A linear (REQUEST_RATE) phase still counts completed/cancelled sessions.""" + config = CreditPhaseConfig( + phase=CreditPhase.PROFILING, + timing_mode=TimingMode.REQUEST_RATE, + ) + counter = CreditCounter(config) + counter.increment_sent(_linear_turn(0, 1)) + counter.increment_sent(_linear_turn(0, 1)) + + counter.increment_returned(is_final_turn=True, cancelled=False) + counter.increment_returned(is_final_turn=True, cancelled=True) + + assert counter.completed_sessions == 1 + assert counter.cancelled_sessions == 1 + + +def test_non_graph_dag_child_accounting_unchanged(): + """A DAG child (agent_depth > 0, trace_id None) still bumps requests_sent + only and honors the request cap -- unchanged by the graph gate.""" + counter = _counter(total_expected_requests=2) + child = TurnToSend( + conversation_id="c0", + x_correlation_id="x0::child", + turn_index=0, + num_turns=1, + agent_depth=1, + ) + _, final0 = counter.increment_sent(child) + assert counter.requests_sent == 1 + assert counter.sent_sessions == 0 + assert final0 is False + + _, final1 = counter.increment_sent(child) + assert counter.requests_sent == 2 + assert final1 is True # request cap crossed on a child diff --git a/tests/unit/graph/test_credit_dispatch_adapter.py b/tests/unit/graph/test_credit_dispatch_adapter.py new file mode 100644 index 0000000000..56e48e67e2 --- /dev/null +++ b/tests/unit/graph/test_credit_dispatch_adapter.py @@ -0,0 +1,696 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""R4/R2 — CreditDispatchAdapter (Future bridge) + unconditional return hook. + +Component-level tests with a FAKE issuer + return loop. No real worker, no ZMQ. + +The adapter bridges the dataflow ``TraceExecutor.dispatch(node, request, ctx)`` +contract onto the v1 credit pipeline: it issues a real ``TurnToSend`` carrying +``trace_id``/``node_ordinal``/``phase_variant``, parks an ``asyncio.Future`` +keyed by a per-trajectory ``(x_correlation_id, turn_index)`` coordinate, and +awaits it. The strategy resolves that Future on the correlated ``CreditReturn`` +(success/cancel/error). A timeout guard rejects rather than hangs when a return +is dropped. + +Identity contract under test: +* node id is ``{scope}:{turn}`` (0-based turn within the trajectory scope); +* ``x_correlation_id`` is ``{conversation_id}::{uuid4().hex}`` -- ONE per scope + per adapter instance (all turns of a trajectory share it); +* ``turn_index`` is the node's own 0-based turn parsed from the node id. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass + +import pytest + +from aiperf.dataset.graph.graph_path_catalog import CatalogContext +from aiperf.graph.placement import DispatchRequest, PlacementContext + +pytestmark = pytest.mark.asyncio + + +# --------------------------------------------------------------------------- +# Fakes +# --------------------------------------------------------------------------- + + +class FakeIssuer: + """Records issued graph TurnToSends; lets the test echo a CreditReturn. + + Mirrors the surface the adapter calls: a single ``issue_graph_credit`` + entrypoint that records the turn and reports "can send more" (True). + """ + + def __init__(self) -> None: + self.sent: list[object] = [] + + async def issue_graph_credit(self, turn: object) -> bool: + self.sent.append(turn) + return True + + +@dataclass +class FakeCredit: + """Minimal Credit-like carrying the correlation identity the bridge keys on.""" + + x_correlation_id: str + turn_index: int + trace_id: str + node_ordinal: int + phase_variant: str = "profiling" + + +class FalseIssuer: + """Issuer whose ``issue_graph_credit`` always refuses (returns False). + + Mirrors the real ``CreditIssuer.issue_graph_credit`` stop-gate / prefill-gate + refusal: no credit reaches the wire, so no ``CreditReturn`` can ever arrive to + resolve the parked Future. Records the turn it was asked to issue so the test + can confirm the adapter still minted addressing before being refused. + """ + + def __init__(self) -> None: + self.sent: list[object] = [] + + async def issue_graph_credit(self, turn: object) -> bool: + self.sent.append(turn) + return False + + +@dataclass +class FakeLlmNode: + """Stand-in IR node; the adapter does not read node fields for weka.""" + + output: str = "out" + + +def _ctx(parent_trace_id: str, node_id: str) -> PlacementContext: + return PlacementContext(parent_trace_id=parent_trace_id, parent_node_id=node_id) + + +def _request(node_id: str) -> DispatchRequest: + return DispatchRequest(node_id=node_id) + + +def _catalog(trace_id: str, node_key_to_ordinal: dict[str, int]) -> CatalogContext: + return CatalogContext( + catalog={trace_id: dict(node_key_to_ordinal)}, + ) + + +def _make_adapter(issuer: FakeIssuer, trace_id: str, ordinals: dict[str, int], **kw): + from aiperf.graph.credit_dispatch_adapter import CreditDispatchAdapter + + return CreditDispatchAdapter( + credit_issuer=issuer, + catalog_context=_catalog(trace_id, ordinals), + trace_id=trace_id, + phase_variant=kw.pop("phase_variant", "profiling"), + parent_correlation_id=kw.pop("parent_correlation_id", None), + **kw, + ) + + +# --------------------------------------------------------------------------- +# Happy path: issue carries addressing; echoing the return resolves the await +# --------------------------------------------------------------------------- + + +async def test_dispatch_issues_credit_with_graph_addressing(): + issuer = FakeIssuer() + adapter = _make_adapter(issuer, "t0", {"t0:0": 7}, phase_variant="profiling") + + task = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("t0:0"), _ctx("t0", "t0:0")) + ) + await asyncio.sleep(0) # let it issue + park the future + + assert len(issuer.sent) == 1 + turn = issuer.sent[0] + assert turn.trace_id == "t0" + assert turn.node_ordinal == 7 + assert turn.phase_variant == "profiling" + # conversation_id is the root-scope TEMPLATE id; x_correlation_id is + # {conversation_id}::{nonce}; turn_index is the node's own 0-based turn. + assert turn.conversation_id == "t0" + assert turn.x_correlation_id.startswith("t0::") + assert turn.turn_index == 0 + + adapter.resolve(_credit_for(turn), error=None, cancelled=False) + result = await task + assert isinstance(result, str) + + +async def test_dispatch_returns_placeholder_str_on_success(): + issuer = FakeIssuer() + adapter = _make_adapter(issuer, "t0", {"t0:0": 0}) + task = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("t0:0"), _ctx("t0", "t0:0")) + ) + await asyncio.sleep(0) + adapter.resolve(_credit_for(issuer.sent[0]), error=None, cancelled=False) + assert isinstance(await task, str) + + +def _credit_for(turn: object) -> FakeCredit: + return FakeCredit( + x_correlation_id=turn.x_correlation_id, + turn_index=turn.turn_index, + trace_id=turn.trace_id, + node_ordinal=turn.node_ordinal, + phase_variant=turn.phase_variant, + ) + + +# --------------------------------------------------------------------------- +# Session-routing identity: real per-trajectory num_turns + root corr +# --------------------------------------------------------------------------- + + +async def test_turns_carry_recorded_num_turns_and_root_corr(): + """num_turns is the trajectory's RECORDED turn count from the catalog + (is_final_turn = the recorded session-final fact for session-routing + bind/close semantics), and every turn carries the instance's ROOT + trajectory corr as root_correlation_id.""" + issuer = FakeIssuer() + adapter = _make_adapter(issuer, "t0", {"t0:0": 0, "t0:1": 1, "child:0": 2}) + + tasks = [] + for nid in ("t0:0", "t0:1", "child:0"): + tasks.append( + asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request(nid), _ctx("t0", nid)) + ) + ) + await asyncio.sleep(0) + + root0, root1, child0 = issuer.sent + assert (root0.num_turns, root1.num_turns) == (2, 2) + assert root0.is_final_turn is False + assert root1.is_final_turn is True # recorded session-final turn + assert child0.num_turns == 1 + assert child0.is_final_turn is True + # Root corr: one stable id for the whole instance, equal to the root + # trajectory's own corr. + assert root0.root_correlation_id == root0.x_correlation_id + assert child0.root_correlation_id == root0.x_correlation_id + assert child0.x_correlation_id != child0.root_correlation_id + + for turn, task in zip(issuer.sent, tasks, strict=True): + adapter.resolve(_credit_for(turn), error=None, cancelled=False) + assert isinstance(await task, str) + + +async def test_unknown_runtime_scope_reads_non_final(): + """A runtime scope the catalog does not know must never read as final: + a wrong session-close on the wire is worse than a missing one.""" + issuer = FakeIssuer() + adapter = _make_adapter(issuer, "t0", {"t0:0": 0}) + + task = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("mystery:0"), _ctx("t0", "mystery:0")) + ) + await asyncio.sleep(0) + + (turn,) = issuer.sent + assert turn.is_final_turn is False + + adapter.resolve(_credit_for(turn), error=None, cancelled=False) + assert isinstance(await task, str) + + +# --------------------------------------------------------------------------- +# Native-authored node ids (no {scope}:{turn} shape) +# --------------------------------------------------------------------------- + + +async def test_native_bare_node_ids_share_one_root_trajectory(): + """Author-chosen ids like ``"plan"`` map to the ROOT trajectory with the + catalog ordinal as the turn: one corr for the whole graph (dynamic-slot + content stays on one sticky session), unique waiter keys per node. + Regression: this used to crash with ``ValueError: invalid literal for + int()`` at the first dispatch of any native graph.""" + issuer = FakeIssuer() + adapter = _make_adapter(issuer, "t0", {"plan": 0, "review": 1}) + + tasks = [ + asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request(nid), _ctx("t0", nid)) + ) + for nid in ("plan", "review") + ] + await asyncio.sleep(0) + + assert len(issuer.sent) == 2 + plan, review = issuer.sent + assert plan.x_correlation_id == review.x_correlation_id + assert plan.x_correlation_id.startswith("t0::") + assert plan.conversation_id == review.conversation_id == "t0" + assert (plan.turn_index, review.turn_index) == (0, 1) + assert (plan.node_ordinal, review.node_ordinal) == (0, 1) + + for turn, task in zip(issuer.sent, tasks, strict=True): + adapter.resolve(_credit_for(turn), error=None, cancelled=False) + assert isinstance(await task, str) + + +async def test_native_non_int_tail_id_uses_root_fallback(): + """A colon-bearing id whose tail is not an int (``phase:review``) is a + native-authored id, not a {scope}:{turn} coordinate.""" + issuer = FakeIssuer() + adapter = _make_adapter(issuer, "t0", {"phase:review": 3}) + + task = asyncio.create_task( + adapter.dispatch( + FakeLlmNode(), _request("phase:review"), _ctx("t0", "phase:review") + ) + ) + await asyncio.sleep(0) + + (turn,) = issuer.sent + assert turn.conversation_id == "t0" + assert turn.x_correlation_id.startswith("t0::") + assert turn.turn_index == 3 # catalog ordinal, NOT parsed from the id + + adapter.resolve(_credit_for(turn), error=None, cancelled=False) + assert isinstance(await task, str) + + +# --------------------------------------------------------------------------- +# Cancel / error reject the Future (no hang) +# --------------------------------------------------------------------------- + + +async def test_error_return_rejects_future_no_hang(): + issuer = FakeIssuer() + adapter = _make_adapter(issuer, "t0", {"t0:0": 0}) + task = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("t0:0"), _ctx("t0", "t0:0")) + ) + await asyncio.sleep(0) + adapter.resolve(_credit_for(issuer.sent[0]), error="boom", cancelled=False) + from aiperf.graph.credit_dispatch_adapter import GraphDispatchError + + with pytest.raises(GraphDispatchError): + await asyncio.wait_for(task, timeout=1.0) + + +async def test_cancelled_return_rejects_future_no_hang(): + issuer = FakeIssuer() + adapter = _make_adapter(issuer, "t0", {"t0:0": 0}) + task = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("t0:0"), _ctx("t0", "t0:0")) + ) + await asyncio.sleep(0) + adapter.resolve(_credit_for(issuer.sent[0]), error=None, cancelled=True) + from aiperf.graph.credit_dispatch_adapter import GraphDispatchError + + with pytest.raises(GraphDispatchError): + await asyncio.wait_for(task, timeout=1.0) + + +# --------------------------------------------------------------------------- +# Timeout guard rejects rather than hangs +# --------------------------------------------------------------------------- + + +async def test_timeout_rejects_future_when_no_return_arrives(): + issuer = FakeIssuer() + adapter = _make_adapter(issuer, "t0", {"t0:0": 0}, dispatch_timeout_s=0.01) + with pytest.raises(asyncio.TimeoutError): + await adapter.dispatch(FakeLlmNode(), _request("t0:0"), _ctx("t0", "t0:0")) + + +# --------------------------------------------------------------------------- +# Correlation-key uniqueness: distinct scopes (loop iterations, fork branches) +# park distinct Futures. The differentiator now rides the NODE ID's scope +# (``{scope}:{turn}``), not ``ctx.parent_trace_id`` (which ``_mint`` ignores). +# --------------------------------------------------------------------------- + + +async def test_loop_refire_distinct_scopes_do_not_collide(): + """Two in-flight dispatches in distinct loop-iteration scopes must park + distinct Futures; resolving the first must not orphan the second. + + Loop iterations scope their nodes as ``{trace}::loop#N:{turn}`` (the + executor's loop scoping now lives in the node id, so distinct scopes mint + distinct trajectory correlation ids).""" + issuer = FakeIssuer() + adapter = _make_adapter(issuer, "t0", {"t0::loop#1:0": 0, "t0::loop#2:0": 0}) + + t1 = asyncio.create_task( + adapter.dispatch( + FakeLlmNode(), _request("t0::loop#1:0"), _ctx("t0", "t0::loop#1:0") + ) + ) + t2 = asyncio.create_task( + adapter.dispatch( + FakeLlmNode(), _request("t0::loop#2:0"), _ctx("t0", "t0::loop#2:0") + ) + ) + await asyncio.sleep(0) + assert len(issuer.sent) == 2 + # Both in-flight at once -> two distinct waiters. + assert adapter.inflight_count == 2 + # Distinct scopes -> distinct trajectory correlation ids. + assert issuer.sent[0].x_correlation_id != issuer.sent[1].x_correlation_id + + adapter.resolve(_credit_for(issuer.sent[0]), error=None, cancelled=False) + # First-resolved ordering: iteration 1 completes; iteration 2 stays parked + # (its Future is unresolved, so it CANNOT have completed), NOT orphaned. + assert isinstance(await t1, str) + assert not t2.done() + assert adapter.inflight_count == 1 + adapter.resolve(_credit_for(issuer.sent[1]), error=None, cancelled=False) + assert isinstance(await t2, str) + + +async def test_fork_branch_same_ordinal_does_not_collide(): + """Two fork branches reaching nodes that share an ordinal across distinct + branch scopes must park distinct Futures.""" + issuer = FakeIssuer() + # Both branches share ordinal 0, but distinct branch scopes (``::brA`` / + # ``::brB``) so they mint distinct trajectory correlation ids. + adapter = _make_adapter(issuer, "t0", {"t0::brA:0": 0, "t0::brB:0": 0}) + a = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("t0::brA:0"), _ctx("t0", "t0::brA:0")) + ) + b = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("t0::brB:0"), _ctx("t0", "t0::brB:0")) + ) + await asyncio.sleep(0) + assert adapter.inflight_count == 2 + adapter.resolve(_credit_for(issuer.sent[0]), error=None, cancelled=False) + adapter.resolve(_credit_for(issuer.sent[1]), error=None, cancelled=False) + assert isinstance(await a, str) + assert isinstance(await b, str) + + +async def test_duplicate_inflight_same_coordinate_raises(): + """Re-firing the SAME ``(scope, turn)`` node while its first dispatch is still + in flight raises ``RuntimeError`` -- the executor fires each node at most once + per instance run, so a duplicate waiter key is a programming error, never a + silently-shared Future.""" + issuer = FakeIssuer() + adapter = _make_adapter(issuer, "t0", {"t0:0": 0}) + first = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("t0:0"), _ctx("t0", "t0:0")) + ) + await asyncio.sleep(0) + with pytest.raises(RuntimeError): + await adapter.dispatch(FakeLlmNode(), _request("t0:0"), _ctx("t0", "t0:0")) + adapter.resolve(_credit_for(issuer.sent[0]), error=None, cancelled=False) + assert isinstance(await first, str) + + +# --------------------------------------------------------------------------- +# Stall guard (adv2 A1): a refused issue must reject PROMPTLY, not time out +# --------------------------------------------------------------------------- + + +async def test_refused_issue_rejects_promptly_not_after_timeout(): + """When ``issue_graph_credit`` returns False, no credit reaches the wire so + no ``CreditReturn`` can ever resolve the parked Future. The adapter MUST pop + the waiter and reject immediately rather than await out ``dispatch_timeout_s``. + + Guard: a generous timeout (the would-be hang) with a tight ``wait_for`` bound + proves promptness -- the dispatch resolves in well under the timeout. + """ + issuer = FalseIssuer() + # dispatch_timeout_s large enough that timing-out would dwarf the wait_for + # bound; if the fix is missing, dispatch hangs until this 30s elapses. + adapter = _make_adapter(issuer, "t0", {"t0:0": 0}, dispatch_timeout_s=30.0) + + from aiperf.graph.credit_dispatch_adapter import GraphDispatchError + + with pytest.raises(GraphDispatchError): + await asyncio.wait_for( + adapter.dispatch(FakeLlmNode(), _request("t0:0"), _ctx("t0", "t0:0")), + timeout=1.0, + ) + # The adapter minted addressing before being refused, then unwound cleanly. + assert len(issuer.sent) == 1 + assert adapter.inflight_count == 0 + + +async def test_unknown_return_is_dropped_not_raised(): + """A return whose key matches no parked waiter is a graceful no-op.""" + issuer = FakeIssuer() + adapter = _make_adapter(issuer, "t0", {"t0:0": 0}) + stray = FakeCredit( + x_correlation_id="x-nope", turn_index=99, trace_id="t0", node_ordinal=0 + ) + adapter.resolve(stray, error=None, cancelled=False) # must not raise + assert adapter.inflight_count == 0 + + +# --------------------------------------------------------------------------- +# ``on_drained`` fires when the in-flight set empties. +# +# The strategy registers an ``on_drained`` callback so it can defer popping a +# detached-spawn instance's adapter from its de-mux registry until the adapter is +# idle. The adapter must invoke that callback EXACTLY when a return drains the +# last in-flight waiter, and never while another dispatch is still parked. +# --------------------------------------------------------------------------- + + +async def test_on_drained_fires_when_last_waiter_resolves(): + """A single dispatch's return drains the waiter set -> ``on_drained`` fires.""" + issuer = FakeIssuer() + drained: list[object] = [] + adapter = _make_adapter( + issuer, "t0", {"t0:0": 0}, on_drained=lambda a: drained.append(a) + ) + task = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("t0:0"), _ctx("t0", "t0:0")) + ) + await asyncio.sleep(0) + assert adapter.inflight_count == 1 + assert drained == [] # not yet drained + + adapter.resolve(_credit_for(issuer.sent[0]), error=None, cancelled=False) + assert isinstance(await task, str) + assert drained == [adapter] # fired exactly once, with self + assert adapter.inflight_count == 0 + + +async def test_on_drained_not_fired_while_another_dispatch_parked(): + """With two in-flight dispatches, resolving the first must NOT fire + ``on_drained`` (the set is not empty); resolving the second then fires it.""" + issuer = FakeIssuer() + drained: list[object] = [] + adapter = _make_adapter( + issuer, + "t0", + {"t0::brA:0": 0, "t0::brB:0": 0}, + on_drained=lambda a: drained.append(a), + ) + t1 = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("t0::brA:0"), _ctx("t0", "t0::brA:0")) + ) + t2 = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("t0::brB:0"), _ctx("t0", "t0::brB:0")) + ) + await asyncio.sleep(0) + assert adapter.inflight_count == 2 + + adapter.resolve(_credit_for(issuer.sent[0]), error=None, cancelled=False) + await asyncio.sleep(0) + assert drained == [] # second still parked -> NOT drained + assert adapter.inflight_count == 1 + + adapter.resolve(_credit_for(issuer.sent[1]), error=None, cancelled=False) + assert isinstance(await t1, str) + assert isinstance(await t2, str) + assert drained == [adapter] # drained exactly once, when the set emptied + + +async def test_on_drained_fires_on_rejected_return_too(): + """A cancel/error return is a terminal drain too -> ``on_drained`` fires.""" + issuer = FakeIssuer() + drained: list[object] = [] + adapter = _make_adapter( + issuer, "t0", {"t0:0": 0}, on_drained=lambda a: drained.append(a) + ) + task = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("t0:0"), _ctx("t0", "t0:0")) + ) + await asyncio.sleep(0) + adapter.resolve(_credit_for(issuer.sent[0]), error="boom", cancelled=False) + from aiperf.graph.credit_dispatch_adapter import GraphDispatchError + + with pytest.raises(GraphDispatchError): + await asyncio.wait_for(task, timeout=1.0) + assert drained == [adapter] + + +# --------------------------------------------------------------------------- +# context-overflow early-termination +# --------------------------------------------------------------------------- + + +async def test_overflow_error_raises_node_overflow_terminate(): + """A context-overflow error body resolves the parked Future with + ``_NodeOverflowTerminate`` (NOT ``GraphDispatchError``), so the executor + terminates the trajectory cleanly instead of unwinding it as a trace error. + + A non-final turn whose response is a context-overflow stops dispatching the + rest of that trace's turns -- later turns carry more context and would only + overflow too. + """ + issuer = FakeIssuer() + adapter = _make_adapter(issuer, "t0", {"t0:0": 0}) + task = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("t0:0"), _ctx("t0", "t0:0")) + ) + await asyncio.sleep(0) + + overflow_body = ( + '{"error": {"message": "This model\'s maximum context length is 8192 ' + 'tokens (context_length_exceeded)", "code": "context_length_exceeded"}}' + ) + adapter.resolve(_credit_for(issuer.sent[0]), error=overflow_body, cancelled=False) + + from aiperf.graph.context import _NodeExpectedExit + from aiperf.graph.credit_dispatch_adapter import ( + GraphDispatchError, + _NodeOverflowTerminate, + ) + + with pytest.raises(_NodeOverflowTerminate) as exc_info: + await asyncio.wait_for(task, timeout=1.0) + # Subclass of the executor's clean-exit sentinel so ``_run_node`` catches it + # on the expected-exit branch (no TaskGroup error cascade). + assert isinstance(exc_info.value, _NodeExpectedExit) + # And NOT the generic error type, so it is not mistaken for a real failure. + assert not isinstance(exc_info.value, GraphDispatchError) + + +async def test_non_overflow_error_still_raises_graph_dispatch_error(): + """A non-overflow error body keeps the existing ``GraphDispatchError`` path + (unwinds the trace as a real error) -- overflow handling does not regress it.""" + issuer = FakeIssuer() + adapter = _make_adapter(issuer, "t0", {"t0:0": 0}) + task = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("t0:0"), _ctx("t0", "t0:0")) + ) + await asyncio.sleep(0) + adapter.resolve( + _credit_for(issuer.sent[0]), error="HTTP 500 internal", cancelled=False + ) + from aiperf.graph.credit_dispatch_adapter import ( + GraphDispatchError, + _NodeOverflowTerminate, + ) + + with pytest.raises(GraphDispatchError) as exc_info: + await asyncio.wait_for(task, timeout=1.0) + assert not isinstance(exc_info.value, _NodeOverflowTerminate) + + +# --------------------------------------------------------------------------- +# dag identity map: agent_depth + derived parent correlation id +# --------------------------------------------------------------------------- + + +async def test_dispatch_node_identity_map_stamps_depth_and_parent_corr(): + """A ``node_identity`` hit stamps the mapped depth and derives the parent + NODE's correlation id from the same trajectory-scope corr ``_mint`` folds. + + The child node lives in scope ``child`` (turn 0); its triggering parent node + lives in the ROOT scope ``t0`` (turn 0). The derived parent correlation id is + therefore the root scope's ``{conversation_id}::{nonce}`` corr -- distinct + from the child's own corr.""" + issuer = FakeIssuer() + adapter = _make_adapter( + issuer, + "t0", + {"child:0": 0, "t0:0": 1}, + node_identity={"child:0": (1, "t0:0"), "t0:0": (0, None)}, + ) + task = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("child:0"), _ctx("t0", "child:0")) + ) + await asyncio.sleep(0) + + turn = issuer.sent[0] + assert turn.agent_depth == 1 + # The derived parent corr IS the trajectory-scope corr _mint folds for the + # parent node's scope (root scope "t0"), a {conversation_id}::{nonce} shape. + parent_x_corr, _, _ = adapter._mint("t0:0", 1) + assert turn.parent_correlation_id == parent_x_corr + assert parent_x_corr.startswith("t0::") + # ... and distinct from the child's own trajectory corr. + assert turn.parent_correlation_id != turn.x_correlation_id + + adapter.resolve(_credit_for(turn), error=None, cancelled=False) + await task + + +async def test_dispatch_node_identity_parent_none_falls_back_to_ctor_parent(): + """A map hit with no parent node (root / pre-session child) keeps the + constructor's ``parent_correlation_id`` fallback.""" + issuer = FakeIssuer() + adapter = _make_adapter( + issuer, + "t0", + {"t0:0": 0}, + node_identity={"t0:0": (1, None)}, + parent_correlation_id="legacy-parent", + ) + task = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("t0:0"), _ctx("t0", "t0:0")) + ) + await asyncio.sleep(0) + + turn = issuer.sent[0] + assert turn.agent_depth == 1 + assert turn.parent_correlation_id == "legacy-parent" + + adapter.resolve(_credit_for(turn), error=None, cancelled=False) + await task + + +async def test_dispatch_node_identity_miss_uses_root_identity(): + """A node absent from the map dispatches with root-chain identity.""" + issuer = FakeIssuer() + adapter = _make_adapter( + issuer, + "t0", + {"t0:0": 0}, + node_identity={"other:0": (3, "p:0")}, + ) + task = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("t0:0"), _ctx("t0", "t0:0")) + ) + await asyncio.sleep(0) + + turn = issuer.sent[0] + assert turn.agent_depth == 0 + assert turn.parent_correlation_id is None + + adapter.resolve(_credit_for(turn), error=None, cancelled=False) + await task + + +async def test_dispatch_without_node_identity_map_byte_identical_to_today(): + """No map (weka/dynamo default) => depth 0 and the constructor's + ``parent_correlation_id`` pass-through, exactly as before.""" + issuer = FakeIssuer() + adapter = _make_adapter( + issuer, "t0", {"t0:0": 0}, parent_correlation_id="pc-passthrough" + ) + task = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("t0:0"), _ctx("t0", "t0:0")) + ) + await asyncio.sleep(0) + + turn = issuer.sent[0] + assert turn.agent_depth == 0 + assert turn.parent_correlation_id == "pc-passthrough" + + adapter.resolve(_credit_for(turn), error=None, cancelled=False) + await task diff --git a/tests/unit/graph/test_credit_struct_fields.py b/tests/unit/graph/test_credit_struct_fields.py new file mode 100644 index 0000000000..6de1d6f058 --- /dev/null +++ b/tests/unit/graph/test_credit_struct_fields.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Graph addressing fields on TurnToSend/Credit (Task 7).""" + +from aiperf.credit.structs import Credit, TurnToSend + + +def test_turn_to_send_carries_graph_addressing(): + t = TurnToSend( + conversation_id="c", + x_correlation_id="x", + turn_index=0, + num_turns=1, + trace_id="t0", + node_ordinal=3, + phase_variant="profiling", + ) + assert (t.trace_id, t.node_ordinal, t.phase_variant) == ("t0", 3, "profiling") + + +def test_turn_to_send_graph_addressing_defaults(): + t = TurnToSend(conversation_id="c", x_correlation_id="x", turn_index=0, num_turns=1) + assert (t.trace_id, t.node_ordinal, t.phase_variant) == (None, None, "profiling") + + +def test_credit_carries_graph_addressing(): + c = Credit( + id=0, + phase="profiling", + conversation_id="c", + x_correlation_id="x", + turn_index=0, + num_turns=1, + issued_at_ns=0, + trace_id="t0", + node_ordinal=3, + phase_variant="profiling", + ) + assert (c.trace_id, c.node_ordinal, c.phase_variant) == ("t0", 3, "profiling") diff --git a/tests/unit/graph/test_default_pool_model_override.py b/tests/unit/graph/test_default_pool_model_override.py new file mode 100644 index 0000000000..d2a3d2fe38 --- /dev/null +++ b/tests/unit/graph/test_default_pool_model_override.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Regression: a weka trie node honors a run-level ``--model`` override on the wire. + +The segment-trie IR (the sole weka path) carries each node's RECORDED ``model`` +verbatim in its build-time ``dispatch_overrides`` manifest -- there is NO build- +time model map (the legacy ``build_weka_model_map`` remapped recorded models to +``endpoint.model_names``; the trie path does not). The user's ``--model`` override +is layered run-level at the worker via +:func:`~aiperf.graph.worker_materialize.apply_run_level_payload_options` +(``endpoint.extra`` carrying ``("model", ...)`` CLOBBERS the per-node key, agentx +``payload.update(endpoint.extra)`` precedence), so the recorded model NEVER reaches +the wire when an override is configured. + +These tests build the real interned unified trie store from a weka fixture, +materialize a node, and assert the override wins on the wire while the recorded +model leaks only when no override is configured. +""" + +from pathlib import Path + +import pytest + +from aiperf.common.models import EndpointInfo +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.segment_ir.store_builder import ( + build_unified_trie_store_interned, +) +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) +from aiperf.graph.worker_materialize import ( + apply_run_level_payload_options, + materialize_graph_request_unified, +) + +FIX = Path(__file__).parent / "fixtures" / "weka_min.json" + +# The single recorded model in the fixture -- the value the trie manifest carries +# and that must NOT reach the wire once an override is configured. +_RECORDED_MODEL = "claude-opus-4-5-20251101" +# The user's run-level ``--model`` override -- the value that MUST reach the wire. +_OVERRIDE_MODEL = "my-served-model" + + +async def _build_trie_store( + tmp_path: Path, benchmark_id: str +) -> tuple[str, int, GraphSegmentUnifiedClient]: + """Build the interned unified trie store for the fixture's sole trace. + + Returns ``(trace_id, first_node_ordinal, client)`` with the client opened; + the caller closes it. + """ + parsed = from_weka_trace(str(FIX)) + assert parsed.segment_pool is not None, "weka now always builds the trie IR" + + store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id=benchmark_id + ) + catalog = await build_unified_trie_store_interned(parsed, store) + + trace_id = parsed.traces[0].id + node_ordinals = catalog[trace_id] + assert node_ordinals, "expected per-node trie ordinals" + first_ordinal = min(node_ordinals.values()) + + client = GraphSegmentUnifiedClient( + base_path=tmp_path, benchmark_id=benchmark_id + ).open() + return trace_id, first_ordinal, client + + +@pytest.mark.asyncio +async def test_trie_envelope_carries_recorded_model(tmp_path): + """With no override, the recorded model stands on the materialized payload.""" + trace_id, ordinal, client = await _build_trie_store(tmp_path, "rec") + try: + payload = materialize_graph_request_unified( + client, trace_id, ordinal, "profiling" + ) + assert payload is not None + assert payload["model"] == _RECORDED_MODEL + finally: + client.close() + + +@pytest.mark.asyncio +async def test_run_level_model_override_reaches_wire(tmp_path): + """A run-level ``--model`` (``endpoint.extra``) clobbers the recorded model. + + The trie carries the recorded model per-node; the worker layers the run-level + override last, so the override reaches the wire and the recorded model never + leaks. This is the trie-path equivalent of the legacy build-time model map. + """ + trace_id, ordinal, client = await _build_trie_store(tmp_path, "ovr") + try: + payload = materialize_graph_request_unified( + client, trace_id, ordinal, "profiling" + ) + assert payload is not None + # Precondition: before the run-level layer the recorded model is present. + assert payload["model"] == _RECORDED_MODEL + + endpoint = EndpointInfo( + type="chat", + streaming=True, + use_server_token_count=False, + extra=[("model", _OVERRIDE_MODEL)], + ) + apply_run_level_payload_options(payload, endpoint) + + assert payload["model"] == _OVERRIDE_MODEL, "run-level --model wins on the wire" + assert payload["model"] != _RECORDED_MODEL, "recorded model must not leak" + finally: + client.close() diff --git a/tests/unit/graph/test_dynamic_pool.py b/tests/unit/graph/test_dynamic_pool.py new file mode 100644 index 0000000000..2b100c7250 --- /dev/null +++ b/tests/unit/graph/test_dynamic_pool.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GraphDynamicPool lifecycle: capture values, deferred eviction, LRU backstop.""" + +from aiperf.graph.dynamic_pool import ( + GraphCapturedReply, + GraphDynamicPool, + GraphPoolSentinel, +) + + +def _reply(text: str) -> GraphCapturedReply: + return GraphCapturedReply(text=text) + + +def test_put_get_roundtrip_and_sentinels() -> None: + pool = GraphDynamicPool(max_bytes=1024) + pool.put("t-1#0", "profiling", 0, _reply("hello")) + pool.put("t-1#0", "profiling", 1, GraphPoolSentinel.FAILED) + pool.put("t-1#0", "profiling", 2, GraphPoolSentinel.EMPTY) + + assert pool.get("t-1#0", "profiling", 0) == _reply("hello") + assert pool.get("t-1#0", "profiling", 1) is GraphPoolSentinel.FAILED + assert pool.get("t-1#0", "profiling", 2) is GraphPoolSentinel.EMPTY + assert pool.get("t-1#0", "profiling", 9) is None + assert pool.get("t-9#0", "profiling", 0) is None + + +def test_phase_variants_are_distinct_entries() -> None: + pool = GraphDynamicPool(max_bytes=1024) + pool.put("t-1#0", "warmup", 0, _reply("w")) + pool.put("t-1#0", "profiling", 0, _reply("p")) + pool.trace_end("t-1#0", "warmup") + + assert pool.get("t-1#0", "warmup", 0) is None + assert pool.get("t-1#0", "profiling", 0) == _reply("p") + + +def test_trace_end_evicts_when_idle() -> None: + pool = GraphDynamicPool(max_bytes=1024) + pool.put("t-1#0", "profiling", 0, _reply("hello")) + pool.trace_end("t-1#0", "profiling") + + assert pool.get("t-1#0", "profiling", 0) is None + assert pool.total_bytes == 0 + + +def test_trace_end_defers_while_inflight() -> None: + """A cancelled credit's FAILED write must land in a live entry (§5.5).""" + pool = GraphDynamicPool(max_bytes=1024) + pool.credit_started("t-1#0", "profiling") + pool.put("t-1#0", "profiling", 0, _reply("first")) + pool.trace_end("t-1#0", "profiling") + + assert pool.get("t-1#0", "profiling", 0) == _reply("first") + pool.put("t-1#0", "profiling", 1, GraphPoolSentinel.FAILED) + pool.credit_finished("t-1#0", "profiling") + + assert pool.get("t-1#0", "profiling", 0) is None + assert pool.get("t-1#0", "profiling", 1) is None + assert pool.total_bytes == 0 + + +def test_trace_end_idempotent_and_unknown_noop() -> None: + pool = GraphDynamicPool(max_bytes=1024) + pool.put("t-1#0", "profiling", 0, _reply("x")) + pool.trace_end("t-1#0", "profiling") + pool.trace_end("t-1#0", "profiling") + pool.trace_end("never", "profiling") + assert pool.total_bytes == 0 + + +def test_overwrite_same_ordinal_accounts_bytes_once() -> None: + pool = GraphDynamicPool(max_bytes=1024) + pool.put("t-1#0", "profiling", 0, _reply("aaaa")) + pool.put("t-1#0", "profiling", 0, _reply("bb")) + assert pool.total_bytes == 2 + + +def test_lru_backstop_evicts_least_recently_used_trace() -> None: + pool = GraphDynamicPool(max_bytes=10) + pool.put("t-1#0", "profiling", 0, _reply("aaaa")) # 4 bytes + pool.put("t-2#0", "profiling", 0, _reply("bbbb")) # 4 bytes + pool.get("t-1#0", "profiling", 0) # touch t-1: t-2 becomes LRU + pool.put("t-3#0", "profiling", 0, _reply("cccc")) # 12 > 10: evict t-2 + + assert pool.get("t-2#0", "profiling", 0) is None + assert pool.get("t-1#0", "profiling", 0) == _reply("aaaa") + assert pool.get("t-3#0", "profiling", 0) == _reply("cccc") + assert pool.total_bytes == 8 + + +def test_lru_evicts_writer_last_when_it_alone_exceeds_cap() -> None: + pool = GraphDynamicPool(max_bytes=4) + pool.put("t-1#0", "profiling", 0, _reply("way more than four bytes")) + assert pool.get("t-1#0", "profiling", 0) is None + assert pool.total_bytes == 0 + + +# --------------------------------------------------------------------------- +# Byte accounting for structured values +# --------------------------------------------------------------------------- + + +def test_text_only_value_costs_text_bytes() -> None: + pool = GraphDynamicPool(max_bytes=1024) + pool.put("t-1#0", "profiling", 0, _reply("abcd")) + assert pool.total_bytes == 4 + + +def test_message_json_value_costs_text_plus_json_bytes() -> None: + pool = GraphDynamicPool(max_bytes=1024) + message_json = '{"role":"assistant","content":null,"tool_calls":[]}' + pool.put( + "t-1#0", + "profiling", + 0, + GraphCapturedReply(text="ab", message_json=message_json), + ) + assert pool.total_bytes == 2 + len(message_json.encode()) + + +def test_sentinel_value_costs_flat_bytes() -> None: + pool = GraphDynamicPool(max_bytes=1024) + pool.put("t-1#0", "profiling", 0, GraphPoolSentinel.FAILED) + assert pool.total_bytes == 8 # _SENTINEL_COST_BYTES + + +def test_message_json_bytes_count_toward_lru_cap() -> None: + pool = GraphDynamicPool(max_bytes=25) + pool.put("t-1#0", "profiling", 0, _reply("aaaa")) # 4 bytes + # 0 text bytes + 22 json bytes: 26 > 25 evicts the LRU entry t-1. + pool.put( + "t-2#0", + "profiling", + 0, + GraphCapturedReply(text="", message_json='{"tool_calls":[1,2,3]}'), + ) + + assert pool.get("t-1#0", "profiling", 0) is None + assert pool.get("t-2#0", "profiling", 0) is not None + assert pool.total_bytes == 22 + + +# --------------------------------------------------------------------------- +# R5 -- no resurrection after trace_end (zero-byte immortal entries) +# --------------------------------------------------------------------------- + + +def test_credit_started_after_trace_end_leaves_no_entry() -> None: + """A straggler credit_started behind GraphTraceEnd must not resurrect the + (trace_id, phase_variant) entry -- resurrected entries have no end message + left to evict them and accumulate forever in recycle-heavy runs.""" + pool = GraphDynamicPool(max_bytes=1024) + pool.put("t-1#0", "profiling", 0, _reply("hello")) + pool.trace_end("t-1#0", "profiling") + assert pool.entry_count == 0 + + pool.credit_started("t-1#0", "profiling") + + assert pool.entry_count == 0 + assert pool.total_bytes == 0 + + +def test_put_after_trace_end_is_dropped() -> None: + """A straggler capture behind GraphTraceEnd is dropped, not resurrected.""" + pool = GraphDynamicPool(max_bytes=1024) + pool.put("t-1#0", "profiling", 0, _reply("hello")) + pool.trace_end("t-1#0", "profiling") + + pool.put("t-1#0", "profiling", 1, _reply("late capture")) + + assert pool.entry_count == 0 + assert pool.total_bytes == 0 + assert pool.get("t-1#0", "profiling", 1) is None + + +def test_deferred_eviction_also_tombstones() -> None: + """The deferred (end_pending) eviction path blocks resurrection too.""" + pool = GraphDynamicPool(max_bytes=1024) + pool.credit_started("t-1#0", "profiling") + pool.trace_end("t-1#0", "profiling") # deferred: credit in flight + pool.credit_finished("t-1#0", "profiling") # drains -> evict + tombstone + assert pool.entry_count == 0 + + pool.credit_started("t-1#0", "profiling") + pool.put("t-1#0", "profiling", 0, _reply("zombie")) + + assert pool.entry_count == 0 + assert pool.total_bytes == 0 + + +def test_trace_end_for_unknown_key_blocks_out_of_order_start() -> None: + """End arriving BEFORE any event for the key still blocks resurrection.""" + pool = GraphDynamicPool(max_bytes=1024) + pool.trace_end("t-9#0", "profiling") + + pool.credit_started("t-9#0", "profiling") + + assert pool.entry_count == 0 + + +def test_tombstone_fifo_is_bounded() -> None: + """Tombstones age out FIFO at the capacity bound (no unbounded growth).""" + pool = GraphDynamicPool(max_bytes=1024, tombstone_capacity=2) + for i in range(4): + pool.put(f"t-{i}", "profiling", 0, _reply("x")) + pool.trace_end(f"t-{i}", "profiling") + + assert len(pool._tombstones) == 2 + # The oldest key aged out: a straggler for it CAN recreate an entry (the + # accepted bounded-memory trade); recent keys stay blocked. + pool.credit_started("t-0", "profiling") + assert pool.entry_count == 1 + pool.credit_started("t-3", "profiling") + assert pool.entry_count == 1 + + +def test_distinct_phase_variant_not_blocked_by_other_variants_end() -> None: + """Ending the warmup variant must not tombstone the profiling variant.""" + pool = GraphDynamicPool(max_bytes=1024) + pool.put("t-1#0", "warmup", 0, _reply("w")) + pool.trace_end("t-1#0", "warmup") + + pool.credit_started("t-1#0", "profiling") + pool.put("t-1#0", "profiling", 0, _reply("p")) + + assert pool.get("t-1#0", "profiling", 0) == _reply("p") + assert pool.entry_count == 1 diff --git a/tests/unit/graph/test_dynamic_slots_e2e.py b/tests/unit/graph/test_dynamic_slots_e2e.py new file mode 100644 index 0000000000..f65a04e815 --- /dev/null +++ b/tests/unit/graph/test_dynamic_slots_e2e.py @@ -0,0 +1,500 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dynamic-content slots end-to-end through the real content plane. + +The marquee assertions of the dynamic-content pool: a native +two-node graph is lowered (`parse_native`), drained into a REAL unified store +(`build_unified_trie_store_interned`), and driven through the REAL worker +credit path (mock-self harness) with a REAL `GraphDynamicPool` — proving +request 2's wire payload contains request 1's actual response text, omission +shapes for FAILED producers on both `{"s"}` and `{"sv"}` slots, and the loud +`pool_missing` error attribution when stickiness is broken. +""" + +import types +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +import orjson +import pytest + +from aiperf.common.enums import CacheBustTarget, CreditPhase +from aiperf.credit.structs import Credit, CreditContext +from aiperf.dataset.graph.parser import parse_native +from aiperf.dataset.graph.segment_ir.store_builder import ( + build_unified_trie_store_interned, +) +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) +from aiperf.graph.dynamic_pool import ( + GraphCapturedReply, + GraphDynamicPool, + GraphPoolSentinel, +) +from aiperf.workers.worker import Worker + +PLANNER_REVIEWER = """ +graph: + nodes: + plan: + prompt: [{role: user, content: "Make a plan."}] + output: plan_out + review: + prompt: + - role: user + content: ["Review this plan: ", "@plan_out"] + output: review_out + edges: + - {source: START, target: plan} + - {source: plan, target: review} + - {source: review, target: END} +traces: + - id: t1 +""" + +ACCUMULATE_CHAIN = """ +graph: + state: + hist: {type: messages, reducer: add_messages} + nodes: + a: + prompt: ["@hist", {role: user, content: "turn one"}] + output: hist + b: + prompt: ["@hist", {role: user, content: "turn two"}] + output: hist + c: + prompt: ["@hist", {role: user, content: "turn three"}] + output: c_out + edges: + - {source: START, target: a} + - {source: a, target: b} + - {source: b, target: c} + - {source: c, target: END} +traces: + - id: t1 + initial_state: + hist: + - {role: system, content: "be brief"} +""" + + +MERGE_CHAIN = """ +graph: + state: + hist: {type: messages, reducer: add_messages} + nodes: + a: + prompt: ["@hist", {role: user, content: "warm up"}] + output: hist + b: + prompt: ["@hist", {role: user, content: "second question"}] + output: b_out + edges: + - {source: START, target: a} + - {source: a, target: b} + - {source: b, target: END} +traces: + - id: t1 + initial_state: + hist: + - {role: user, content: "first question"} +""" + + +async def _built_client( + tmp_path: Path, yaml_text: str +) -> tuple[GraphSegmentUnifiedClient, dict[str, int]]: + p = tmp_path / "workload.yaml" + p.write_text(yaml_text) + parsed = parse_native(p) + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id="bench") + catalog = await build_unified_trie_store_interned(parsed, store) + return GraphSegmentUnifiedClient(tmp_path, "bench").open(), catalog["t1"] + + +def _worker(client: GraphSegmentUnifiedClient) -> MagicMock: + self = MagicMock() + self._graph_store_reader = MagicMock(return_value=client) + self._graph_dynamic_pool = GraphDynamicPool(max_bytes=1024 * 1024) + self.model_endpoint.primary_model_name = "mock-model" + endpoint = self.model_endpoint.endpoint + endpoint.cache_bust = CacheBustTarget.NONE + endpoint.use_legacy_max_tokens = False + endpoint.streaming = False + endpoint.extra = None + endpoint.use_server_token_count = False + self._build_graph_request_info = MagicMock(return_value=MagicMock()) + self._send_inference_result_message = AsyncMock() + self._send_graph_error_record = AsyncMock() + self._set_graph_envelope_missing = MagicMock() + self.task_stats = MagicMock() + self.run.benchmark_id = "bench" + self._process_graph_credit = types.MethodType(Worker._process_graph_credit, self) + self._dispatch_graph_request = types.MethodType( + Worker._dispatch_graph_request, self + ) + self._graph_capture_value = types.MethodType(Worker._graph_capture_value, self) + return self + + +def _respond_with(self: MagicMock, text: str) -> None: + record = MagicMock() + record.error = None + self.inference_client.send_request = AsyncMock(return_value=record) + self.inference_client.endpoint.build_assistant_turn = MagicMock( + return_value=types.SimpleNamespace( + texts=[types.SimpleNamespace(contents=[text])], + raw_messages=None, + ) + ) + + +# Instance id is ``{template}::{nonce}``; the worker strips ``::{nonce}`` back +# to the base template id ("t1") for catalog / store / dynamic-pool keying. +INSTANCE = "t1::inst0" + + +def _ctx(ordinal: int) -> CreditContext: + return CreditContext( + credit=Credit( + id=ordinal + 1, + phase=CreditPhase.PROFILING, + conversation_id="t1", + x_correlation_id="t1::corr", + turn_index=0, + num_turns=1, + issued_at_ns=0, + trace_id=INSTANCE, + node_ordinal=ordinal, + ), + drop_perf_ns=0, + ) + + +def _sent_payload(self: MagicMock) -> dict: + return self._build_graph_request_info.call_args[0][1] + + +@pytest.mark.asyncio +async def test_request_two_contains_response_one(tmp_path: Path) -> None: + client, ordinals = await _built_client(tmp_path, PLANNER_REVIEWER) + self = _worker(client) + + _respond_with(self, "1. write tests 2. ship") + await self._process_graph_credit(_ctx(ordinals["plan"]), "x-req-1", None) + + ctx = _ctx(ordinals["review"]) + await self._process_graph_credit(ctx, "x-req-2", None) + + payload = _sent_payload(self) + assert payload["messages"] == [ + {"role": "user", "content": "Review this plan: 1. write tests 2. ship"} + ] + assert ctx.error is None + + +@pytest.mark.asyncio +async def test_failed_producer_composes_static_prefix_only(tmp_path: Path) -> None: + client, ordinals = await _built_client(tmp_path, PLANNER_REVIEWER) + self = _worker(client) + self._graph_dynamic_pool.put( + INSTANCE, "profiling", ordinals["plan"], GraphPoolSentinel.FAILED + ) + _respond_with(self, "unused") + + await self._process_graph_credit(_ctx(ordinals["review"]), "x-req-2", None) + + assert _sent_payload(self)["messages"] == [ + {"role": "user", "content": "Review this plan: "} + ] + + +@pytest.mark.asyncio +async def test_accumulated_history_chain_splices_in_order(tmp_path: Path) -> None: + client, ordinals = await _built_client(tmp_path, ACCUMULATE_CHAIN) + self = _worker(client) + + _respond_with(self, "answer one") + await self._process_graph_credit(_ctx(ordinals["a"]), "x-req-1", None) + + _respond_with(self, "answer two") + await self._process_graph_credit(_ctx(ordinals["b"]), "x-req-2", None) + b_payload = _sent_payload(self) + # Full alternation: b sees the seed, a's authored turn, a's live reply, + # then its own turn. + assert b_payload["messages"] == [ + {"role": "system", "content": "be brief"}, + {"role": "user", "content": "turn one"}, + {"role": "assistant", "content": "answer one"}, + {"role": "user", "content": "turn two"}, + ] + + _respond_with(self, "unused") + await self._process_graph_credit(_ctx(ordinals["c"]), "x-req-3", None) + c_payload = _sent_payload(self) + assert c_payload["messages"] == [ + {"role": "system", "content": "be brief"}, + {"role": "user", "content": "turn one"}, + {"role": "assistant", "content": "answer one"}, + {"role": "user", "content": "turn two"}, + {"role": "assistant", "content": "answer two"}, + {"role": "user", "content": "turn three"}, + ] + + +@pytest.mark.asyncio +async def test_tool_calls_capture_splices_verbatim_assistant_message( + tmp_path: Path, +) -> None: + """A structured (tool_calls) capture's ``{"s"}`` splice is the verbatim + recorded assistant message -- tool_calls survive into the successor's wire + body, byte-matching the legacy child-seed rendering.""" + client, ordinals = await _built_client(tmp_path, ACCUMULATE_CHAIN) + self = _worker(client) + tool_calls_msg = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + } + pool = self._graph_dynamic_pool + pool.put( + INSTANCE, + "profiling", + ordinals["a"], + GraphCapturedReply(text="", message_json=orjson.dumps(tool_calls_msg).decode()), + ) + pool.put( + INSTANCE, "profiling", ordinals["b"], GraphCapturedReply(text="answer two") + ) + _respond_with(self, "unused") + + await self._process_graph_credit(_ctx(ordinals["c"]), "x-req-3", None) + + assert _sent_payload(self)["messages"] == [ + {"role": "system", "content": "be brief"}, + {"role": "user", "content": "turn one"}, + tool_calls_msg, + {"role": "user", "content": "turn two"}, + {"role": "assistant", "content": "answer two"}, + {"role": "user", "content": "turn three"}, + ] + + +@pytest.mark.asyncio +async def test_composed_slot_concatenates_text_of_structured_capture( + tmp_path: Path, +) -> None: + """A ``{"sv"}`` block part uses the structured capture's ``.text`` (the + reply's replayable text), never its message JSON.""" + client, ordinals = await _built_client(tmp_path, PLANNER_REVIEWER) + self = _worker(client) + plan_msg = { + "role": "assistant", + "content": "1. write tests 2. ship", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + } + self._graph_dynamic_pool.put( + INSTANCE, + "profiling", + ordinals["plan"], + GraphCapturedReply( + text="1. write tests 2. ship", + message_json=orjson.dumps(plan_msg).decode(), + ), + ) + _respond_with(self, "unused") + + await self._process_graph_credit(_ctx(ordinals["review"]), "x-req-2", None) + + assert _sent_payload(self)["messages"] == [ + {"role": "user", "content": "Review this plan: 1. write tests 2. ship"} + ] + + +@pytest.mark.asyncio +async def test_failed_array_slot_omits_assistant_turn(tmp_path: Path) -> None: + client, ordinals = await _built_client(tmp_path, ACCUMULATE_CHAIN) + self = _worker(client) + pool = self._graph_dynamic_pool + pool.put( + INSTANCE, "profiling", ordinals["a"], GraphCapturedReply(text="answer one") + ) + pool.put(INSTANCE, "profiling", ordinals["b"], GraphPoolSentinel.FAILED) + _respond_with(self, "unused") + + await self._process_graph_credit(_ctx(ordinals["c"]), "x-req-3", None) + + # b's reply FAILED -> its assistant turn is omitted, but its authored user + # turn ("turn two") stays visible/unanswered, adjacent to "turn three". + assert _sent_payload(self)["messages"] == [ + {"role": "system", "content": "be brief"}, + {"role": "user", "content": "turn one"}, + {"role": "assistant", "content": "answer one"}, + {"role": "user", "content": "turn two"}, + {"role": "user", "content": "turn three"}, + ] + + +@pytest.mark.asyncio +async def test_merge_consecutive_user_off_by_default(tmp_path: Path) -> None: + """A FAILED producer leaves its user turn unanswered, adjacent to the next + user turn; default policy sends them as-is.""" + client, ordinals = await _built_client(tmp_path, MERGE_CHAIN) + self = _worker(client) + self._graph_dynamic_pool.put( + INSTANCE, "profiling", ordinals["a"], GraphPoolSentinel.FAILED + ) + _respond_with(self, "unused") + + await self._process_graph_credit(_ctx(ordinals["b"]), "x-req-2", None) + + # init "first question", a's delta "warm up" (a's reply omitted, FAILED), + # b's delta "second question" -> three consecutive user turns, as-is. + assert _sent_payload(self)["messages"] == [ + {"role": "user", "content": "first question"}, + {"role": "user", "content": "warm up"}, + {"role": "user", "content": "second question"}, + ] + + +@pytest.mark.asyncio +async def test_merge_consecutive_user_when_enabled(tmp_path: Path, monkeypatch) -> None: + from aiperf.common.environment import Environment + + monkeypatch.setattr( + Environment.GRAPH, "MERGE_CONSECUTIVE_USER", True, raising=False + ) + client, ordinals = await _built_client(tmp_path, MERGE_CHAIN) + self = _worker(client) + self._graph_dynamic_pool.put( + INSTANCE, "profiling", ordinals["a"], GraphPoolSentinel.FAILED + ) + _respond_with(self, "unused") + + await self._process_graph_credit(_ctx(ordinals["b"]), "x-req-2", None) + + assert _sent_payload(self)["messages"] == [ + {"role": "user", "content": "first question\nwarm up\nsecond question"}, + ] + + +@pytest.mark.asyncio +async def test_merge_leaves_assistant_separated_users_untouched( + tmp_path: Path, monkeypatch +) -> None: + """With the knob on, an assistant-separated user pair is NOT merged. Uses a + system-seeded chain so the only user runs are genuine (no init/delta user + boundary), and all replies succeed -> clean alternation stays intact.""" + from aiperf.common.environment import Environment + + monkeypatch.setattr( + Environment.GRAPH, "MERGE_CONSECUTIVE_USER", True, raising=False + ) + client, ordinals = await _built_client(tmp_path, ACCUMULATE_CHAIN) + self = _worker(client) + pool = self._graph_dynamic_pool + pool.put( + INSTANCE, "profiling", ordinals["a"], GraphCapturedReply(text="answer one") + ) + pool.put( + INSTANCE, "profiling", ordinals["b"], GraphCapturedReply(text="answer two") + ) + _respond_with(self, "unused") + + await self._process_graph_credit(_ctx(ordinals["c"]), "x-req-3", None) + + # Every user turn is separated by an assistant reply -> nothing merges. + assert _sent_payload(self)["messages"] == [ + {"role": "system", "content": "be brief"}, + {"role": "user", "content": "turn one"}, + {"role": "assistant", "content": "answer one"}, + {"role": "user", "content": "turn two"}, + {"role": "assistant", "content": "answer two"}, + {"role": "user", "content": "turn three"}, + ] + + +@pytest.mark.asyncio +async def test_missing_pool_value_errors_loudly_without_dispatch( + tmp_path: Path, +) -> None: + client, ordinals = await _built_client(tmp_path, PLANNER_REVIEWER) + self = _worker(client) + _respond_with(self, "unused") + + ctx = _ctx(ordinals["review"]) + await self._process_graph_credit(ctx, "x-req-2", None) + + assert ctx.error is not None + assert ctx.error.startswith(f"aiperf.graph.pool_missing: {INSTANCE}/") + self.inference_client.send_request.assert_not_awaited() + # WK2: the pre-dispatch failure still emits a synthetic error record so the + # RecordsManager completion barrier stays in lockstep with the credit side. + self._send_graph_error_record.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_pool_missing_error_is_a_trace_stop() -> None: + """The adapter sniffs the prefix; the executor never contains it (§6.3).""" + import asyncio + + from aiperf.dataset.graph.models import ( + ChannelRequirement, + ChannelSpec, + GraphRecord, + LlmNode, + ParsedGraph, + StaticEdge, + TraceRecord, + ) + from aiperf.dataset.graph.segment_ir.pool import SegmentPool + from aiperf.graph.credit_dispatch_adapter import GraphStickinessError + from aiperf.graph.executor import TraceExecutor + + class _PoolMissingIssuer: + async def dispatch(self, node, request, ctx, first_token_cb=None): + if request.node_id == "b": + raise GraphStickinessError("aiperf.graph.pool_missing: t1#0/0") + return "" + + graph = GraphRecord( + state={"a_out": ChannelSpec(), "b_out": ChannelSpec()}, + nodes={ + "a": LlmNode(prompt=[{"role": "user", "content": "q"}], output="a_out"), + "b": LlmNode( + prompt=[{"role": "user", "content": "q"}], + output="b_out", + inputs=[ChannelRequirement(channel="a_out", count=1)], + ), + }, + edges=[ + StaticEdge(source="START", target="a"), + StaticEdge(source="a", target="b"), + StaticEdge(source="b", target="END"), + ], + ) + parsed = ParsedGraph( + graph=graph, traces=[TraceRecord(id="t1")], segment_pool=SegmentPool() + ) + executor = TraceExecutor(parsed, credit_issuer=_PoolMissingIssuer()) + with pytest.raises(ExceptionGroup) as excinfo: + async with asyncio.TaskGroup(): + await executor.run(parsed.traces[0]) + # The trace-stop must carry the stickiness error itself -- proof the + # executor propagated (never contained) the pool-missing failure. + assert excinfo.group_contains(GraphStickinessError) diff --git a/tests/unit/graph/test_effective_root_seed.py b/tests/unit/graph/test_effective_root_seed.py new file mode 100644 index 0000000000..5767b797c9 --- /dev/null +++ b/tests/unit/graph/test_effective_root_seed.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from aiperf.common import random_generator as rng +from aiperf.dataset.graph.adapters.shared.content import ( + resolve_effective_root_seed, +) + + +def test_resolve_effective_root_seed_explicit_seed_wins() -> None: + rng.reset() + rng.init(777) + assert resolve_effective_root_seed(1234) == 1234 + + +def test_resolve_effective_root_seed_none_uses_ambient_manager_seed() -> None: + rng.reset() + rng.init(777) + assert resolve_effective_root_seed(None) == 777 + + +def test_resolve_effective_root_seed_no_manager_generates_per_run_seed() -> None: + rng.reset() + first = resolve_effective_root_seed(None) + second = resolve_effective_root_seed(None) + assert isinstance(first, int) + assert isinstance(second, int) + assert first != second + + +def test_resolve_effective_root_seed_unseeded_manager_generates_per_run_seed() -> None: + # Ambient manager exists but was initialized with None (unseeded run): + # each resolution generates a fresh OS-entropy seed, so distinct unseeded + # runs differ while the resolved int keeps ONE run internally consistent. + rng.reset() + rng.init(None) + first = resolve_effective_root_seed(None) + second = resolve_effective_root_seed(None) + assert first != second diff --git a/tests/unit/graph/test_executor_runs_weka.py b/tests/unit/graph/test_executor_runs_weka.py new file mode 100644 index 0000000000..abb53d37c0 --- /dev/null +++ b/tests/unit/graph/test_executor_runs_weka.py @@ -0,0 +1,434 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end smoke test: drive the dataflow `TraceExecutor` over a real weka graph. + +This is the first test on the branch that actually instantiates and runs +`TraceExecutor`. It ingests the weka fixture via `from_weka_trace`, then drives +the executor's real entrypoint (`TraceExecutor.run(trace)`) with a stub +credit issuer that records each `dispatch` call and returns a placeholder string +(downstream weka prompts do not read LLM output, so a placeholder is contractually +fine). Every LLM node must dispatch exactly once and the run must complete. +""" + +import asyncio +from collections import Counter +from pathlib import Path +from typing import Any + +import msgspec +import pytest + +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.adapters.weka.trace_models import WekaTrace +from aiperf.dataset.graph.adapters.weka.trie_build import ( + ReconCallbacks, + build_trie_graph, +) +from aiperf.dataset.graph.models import ( + LlmNode, + TraceRecord, +) +from aiperf.graph.executor import TraceExecutor + +FIX = Path(__file__).parent / "fixtures" / "weka_min.json" + + +class _RecordingCreditIssuer: + """Stub credit issuer that records dispatch calls and returns a placeholder. + + Mirrors the contract that `dispatch/llm.py` relies on: + `await issuer.dispatch(node, request, placement_ctx)`. + Returns a `str` placeholder; downstream weka prompts do not read prior LLM + output, so the value is never inspected by the runtime. + """ + + def __init__(self) -> None: + self.dispatched: list[str] = [] + + async def dispatch( + self, + node: Any, + request: Any, + ctx: Any, + **kwargs: Any, + ) -> str: + self.dispatched.append(request.node_id) + return f"placeholder::{request.node_id}" + + +@pytest.mark.asyncio +async def test_trace_executor_runs_weka_graph_end_to_end(): + # Real-content synthesis is always on; the synthetic fixture has hash_ids + + # block_size > 0, so it synthesizes offline (gpt2 -> builtin tokenizer + # fallback) without reaching a HuggingFace download. + parsed = from_weka_trace(str(FIX)) + assert parsed.traces, "expected at least one trace" + graph = parsed.graph if not parsed.graphs else next(iter(parsed.graphs.values())) + + expected_dispatch_ids = { + nid for nid, n in graph.nodes.items() if isinstance(n, LlmNode) + } + assert expected_dispatch_ids, "expected at least one dispatchable LLM node" + + issuer = _RecordingCreditIssuer() + executor = TraceExecutor(parsed, credit_issuer=issuer) + + async with asyncio.TaskGroup(): + for trace in parsed.traces: + result = await executor.run(trace) + assert result.trace_id == trace.id + + assert len(issuer.dispatched) == len(expected_dispatch_ids), ( + f"expected {len(expected_dispatch_ids)} dispatches " + f"({sorted(expected_dispatch_ids)}), got {sorted(issuer.dispatched)}" + ) + assert set(issuer.dispatched) == expected_dispatch_ids + + +class _OverflowOnceCreditIssuer: + """Stub issuer that raises ``_NodeOverflowTerminate`` on the FIRST dispatch. + + Mirrors what ``CreditDispatchAdapter`` does when the worker returns a + context-overflow error: the first dispatched LLM node terminates the + trajectory early. Records every dispatch attempt so the test can assert the + downstream turns never dispatch. + """ + + def __init__(self) -> None: + self.dispatched: list[str] = [] + + async def dispatch( + self, + node: Any, + request: Any, + ctx: Any, + **kwargs: Any, + ) -> str: + from aiperf.graph.credit_dispatch_adapter import _NodeOverflowTerminate + + self.dispatched.append(request.node_id) + if len(self.dispatched) == 1: + raise _NodeOverflowTerminate("context_length_exceeded") + return f"placeholder::{request.node_id}" + + +@pytest.mark.asyncio +async def test_overflow_terminates_trajectory_early_no_error(): + """When the first LLM node overflows, the executor stops dispatching the + rest of the trace's turns and the run completes WITHOUT raising (the trace is + not an errored trace). + """ + parsed = from_weka_trace(str(FIX)) + graph = parsed.graph if not parsed.graphs else next(iter(parsed.graphs.values())) + dispatchable = {nid for nid, n in graph.nodes.items() if isinstance(n, LlmNode)} + # Sanity: the fixture is a linear chain of >1 dispatchable node so "stop the + # rest" is observable. + assert len(dispatchable) > 1 + + issuer = _OverflowOnceCreditIssuer() + executor = TraceExecutor(parsed, credit_issuer=issuer) + + # The run must NOT raise: the overflow is a clean early-termination, and the + # downstream orphan cascade is swallowed (ctx.overflow_terminated). + async with asyncio.TaskGroup(): + for trace in parsed.traces: + result = await executor.run(trace) + assert result.trace_id == trace.id + + # Exactly ONE dispatch happened (the overflowed node); every downstream turn + # of the trajectory was suppressed. + assert len(issuer.dispatched) == 1, ( + f"overflow must stop dispatching the rest of the trace; " + f"got dispatches {issuer.dispatched}" + ) + assert set(issuer.dispatched) < dispatchable + + +# --- AND-fan-in regression (T8-join) -------------------------------------- + +_BLOCK_SIZE = 64 + + +def _stub_decode_block_tokens(hash_ids: list[int]) -> list[int]: + out: list[int] = [] + for h in hash_ids: + out.extend(range(h * 100, h * 100 + _BLOCK_SIZE)) + return out + + +def _stub_partial_tail_tokens(n_tokens: int, seed: str) -> list[int]: + base = sum(ord(c) for c in seed) * 1000 + return list(range(base, base + n_tokens)) + + +def _stub_decode_tokens_to_text(tokens: list[int]) -> str: + return "|".join(str(t) for t in tokens) + + +_STUB_CALLBACKS = ReconCallbacks( + decode_block_tokens=_stub_decode_block_tokens, + sample_partial_tail_tokens=_stub_partial_tail_tokens, + decode_tokens_to_text=_stub_decode_tokens_to_text, +) + + +def _two_predecessor_trie_graph(): + """Build a trie graph whose resume turn AND-fans-in on MULTIPLE predecessors. + + Parent p0 (t=0) spawns two BLOCKING subagents; each subagent's inner turn + derives from p0. The parent's resume turn p1 (t=4) starts after p0 AND both + subagents finish. Under interval-order timing + (:func:`_build_interval_edges`) p1's predecessors are the MAXIMAL + finished-before frontier: the two subagent last turns. The content-parent p0 + (end 0.5) is transitively covered -- it finished before s1_inner started + (0.5 <= 1.0), so ``p0 -> s1_inner -> p1`` drops p0 from p1's frontier. p1 + therefore declares TWO ``ChannelRequirement`` inputs (one per subagent + ``_out``). This is the multi-predecessor early-fire the cycle guard tripped + on. + + Returns ``(parsed_with_trace, p1_node_id, {predecessor node ids})``. + """ + s1 = { + "t": 1.0, + "type": "subagent", + "agent_id": "agent_1", + "subagent_type": "Explore", + "status": "completed", + "duration_ms": 2000, + "models": ["M"], + "requests": [ + {"t": 1.0, "type": "n", "model": "M", "in": 128, "out": 16, + "hash_ids": [50, 51], "api_time": 1.5}, + ], + } # fmt: skip + s2 = { + "t": 1.0, + "type": "subagent", + "agent_id": "agent_2", + "subagent_type": "Explore", + "status": "completed", + "duration_ms": 2000, + "models": ["M"], + "requests": [ + {"t": 1.2, "type": "n", "model": "M", "in": 128, "out": 16, + "hash_ids": [60, 61], "api_time": 1.5}, + ], + } # fmt: skip + trace = WekaTrace.model_validate( + { + "id": "trie_join", + "models": ["M"], + "block_size": _BLOCK_SIZE, + "hash_id_scope": "local", + "requests": [ + { + "t": 0.0, + "type": "n", + "model": "M", + "in": 128, + "out": 64, + "hash_ids": [1, 2], + "api_time": 0.5, + }, + s1, + s2, + { + "t": 4.0, + "type": "n", + "model": "M", + "in": 192, + "out": 32, + "hash_ids": [1, 2, 3], + "api_time": 1.0, + }, + ], + } + ) + parsed, pool = build_trie_graph(trace, callbacks=_STUB_CALLBACKS) + nodes = {nid: n for nid, n in parsed.graph.nodes.items() if isinstance(n, LlmNode)} + by_offset = {n.arrival_offset_us: nid for nid, n in nodes.items()} + p1 = by_offset[int(4.0 * 1e6)] + # p1's interval-order frontier predecessors are the two subagent last turns; + # the content-parent p0 (offset 0.0) is transitively covered and dropped. + preds = { + by_offset[int(1.0 * 1e6)], + by_offset[int(1.2 * 1e6)], + } + + # Precondition: p1 really has the multi-predecessor AND-join the bug fired + # past (both subagent joins; p0 is frontier-dropped under interval-order). + assert len(nodes[p1].inputs) == 2, nodes[p1].inputs + assert {req.channel for req in nodes[p1].inputs} == {f"{pid}_out" for pid in preds} + + trie_trace = TraceRecord(id=trace.id) + parsed = msgspec.structs.replace(parsed, traces=[trie_trace], segment_pool=pool) + return parsed, p1, preds + + +@pytest.mark.asyncio +async def test_trie_and_fanin_join_node_fires_once_after_both_predecessors(): + """A 2-predecessor trie join runs with NO cycle error and fires the join ONCE. + + Before T8-join the join node declared no input channels, so the executor's + ``await_inputs`` gate was a no-op: the node fired on its FIRST completing + predecessor, finished, then the SECOND predecessor re-scheduled it -> + ``RuntimeError("cycle detected: node ... re-scheduled after completing")``. + + With each node now reading its predecessors' ``{src}_out`` channels, the + join WAITS for BOTH predecessors and fires exactly once. This is the + EXECUTOR-level proof the static-structure tests miss (see + gotcha_graph_adapter_tests_skip_validator). + """ + parsed, join_id, pred_ids = _two_predecessor_trie_graph() + all_ids = {nid for nid, n in parsed.graph.nodes.items() if isinstance(n, LlmNode)} + + issuer = _RecordingCreditIssuer() + executor = TraceExecutor(parsed, credit_issuer=issuer) + + # The run must NOT raise the cycle RuntimeError (the regression). asyncio's + # TaskGroup would wrap it in an ExceptionGroup; reaching the asserts at all + # means no node re-scheduled after completing. + async with asyncio.TaskGroup(): + for trace in parsed.traces: + result = await executor.run(trace) + assert result.trace_id == trace.id + + counts = Counter(issuer.dispatched) + # Every node dispatched exactly once -- the join did not double-fire. + assert set(counts) == all_ids + assert counts[join_id] == 1, ( + f"join node {join_id!r} must fire EXACTLY once; got {counts[join_id]}" + ) + for pid in pred_ids: + assert counts[pid] == 1 + + +@pytest.mark.asyncio +async def test_trie_and_fanin_join_waits_for_both_predecessors(): + """The join node fires only AFTER both predecessors have dispatched. + + Records dispatch ORDER and asserts the join id appears strictly after both + predecessor ids -- the join's two ``count=1`` requirements force it to wait + for both ``_out`` writes (each predecessor's LLM dispatch), not just the + first. + """ + parsed, join_id, pred_ids = _two_predecessor_trie_graph() + + issuer = _RecordingCreditIssuer() + executor = TraceExecutor(parsed, credit_issuer=issuer) + + async with asyncio.TaskGroup(): + for trace in parsed.traces: + await executor.run(trace) + + order = issuer.dispatched + join_pos = order.index(join_id) + for pid in pred_ids: + assert pid in order[:join_pos], ( + f"join {join_id!r} dispatched before predecessor {pid!r}; order={order}" + ) + + +# --- virtual-time replay (Clock abstraction) ------------------------------ + + +class _VTimeIssuer: + """Records each node's VIRTUAL dispatch time, then consumes the node's + recorded ``api_time`` in virtual time. + + ``api_by_id`` maps ``request.node_id`` -> recorded processing seconds. The + executor records the node's finish as ``dispatch_start + api_time`` (the + issuer's virtual sleep), so a successor's firing gate clears at the recorded + end-to-start instant. + """ + + def __init__(self, clock: Any, api_by_id: dict[str, float]) -> None: + self._clock = clock + self._api = api_by_id + self.dispatched_at: dict[str, float] = {} + + async def dispatch( + self, + node: Any, + request: Any, + ctx: Any, + **kwargs: Any, + ) -> str: + nid = request.node_id + self.dispatched_at[nid] = self._clock.now_ns() / 1e9 + api_s = self._api.get(nid, 0.0) + if api_s > 0.0: + await self._clock.sleep_ns(int(api_s * 1e9)) + return f"placeholder::{nid}" + + +async def _drive_virtual(clock: Any, task: Any) -> Any: + """Pump a virtual-clock replay: drain ready callbacks, then fast-forward sim + time to the earliest parked waiter whenever the loop goes idle.""" + loop = asyncio.get_running_loop() + ready = loop._ready # noqa: SLF001 -- idle detection for the pump + while not task.done(): + while ready and not task.done(): + await asyncio.sleep(0) + if task.done(): + break + nxt = clock.peek_min_waiter_ns() + if nxt is None: + await asyncio.sleep(0) + if not ready and clock.peek_min_waiter_ns() is None and not task.done(): + raise RuntimeError("virtual-time replay stalled") + continue + await clock.advance_to(nxt) + return task.result() + + +@pytest.mark.asyncio +async def test_trie_replays_recorded_timeline_on_virtual_time(): + """The REAL ``TraceExecutor`` on a ``VirtualClock`` reproduces the recorded + per-node start times byte-exact, fast. + + The 2-predecessor trie has recorded starts p0=0, s1=1.0, s2=1.2, p1=4.0 with + api_times 0.5/1.5/1.5/1.0 and no idle gaps, so the warped timeline equals the + raw starts. Driven by a virtual clock pumped to each parked waiter, every + node must dispatch at its recorded start -- proving the production firing + loop (input gates, AND-join, edge-gate ``max``) reconstructs the timeline, + not just a hand-rolled sim. + """ + from aiperf.common.clock import VirtualClock + + parsed, _join_id, _pred_ids = _two_predecessor_trie_graph() + nodes = {nid: n for nid, n in parsed.graph.nodes.items() if isinstance(n, LlmNode)} + # recorded start (== raw t, no idle gap > cap) and api_time keyed by the + # node id the executor dispatches (``request.node_id`` == graph node id). + by_offset_us = {n.arrival_offset_us: nid for nid, n in nodes.items()} + expected_start = { + by_offset_us[0]: 0.0, + by_offset_us[int(1.0 * 1e6)]: 1.0, + by_offset_us[int(1.2 * 1e6)]: 1.2, + by_offset_us[int(4.0 * 1e6)]: 4.0, + } + api_by_id = { + by_offset_us[0]: 0.5, + by_offset_us[int(1.0 * 1e6)]: 1.5, + by_offset_us[int(1.2 * 1e6)]: 1.5, + by_offset_us[int(4.0 * 1e6)]: 1.0, + } + + clock = VirtualClock() + issuer = _VTimeIssuer(clock, api_by_id) + executor = TraceExecutor(parsed, credit_issuer=issuer, clock=clock) + + async def _phase() -> None: + async with asyncio.TaskGroup(): + for trace in parsed.traces: + await executor.run(trace) + + phase_task = asyncio.ensure_future(_phase()) + await _drive_virtual(clock, phase_task) + + assert set(issuer.dispatched_at) == set(expected_start) + for nid, want in expected_start.items(): + assert issuer.dispatched_at[nid] == pytest.approx(want, abs=1e-3), ( + f"node {nid!r} dispatched at {issuer.dispatched_at[nid]:.4f}s, " + f"expected recorded start {want:.4f}s" + ) diff --git a/tests/unit/graph/test_executor_watchdog.py b/tests/unit/graph/test_executor_watchdog.py new file mode 100644 index 0000000000..90759ec961 --- /dev/null +++ b/tests/unit/graph/test_executor_watchdog.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""EXECUTOR_WATCHDOG_TIMEOUT converts a firing-loop deadlock into a failure. + +DISPATCH_TIMEOUT only bounds a node that ISSUED a credit; a node wedged on an +unsatisfiable channel input never dispatched and has no per-dispatch guard. The +opt-in executor watchdog wraps `TraceExecutor.run` in `asyncio.wait_for` so such +a wedge raises `TimeoutError` instead of hanging. Default None preserves the +faithful unbounded idle-gap replay, so this test explicitly opts in. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from aiperf.common.environment import Environment +from aiperf.dataset.graph.models import ( + ChannelRequirement, + ChannelSpec, + GraphRecord, + LlmNode, + ParsedGraph, + StaticEdge, + TraceRecord, +) +from aiperf.graph.executor import TraceExecutor + + +async def _run_deadlock(executor: TraceExecutor, trace: TraceRecord) -> BaseException: + """Drive a wedged run and return the exception it surfaced. + + A wedged ``run`` surfaces a body exception directly or wrapped in an + ``ExceptionGroup`` depending on how the internal cancellation unwinds, so + the caller inspects the returned exception rather than relying on the exact + top-level type. + """ + try: + async with asyncio.TaskGroup(): + await executor.run(trace) + except BaseException as exc: # noqa: BLE001 - test inspects the surfaced type + return exc + raise AssertionError("run completed without raising; expected a deadlock") + + +def _has_timeout(exc: BaseException) -> bool: + """True when ``exc`` is (or wraps, via ExceptionGroup) a TimeoutError.""" + if isinstance(exc, TimeoutError): + return True + nested = getattr(exc, "exceptions", None) + if nested is not None: + return any(_has_timeout(sub) for sub in nested) + return False + + +def _mutual_deadlock_graph() -> ParsedGraph: + """Two entry LLMs each AND-fan-in gated on the OTHER's output channel. + + A awaits ``b`` (produced only by B) and B awaits ``a`` (produced only by A), + so neither ever becomes input-ready: a genuine liveness deadlock the orphan + check cannot break (each channel has a live, never-completing producer). + """ + nodes: dict[str, object] = { + "A": LlmNode( + prompt=["@a"], output="a", inputs=[ChannelRequirement(channel="b", count=1)] + ), + "B": LlmNode( + prompt=["@b"], output="b", inputs=[ChannelRequirement(channel="a", count=1)] + ), + } + edges = [ + StaticEdge(source="START", target="A"), + StaticEdge(source="START", target="B"), + ] + graph = GraphRecord( + nodes=nodes, + edges=edges, + state={"a": ChannelSpec(), "b": ChannelSpec()}, + ) + return ParsedGraph(graph=graph, traces=[TraceRecord(id="t")]) + + +@pytest.mark.asyncio +async def test_watchdog_times_out_a_wedged_firing_loop(monkeypatch): + """With the watchdog set, run() self-bounds the deadlock with a TimeoutError.""" + monkeypatch.setattr( + Environment.GRAPH, "EXECUTOR_WATCHDOG_TIMEOUT", 0.2, raising=False + ) + parsed = _mutual_deadlock_graph() + executor = TraceExecutor(parsed) + exc = await _run_deadlock(executor, parsed.traces[0]) + assert _has_timeout(exc) + + +@pytest.mark.asyncio +async def test_no_watchdog_by_default_does_not_self_bound(monkeypatch): + """With the default (None) run() does NOT self-bound; only an external + wait_for ends the deadlock (proving the watchdog, not some other guard, is + what produced the timeout in the test above).""" + monkeypatch.setattr( + Environment.GRAPH, "EXECUTOR_WATCHDOG_TIMEOUT", None, raising=False + ) + parsed = _mutual_deadlock_graph() + executor = TraceExecutor(parsed) + + async def _drive() -> None: + async with asyncio.TaskGroup(): + await executor.run(parsed.traces[0]) + + # run() itself never returns; only our external wait_for breaks the wedge. + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(_drive(), timeout=0.3) diff --git a/tests/unit/graph/test_first_token_anchor_lowering.py b/tests/unit/graph/test_first_token_anchor_lowering.py new file mode 100644 index 0000000000..653568cb76 --- /dev/null +++ b/tests/unit/graph/test_first_token_anchor_lowering.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Post-TTFT overlaps lower to first-token-anchored edges. + +``apply_start_anchors`` already collapses an overlapped node's edges to one +start-anchored ``StaticEdge`` (``delay_after_predecessor_start_us``). This +module locks the post-TTFT REFINEMENT: when the recorded child starts at/after +the streaming parent's recorded first token (``node.raw_start - parent.raw_start +>= parent.request.ttft``), the same edge ALSO carries +``delay_after_predecessor_first_token_us == delay_after_predecessor_start_us - +ttft*1e6``. A pre-TTFT child (start before the recorded first token) or a +non-streaming parent (``ttft is None``) keeps the pure dispatch anchor. + +The lowering is exercised at three levels: hand-built ``TrieNode``s through +``apply_start_anchors`` (mirroring ``test_start_anchor_edges.py``), the weka +adapter end-to-end via ``build_trie_graph`` (defeating the +adapter-tests-skip-validator trap), and the ``chop_trie_at_tstar`` carrier +that must round-trip the third field None-preservingly. +""" + +from __future__ import annotations + +from collections import defaultdict + +import msgspec +import pytest + +from aiperf.dataset.graph.adapters.weka.trace_models import WekaTrace +from aiperf.dataset.graph.adapters.weka.trie_build import ( + ReconCallbacks, + build_trie_graph, +) +from aiperf.dataset.graph.models import ( + GraphRecord, + LlmNode, + ParsedGraph, + StaticEdge, + TraceRecord, +) +from aiperf.dataset.graph.segment_ir.interval_order import ( + apply_start_anchors, + build_interval_edges, + compute_ranks, +) +from aiperf.dataset.graph.segment_ir.trie_content import ( + TrieNode, + TrieRequest, +) +from aiperf.dataset.graph.validator import ValidationSeverity, validate +from aiperf.timing.snapshot_chop import chop_trie_at_tstar + +_BLOCK_SIZE = 64 + + +def _node(nid, t, api, causal=None, ttft=None): + n = TrieNode( + node_id=nid, + request=TrieRequest( + hash_ids=[], + input_length=64, + output_length=8, + t=t, + api_time=api, + ttft=ttft, + ), + order=0, + causal_parent_id=causal, + ) + n.warped_start = t + return n + + +# --- hand-built TrieNode edges -------------------------------------------- + + +def test_post_ttft_overlap_gets_both_fields(): + p = _node("p", 0.0, 8.0, ttft=2.0) # streaming parent + c = _node("c", 4.0, 1.0, causal="p") # 4.0 >= 0.0 + 2.0: post-TTFT + nodes = [p, c] + compute_ranks(nodes) + edges = build_interval_edges(nodes) + apply_start_anchors(nodes, edges) + (e,) = edges["c"] + assert e.source == "p" + assert e.delay_after_predecessor_start_us == pytest.approx(4.0e6) + assert e.delay_after_predecessor_first_token_us == pytest.approx(2.0e6) + + +def test_pre_ttft_overlap_keeps_pure_dispatch_anchor(): + p = _node("p", 0.0, 8.0, ttft=2.0) + c = _node("c", 1.0, 1.0, causal="p") # 1.0 < ttft: pre-TTFT + nodes = [p, c] + compute_ranks(nodes) + edges = build_interval_edges(nodes) + apply_start_anchors(nodes, edges) + (e,) = edges["c"] + assert e.delay_after_predecessor_start_us == pytest.approx(1.0e6) + assert e.delay_after_predecessor_first_token_us is None + + +def test_no_ttft_parent_keeps_pure_dispatch_anchor(): + p = _node("p", 0.0, 8.0, ttft=None) # non-streaming (n-type) parent + c = _node("c", 4.0, 1.0, causal="p") # overlaps, but parent has no ttft + nodes = [p, c] + compute_ranks(nodes) + edges = build_interval_edges(nodes) + apply_start_anchors(nodes, edges) + (e,) = edges["c"] + assert e.delay_after_predecessor_start_us == pytest.approx(4.0e6) + assert e.delay_after_predecessor_first_token_us is None + + +# --- weka end-to-end ------------------------------------------------------- + + +def _stub_decode_block_tokens(hash_ids: list[int]) -> list[int]: + out: list[int] = [] + for h in hash_ids: + out.extend(range(h * 100, h * 100 + _BLOCK_SIZE)) + return out + + +def _stub_partial_tail_tokens(n_tokens: int, seed: str) -> list[int]: + base = sum(ord(c) for c in seed) * 1000 + return list(range(base, base + n_tokens)) + + +def _stub_decode_tokens_to_text(tokens: list[int]) -> str: + return "|".join(str(t) for t in tokens) + + +_STUB_CALLBACKS = ReconCallbacks( + decode_block_tokens=_stub_decode_block_tokens, + sample_partial_tail_tokens=_stub_partial_tail_tokens, + decode_tokens_to_text=_stub_decode_tokens_to_text, +) + +# ttft_anchor:0: streaming (type "s") parent, ttft 2.0, api 8.0, stops on +# tool_use. Completed subagent whose first inner request a1:0 starts at t=4.0 -- +# inside ttft_anchor:0's [0, 8) interval and at/after its recorded first token +# (t >= ttft). +_TTFT_TRACE = { + "id": "ttft_anchor", "models": ["M"], "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "s", "model": "M", "in": 128, "out": 64, + "hash_ids": [1, 2], "api_time": 8.0, "ttft": 2.0, "stop": "tool_use"}, + {"t": 2.0, "type": "subagent", "agent_id": "a1", + "subagent_type": "Explore", "status": "completed", "models": ["M"], + "requests": [ + {"t": 4.0, "type": "n", "model": "M", "in": 128, "out": 16, + "hash_ids": [50, 51], "api_time": 1.0}, + ]}, + ], +} # fmt: skip + + +def test_weka_streaming_parent_stamps_ttft_end_to_end(): + parsed, pool = build_trie_graph( + WekaTrace.model_validate(_TTFT_TRACE), callbacks=_STUB_CALLBACKS + ) + incoming = defaultdict(list) + for e in parsed.graph.edges: + incoming[e.target].append(e) + (edge,) = incoming["a1:0"] + assert edge.source == "ttft_anchor:0" + assert edge.delay_after_predecessor_start_us == pytest.approx(4.0e6) + assert edge.delay_after_predecessor_first_token_us == pytest.approx(2.0e6) + # Validator must accept the graph (adapter-tests-skip-validator trap): the + # first-token anchor rides its dispatch fallback on a non-START source. + validated = msgspec.structs.replace( + parsed, traces=[TraceRecord(id="ttft_anchor")], segment_pool=pool + ) + blocking = [ + i for i in validate(validated) if i.severity is ValidationSeverity.ERROR + ] + assert blocking == [], blocking + + +# --- snapshot chop carrier --------------------------------------------------- + + +def test_chop_round_trips_first_token_anchor_on_surviving_edge(): + """``chop_trie_at_tstar`` keeps a surviving first-token-anchored edge verbatim. + + The t* chop is the snapshot carrier that rebuilds the edge set; both anchor + fields (dispatch fallback + first-token refinement) must round-trip + None-preservingly through the REAL chop path (t*>0, both endpoints kept). + """ + nodes = { + "p": LlmNode(prompt=["hi"], output="p_out", arrival_offset_us=2_000_000), + "c": LlmNode(prompt=["hi"], output="c_out", arrival_offset_us=6_000_000), + } + edges = [ + StaticEdge(source="START", target="p"), + StaticEdge( + source="p", + target="c", + delay_after_predecessor_start_us=4.0e6, + delay_after_predecessor_first_token_us=2.0e6, + ), + ] + parsed = ParsedGraph( + graph=GraphRecord(nodes=nodes, edges=edges, state={}), + traces=[TraceRecord(id="t")], + ) + + chopped = chop_trie_at_tstar(parsed, t_star_us=1_000_000) + + assert chopped is not parsed, "t*>0 must run the real chop path" + (edge,) = [e for e in chopped.graph.edges if e.target == "c"] + assert edge.source == "p" + assert edge.delay_after_predecessor_start_us == pytest.approx(4.0e6) + assert edge.delay_after_predecessor_first_token_us == pytest.approx(2.0e6) + assert edge.delay_after_predecessor_us is None diff --git a/tests/unit/graph/test_first_token_fidelity_tool.py b/tests/unit/graph/test_first_token_fidelity_tool.py new file mode 100644 index 0000000000..8c8e5f3a46 --- /dev/null +++ b/tests/unit/graph/test_first_token_fidelity_tool.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The offline fidelity proof validates POST-TTFT first-token-anchored nodes as +``first_token + D'`` offsets, falling back to ``dispatch + D`` (with a LOUD +report line) when the parent's observed TTFT is not recoverable from the export. + +The weka trie builder stamps a post-TTFT overlap node's single +``StaticEdge`` with BOTH ``delay_after_predecessor_start_us`` (D) and +``delay_after_predecessor_first_token_us`` (D' = D - ttft*1e6). This locks the +matching behavior in ``tools/weka_trace_fidelity.py``: + +* :func:`build_recorded_trace` must widen ``_RecordedNode.start_anchor`` to the + triple ``(source, start_delay_us, first_token_delay_us | None)`` -- pre-TTFT + and no-ttft-parent overlaps carry ``None`` in the third slot. +* :func:`causality_timing_vs_real_trace` must, for a post-TTFT node, (a) compare + the child's observed dispatch to ``parent_first_token + D'`` when the parent's + observed first token is recoverable from the export (``responses[0].perf_ns - + start_perf_ns`` added to the parent's dispatch wall clock); (b) else fall back + to ``parent_dispatch + D`` AND emit a LOUD ``FALLBACK`` note -- never silently + skipping the edge. + +Geometry is byte-identical to ``_TTFT_TRACE`` in +``tests/unit/graph/test_first_token_runtime.py``: streaming P (ttft_anchor:0) ttft +2.0 api 8.0; PRE-TTFT child (a1:0) at t=1.0 (pure dispatch anchor); POST-TTFT child +(a2:0) at t=4.0 (D=4.0, D'=2.0); end-anchored tail (ttft_anchor:1) at t=9.0. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +from tools.weka_trace_fidelity import ( + build_recorded_trace, + causality_timing_vs_real_trace, +) +from tools.weka_trie_timing_sim import simulate_trace + +# Byte-identical to ``test_first_token_runtime._TTFT_TRACE``. +_TTFT_TRACE = { + "id": "ttft_anchor", "models": ["M"], "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "s", "ttft": 2.0, "api_time": 8.0, "in": 128, "out": 64, + "hash_ids": [1, 2], "stop": "tool_use", "model": "M"}, + {"t": 0.5, "type": "subagent", "agent_id": "a1", + "subagent_type": "Explore", "status": "completed", "models": ["M"], + "requests": [ + {"t": 1.0, "type": "n", "model": "M", "in": 128, "out": 16, + "hash_ids": [50, 51], "api_time": 1.0}, + ]}, + {"t": 3.5, "type": "subagent", "agent_id": "a2", + "subagent_type": "Explore", "status": "completed", "models": ["M"], + "requests": [ + {"t": 4.0, "type": "n", "model": "M", "in": 128, "out": 16, + "hash_ids": [60, 61], "api_time": 1.0}, + ]}, + {"t": 9.0, "type": "n", "model": "M", "in": 256, "out": 32, + "hash_ids": [1, 2, 3, 4], "api_time": 0.5}, + ], +} # fmt: skip + +_P = "ttft_anchor:0" +_C_PRE = "a1:0" +_C_POST = "a2:0" +_TAIL = "ttft_anchor:1" + +_S = 1_000_000_000 # 1e9 ns per second +_ORIGIN_NS = 1 * _S # arbitrary fresh run-origin (absolute wall-clock differs) + + +def _write_trace(tmp_path: Path) -> Path: + """Materialize ``_TTFT_TRACE`` to a JSON file the tool can rebuild from.""" + path = tmp_path / "ttft_anchor.json" + path.write_text(json.dumps(_TTFT_TRACE)) + return path + + +def _record( + node_id: str, + request_start_ns: int, + *, + start_perf_ns: int | None = None, + responses: list[dict] | None = None, +) -> dict: + """One raw-export JSONL line for a profiling dispatch of ``node_id``. + + ``start_perf_ns`` + ``responses`` (each with a ``perf_ns``) let the tool + recover a streaming parent's OBSERVED first token; omit both for a + zero-latency-style export where the observed TTFT is unrecoverable. + """ + return { + "metadata": { + "conversation_id": "ttft_anchor#0", + "x_request_id": f"{node_id}::deadbeefdeadbeefdeadbeefdeadbeef", + "benchmark_phase": "profiling", + "request_start_ns": request_start_ns, + "credit_issued_ns": None, + }, + "start_perf_ns": start_perf_ns, + "responses": responses or [], + "payload": {"messages": []}, + } + + +def _write_raw(tmp_path: Path, records: list[dict]) -> Path: + """Write a raw-export JSONL from pre-built record dicts.""" + path = tmp_path / "profile_export_raw.jsonl" + path.write_text("\n".join(json.dumps(r) for r in records) + "\n") + return path + + +def test_build_recorded_trace_captures_first_token_delay(tmp_path: Path) -> None: + """Post-TTFT edges carry D' as the third ``start_anchor`` slot; pre-TTFT None. + + The overlap children each carry one ``delay_after_predecessor_start_us`` edge + off ``ttft_anchor:0`` (streaming, ttft 2.0). ``a1:0`` started PRE-TTFT (t=1.0 < + 2.0) so its edge has no first-token delay; ``a2:0`` started POST-TTFT (t=4.0 >= 2.0) + so its edge carries D' = D - ttft*1e6 = 4.0e6 - 2.0e6 = 2.0e6. + """ + trace_file = _write_trace(tmp_path) + recorded = build_recorded_trace(trace_file, idle_gap_cap_seconds=60.0) + + assert recorded.nodes[_C_PRE].start_anchor == (_P, 1.0e6, None) + assert recorded.nodes[_C_POST].start_anchor == (_P, 4.0e6, 2.0e6) + assert recorded.nodes[_C_PRE].predecessors == [] + assert recorded.nodes[_C_POST].predecessors == [] + + +def test_fallback_replay_passes_with_loud_note(tmp_path: Path) -> None: + """A zero-latency export (no recoverable parent TTFT) times the post-TTFT child + at ``parent_dispatch + D`` and emits a LOUD FALLBACK note; the faithful placement + at ``parent_start + 4.0s`` PASSES. + """ + trace_file = _write_trace(tmp_path) + raw = _write_raw( + tmp_path, + [ + _record(_P, _ORIGIN_NS), + _record(_C_PRE, _ORIGIN_NS + int(1.0 * _S)), + _record(_C_POST, _ORIGIN_NS + int(4.0 * _S)), + ], + ) + + report = causality_timing_vs_real_trace(raw, trace_file, idle_gap_cap_seconds=60.0) + + assert report.passed, report.render() + # Both start-anchored children validate exactly (pre-TTFT + post-TTFT fallback). + assert report.exact_edges >= 2 + # The post-TTFT edge fell back to dispatch + D and said so, LOUDLY. + assert any("FALLBACK" in n and _C_POST in n for n in report.notes), report.render() + + +def test_shifted_first_token_child_fails(tmp_path: Path) -> None: + """Shifting the post-TTFT child's dispatch +1.0s (to 5.0s after P) FAILS naming + it -- the proof actually checks first-token timing rather than skipping the node. + + +1.0s exceeds the abs tolerance (0.75s) and the relative tolerance + (0.15 * 4.0s = 0.6s), so the fallback comparison must flag it. + """ + trace_file = _write_trace(tmp_path) + raw = _write_raw( + tmp_path, + [ + _record(_P, _ORIGIN_NS), + _record(_C_PRE, _ORIGIN_NS + int(1.0 * _S)), + _record(_C_POST, _ORIGIN_NS + int(5.0 * _S)), + ], + ) + + report = causality_timing_vs_real_trace(raw, trace_file, idle_gap_cap_seconds=60.0) + + assert not report.passed + assert any(_C_POST in m.where for m in report.mismatches), report.render() + + +def test_recoverable_first_token_uses_first_token_anchor(tmp_path: Path) -> None: + """When the parent's observed first token IS recoverable, the proof anchors the + post-TTFT child on ``first_token + D'`` -- NOT the dispatch fallback. + + P streams its first token at an OBSERVED 3.0s (differs from the recorded 2.0), + so first_token_wall = P_dispatch + 3.0 and the expected child dispatch is + 3.0 + D'(2.0) = 5.0s after P. Placing the child there PASSES; placing it at the + dispatch-fallback position (P_dispatch + D = 4.0s) FAILS, proving the observed + first token -- not the recorded gap -- drives the comparison. + """ + trace_file = _write_trace(tmp_path) + p_rec = _record( + _P, _ORIGIN_NS, start_perf_ns=0, responses=[{"perf_ns": int(3.0 * _S)}] + ) + + faithful = _write_raw( + tmp_path, + [p_rec, _record(_C_POST, _ORIGIN_NS + int(5.0 * _S))], + ) + report = causality_timing_vs_real_trace( + faithful, trace_file, idle_gap_cap_seconds=60.0 + ) + assert report.passed, report.render() + assert report.exact_edges >= 1 + # First-token anchor was used, so no fallback note fired for the child. + assert not any("FALLBACK" in n and _C_POST in n for n in report.notes) + + # Twin at the dispatch-fallback position (P + D = 4.0s) must FAIL, because the + # tool anchored on the observed first token (expected 5.0s), not dispatch + D. + raw2 = tmp_path / "profile_export_raw.jsonl" + raw2.write_text( + "\n".join( + json.dumps(r) for r in [p_rec, _record(_C_POST, _ORIGIN_NS + int(4.0 * _S))] + ) + + "\n" + ) + report2 = causality_timing_vs_real_trace( + raw2, trace_file, idle_gap_cap_seconds=60.0 + ) + assert not report2.passed + assert any(_C_POST in m.where for m in report2.mismatches), report2.render() + + +def test_timing_sim_tool_counts_first_token_edge(tmp_path: Path) -> None: + """The sibling ``weka_trie_timing_sim`` simulator gates first-token edges off the + predecessor's DISPATCH (observed ttft == recorded ttft in a pure replay, so + ``first_token + D' == dispatch + D``) and counts them under a distinct label. + + ``_TTFT_TRACE`` has exactly one post-TTFT overlap (a2:0), so the simulator + reconstructs the recorded warped timeline byte-exact for all four nodes and + reports exactly one first-token edge. + """ + trace_file = _write_trace(tmp_path) + + checked, exact, diverged, first_token_edges = simulate_trace( + trace_file, cap=60.0, tol=1e-3 + ) + + assert diverged == [] + assert checked == exact == 4 + assert first_token_edges == 1 diff --git a/tests/unit/graph/test_first_token_routing.py b/tests/unit/graph/test_first_token_routing.py new file mode 100644 index 0000000000..7598403d10 --- /dev/null +++ b/tests/unit/graph/test_first_token_routing.py @@ -0,0 +1,368 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Task 4 — post-TTFT first-token routing: strategy -> adapter -> dispatch cb. + +Locks the delivery chain wired in Task 4: + +* ``first_token_sources(graph)`` derives the set of node ids that SOURCE a + first-token-anchored ``StaticEdge`` (``delay_after_predecessor_first_token_us`` + set), from the same per-trace graph the adapter dispatches. +* ``CreditDispatchAdapter`` stamps ``TurnToSend.first_token_event`` for those + source nodes, parks an optional per-dispatch ``first_token_cb`` under the + waiter key, and fires it AT MOST ONCE from ``on_first_token`` (unknown / None + keys are graceful no-ops). +* ``GraphIRReplayStrategy._on_graph_first_token`` de-muxes each ``FirstToken`` to + the owning adapter by ``trace_id`` (mirroring ``_on_graph_return``), tolerating + an unknown / None trace id. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from pathlib import Path + +import pytest + +from aiperf.credit.messages import FirstToken +from aiperf.dataset.graph.adapters.weka.trace_models import WekaTrace +from aiperf.dataset.graph.adapters.weka.trie_build import ( + ReconCallbacks, + build_trie_graph, +) +from aiperf.dataset.graph.graph_path_catalog import CatalogContext +from aiperf.dataset.graph.models import GraphRecord, LlmNode, StaticEdge +from aiperf.graph.credit_dispatch_adapter import CreditDispatchAdapter +from aiperf.graph.placement import DispatchRequest, PlacementContext +from aiperf.timing.strategies.graph_ir_replay import ( + GraphIRReplayStrategy, + first_token_sources, +) + +_BLOCK_SIZE = 64 +_WEKA_MIN = Path(__file__).parent / "fixtures" / "weka_min.json" + + +# --------------------------------------------------------------------------- +# first_token_sources(graph) — the pure derivation +# --------------------------------------------------------------------------- + + +def _stub_decode_block_tokens(hash_ids: list[int]) -> list[int]: + out: list[int] = [] + for h in hash_ids: + out.extend(range(h * 100, h * 100 + _BLOCK_SIZE)) + return out + + +def _stub_partial_tail_tokens(n_tokens: int, seed: str) -> list[int]: + base = sum(ord(c) for c in seed) * 1000 + return list(range(base, base + n_tokens)) + + +def _stub_decode_tokens_to_text(tokens: list[int]) -> str: + return "|".join(str(t) for t in tokens) + + +_STUB_CALLBACKS = ReconCallbacks( + decode_block_tokens=_stub_decode_block_tokens, + sample_partial_tail_tokens=_stub_partial_tail_tokens, + decode_tokens_to_text=_stub_decode_tokens_to_text, +) + +# ttft_anchor:0: streaming parent (ttft 2.0) whose completed subagent's first inner +# request starts at t=4.0 -- inside the parent's [0, 8) interval and at/after its +# first token, so the lowered edge carries delay_after_predecessor_first_token_us. +_TTFT_TRACE = { + "id": "ttft_anchor", "models": ["M"], "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "s", "model": "M", "in": 128, "out": 64, + "hash_ids": [1, 2], "api_time": 8.0, "ttft": 2.0, "stop": "tool_use"}, + {"t": 2.0, "type": "subagent", "agent_id": "a1", + "subagent_type": "Explore", "status": "completed", "models": ["M"], + "requests": [ + {"t": 4.0, "type": "n", "model": "M", "in": 128, "out": 16, + "hash_ids": [50, 51], "api_time": 1.0}, + ]}, + ], +} # fmt: skip + + +def test_first_token_sources_computed_from_parsed_graph(): + """The post-TTFT weka trace lowers exactly one first-token-anchored edge, + sourced at ``ttft_anchor:0`` -- so the derivation returns ``{"ttft_anchor:0"}``.""" + parsed, _pool = build_trie_graph( + WekaTrace.model_validate(_TTFT_TRACE), callbacks=_STUB_CALLBACKS + ) + assert first_token_sources(parsed.graph) == frozenset({"ttft_anchor:0"}) + + +def test_first_token_sources_excludes_plain_and_start_only_edges(): + """Only edges carrying ``delay_after_predecessor_first_token_us`` count; a + plain delay edge and a start-anchored-only edge are excluded.""" + nodes = { + "A": LlmNode(prompt=["@a"], output="a"), + "B": LlmNode(prompt=["@b"], output="b"), + "C": LlmNode(prompt=["@c"], output="c"), + "D": LlmNode(prompt=["@d"], output="d"), + } + edges = [ + StaticEdge(source="START", target="A"), + StaticEdge(source="A", target="B", delay_after_predecessor_us=3.0e6), + StaticEdge(source="B", target="C", delay_after_predecessor_start_us=4.0e6), + StaticEdge( + source="C", + target="D", + delay_after_predecessor_start_us=4.0e6, + delay_after_predecessor_first_token_us=2.0e6, + ), + ] + graph = GraphRecord(nodes=nodes, edges=edges, state={}) + assert first_token_sources(graph) == frozenset({"C"}) + + +def test_first_token_sources_empty_for_gap_free_graph(): + graph = GraphRecord( + nodes={"A": LlmNode(prompt=["@a"], output="a")}, + edges=[StaticEdge(source="START", target="A")], + state={}, + ) + assert first_token_sources(graph) == frozenset() + + +# --------------------------------------------------------------------------- +# Adapter — first_token_event stamping + per-dispatch cb routing +# --------------------------------------------------------------------------- + + +class FakeIssuer: + def __init__(self) -> None: + self.sent: list[object] = [] + + async def issue_graph_credit(self, turn: object) -> bool: + self.sent.append(turn) + return True + + +@dataclass +class FakeCredit: + x_correlation_id: str + turn_index: int + trace_id: str + node_ordinal: int + phase_variant: str = "profiling" + + +@dataclass +class FakeLlmNode: + output: str = "out" + + +def _ctx(parent_trace_id: str, node_id: str) -> PlacementContext: + return PlacementContext(parent_trace_id=parent_trace_id, parent_node_id=node_id) + + +def _request(node_id: str) -> DispatchRequest: + return DispatchRequest(node_id=node_id) + + +def _catalog(trace_id: str, node_key_to_ordinal: dict[str, int]) -> CatalogContext: + return CatalogContext( + catalog={trace_id: dict(node_key_to_ordinal)}, + ) + + +def _make_adapter(issuer: FakeIssuer, trace_id: str, ordinals: dict[str, int], **kw): + return CreditDispatchAdapter( + credit_issuer=issuer, + catalog_context=_catalog(trace_id, ordinals), + trace_id=trace_id, + **kw, + ) + + +def _credit_for(turn: object) -> FakeCredit: + return FakeCredit( + x_correlation_id=turn.x_correlation_id, + turn_index=turn.turn_index, + trace_id=turn.trace_id, + node_ordinal=turn.node_ordinal, + phase_variant=turn.phase_variant, + ) + + +@pytest.mark.asyncio +async def test_adapter_registers_and_routes_first_token_cb(): + """A dispatch on a first-token source stamps ``first_token_event`` and parks + the ``first_token_cb``; ``on_first_token`` fires it exactly once, and a repeat + call for the same key after resolve is a graceful no-op.""" + issuer = FakeIssuer() + adapter = _make_adapter( + issuer, "t0", {"t0:0": 0}, first_token_sources=frozenset({"t0:0"}) + ) + fired: list[int] = [] + task = asyncio.create_task( + adapter.dispatch( + FakeLlmNode(), + _request("t0:0"), + _ctx("t0", "t0:0"), + first_token_cb=lambda: fired.append(1), + ) + ) + await asyncio.sleep(0) + + turn = issuer.sent[0] + assert turn.first_token_event is True + + adapter.on_first_token(turn.x_correlation_id, turn.turn_index) + assert fired == [1] + + adapter.resolve(_credit_for(turn), error=None, cancelled=False) + assert isinstance(await task, str) + + # A second first-token for the same key (or after resolve) must not re-fire. + adapter.on_first_token(turn.x_correlation_id, turn.turn_index) + assert fired == [1] + + +@pytest.mark.asyncio +async def test_first_token_event_false_for_non_source_node(): + """A node that sources no first-token-anchored edge issues its turn with + ``first_token_event=False`` (no wasted TTFT emission).""" + issuer = FakeIssuer() + adapter = _make_adapter( + issuer, "t0", {"t0:0": 0}, first_token_sources=frozenset({"other:0"}) + ) + task = asyncio.create_task( + adapter.dispatch(FakeLlmNode(), _request("t0:0"), _ctx("t0", "t0:0")) + ) + await asyncio.sleep(0) + assert issuer.sent[0].first_token_event is False + adapter.resolve(_credit_for(issuer.sent[0]), error=None, cancelled=False) + assert isinstance(await task, str) + + +@pytest.mark.asyncio +async def test_on_first_token_unknown_key_is_noop(): + """An ``on_first_token`` for a key that parked no callback (or None fields) + is a graceful no-op.""" + issuer = FakeIssuer() + adapter = _make_adapter(issuer, "t0", {"t0:0": 0}) + adapter.on_first_token("nope", 99) # never registered + adapter.on_first_token(None, None) # None fast-path fields + # Nothing raised, no cb to fire. + + +@pytest.mark.asyncio +async def test_first_token_cb_dropped_when_dispatch_resolves_without_ttft(): + """A resolve that arrives with no preceding TTFT drops the parked cb; a late + ``on_first_token`` afterwards must not fire it.""" + issuer = FakeIssuer() + adapter = _make_adapter( + issuer, "t0", {"t0:0": 0}, first_token_sources=frozenset({"t0:0"}) + ) + fired: list[int] = [] + task = asyncio.create_task( + adapter.dispatch( + FakeLlmNode(), + _request("t0:0"), + _ctx("t0", "t0:0"), + first_token_cb=lambda: fired.append(1), + ) + ) + await asyncio.sleep(0) + turn = issuer.sent[0] + adapter.resolve(_credit_for(turn), error=None, cancelled=False) + assert isinstance(await task, str) + adapter.on_first_token(turn.x_correlation_id, turn.turn_index) + assert fired == [] + + +# --------------------------------------------------------------------------- +# Strategy — _on_graph_first_token de-mux by trace_id +# --------------------------------------------------------------------------- + + +class _FakeAdapter: + def __init__(self) -> None: + self.calls: list[tuple[str | None, int | None]] = [] + + def on_first_token( + self, x_correlation_id: str | None, turn_index: int | None + ) -> None: + self.calls.append((x_correlation_id, turn_index)) + + +def _min_strategy() -> GraphIRReplayStrategy: + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + + parsed = from_weka_trace(str(_WEKA_MIN)) + return GraphIRReplayStrategy( + credit_issuer=object(), + parsed_graph=parsed, + register_observer=lambda obs: None, + ) + + +def _ft(trace_id, x_correlation_id="x", turn_index=0) -> FirstToken: + from aiperf.common.enums import CreditPhase + + return FirstToken( + credit_id=1, + phase=CreditPhase.PROFILING, + ttft_ns=5, + trace_id=trace_id, + x_correlation_id=x_correlation_id, + turn_index=turn_index, + ) + + +def test_strategy_demux_routes_by_trace_id(): + """A FirstToken reaches ONLY the adapter registered under its trace_id.""" + strategy = _min_strategy() + a, b = _FakeAdapter(), _FakeAdapter() + strategy._adapters = {"A": a, "B": b} + + strategy._on_graph_first_token(_ft("A", x_correlation_id="xa", turn_index=3)) + assert a.calls == [("xa", 3)] + assert b.calls == [] + + +def test_strategy_demux_unknown_trace_id_is_noop(): + strategy = _min_strategy() + a = _FakeAdapter() + strategy._adapters = {"A": a} + strategy._on_graph_first_token(_ft("ZZZ", x_correlation_id="xz", turn_index=1)) + assert a.calls == [] + + +def test_strategy_demux_none_trace_id_is_noop(): + strategy = _min_strategy() + a = _FakeAdapter() + strategy._adapters = {"A": a} + strategy._on_graph_first_token(_ft(None)) + assert a.calls == [] + + +@pytest.mark.asyncio +async def test_setup_phase_installs_first_token_observer_when_wired(): + """When a first-token registrar is supplied, setup_phase installs the de-mux + observer; teardown detaches it with None.""" + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + + parsed = from_weka_trace(str(_WEKA_MIN)) + installed_return: list[object] = [] + installed_ft: list[object] = [] + strategy = GraphIRReplayStrategy( + credit_issuer=object(), + parsed_graph=parsed, + register_observer=installed_return.append, + register_first_token_observer=installed_ft.append, + ) + await strategy.setup_phase() + assert len(installed_ft) == 1 + assert callable(installed_ft[0]) + + await strategy.teardown_phase() + assert installed_ft[-1] is None diff --git a/tests/unit/graph/test_first_token_runtime.py b/tests/unit/graph/test_first_token_runtime.py new file mode 100644 index 0000000000..3c8878eed7 --- /dev/null +++ b/tests/unit/graph/test_first_token_runtime.py @@ -0,0 +1,443 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Runtime proof: post-TTFT first-token-anchored edges gate a successor on the +predecessor's OBSERVED first token, falling back to the dispatch anchor when the +predecessor terminates without streaming one. + +Drives the REAL ``TraceExecutor`` on a ``VirtualClock`` over a weka trie graph +whose overlap geometry produces a ``StaticEdge`` carrying BOTH +``delay_after_predecessor_start_us`` (D) and +``delay_after_predecessor_first_token_us`` (D') (Task 2 lowering). The harness +(``_STUB_CALLBACKS``, ``_drive_virtual``) mirrors ``test_start_anchor_runtime.py``; +``_TTFTIssuer`` extends ``_VTimeIssuer`` so a node with a recorded ttft sleeps +``ttft`` virtual seconds, invokes the ``first_token_cb`` it received as a dispatch +kwarg, then sleeps the remainder. + +Contract: +1. Recorded-speed replay reproduces the recorded starts (observed == fallback). +2. Inflating the parent's ttft moves the first-token-anchored child by exactly + the inflation while the pre-TTFT dispatch-anchored child stays put. +3. A parent that errors before its first token gates the child at the dispatch + fallback (``_finalize_node`` latches the event; no wall entry). +4. ``AIPERF_GRAPH_IGNORE_EDGE_DELAYS`` short-circuits the wait entirely. +5. A parent that resolves WITHOUT streaming a first token falls back; a late / + duplicate stamp is a no-op (wall not overwritten). +""" + +import asyncio +from typing import Any + +import msgspec +import pytest + +from aiperf.common.clock import VirtualClock +from aiperf.common.environment import Environment +from aiperf.dataset.graph.adapters.weka.trace_models import WekaTrace +from aiperf.dataset.graph.adapters.weka.trie_build import ( + ReconCallbacks, + build_trie_graph, +) +from aiperf.dataset.graph.models import TraceRecord +from aiperf.graph.credit_dispatch_adapter import GraphDispatchError +from aiperf.graph.executor import TraceExecutor + +_BLOCK_SIZE = 64 + + +def _stub_decode_block_tokens(hash_ids: list[int]) -> list[int]: + out: list[int] = [] + for h in hash_ids: + out.extend(range(h * 100, h * 100 + _BLOCK_SIZE)) + return out + + +def _stub_partial_tail_tokens(n_tokens: int, seed: str) -> list[int]: + base = sum(ord(c) for c in seed) * 1000 + return list(range(base, base + n_tokens)) + + +def _stub_decode_tokens_to_text(tokens: list[int]) -> str: + return "|".join(str(t) for t in tokens) + + +_STUB_CALLBACKS = ReconCallbacks( + decode_block_tokens=_stub_decode_block_tokens, + sample_partial_tail_tokens=_stub_partial_tail_tokens, + decode_tokens_to_text=_stub_decode_tokens_to_text, +) + + +class _VTimeIssuer: + """Records each node's VIRTUAL dispatch time, then consumes the node's + recorded ``api_time`` in virtual time (mirrors ``test_start_anchor_runtime``). + """ + + def __init__(self, clock: Any, api_by_id: dict[str, float]) -> None: + self._clock = clock + self._api = api_by_id + self.dispatched_at: dict[str, float] = {} + self.dispatched: list[str] = [] + + async def dispatch( + self, + node: Any, + request: Any, + ctx: Any, + first_token_cb: Any = None, + **kwargs: Any, + ) -> str: + nid = request.node_id + self.dispatched.append(nid) + self.dispatched_at[nid] = self._clock.now_ns() / 1e9 + api_s = self._api.get(nid, 0.0) + if api_s > 0.0: + await self._clock.sleep_ns(int(api_s * 1e9)) + return f"placeholder::{nid}" + + +class _TTFTIssuer(_VTimeIssuer): + """A node with a recorded ttft streams a first token at ``ttft`` virtual + seconds (invoking the received ``first_token_cb``), then finishes the + remaining ``api - ttft`` virtual seconds. Nodes without a recorded ttft + behave exactly like ``_VTimeIssuer`` (sleep ``api``, never stamp).""" + + def __init__( + self, clock: Any, api_by_id: dict[str, float], ttft_by_id: dict[str, float] + ) -> None: + super().__init__(clock, api_by_id) + self._ttft = ttft_by_id + + async def dispatch( + self, + node: Any, + request: Any, + ctx: Any, + first_token_cb: Any = None, + **kwargs: Any, + ) -> str: + nid = request.node_id + self.dispatched.append(nid) + self.dispatched_at[nid] = self._clock.now_ns() / 1e9 + ttft = self._ttft.get(nid) + api = self._api.get(nid, 0.0) + if ttft is not None: + await self._clock.sleep_ns(int(ttft * 1e9)) + if first_token_cb is not None: + first_token_cb() + await self._clock.sleep_ns(int(max(0.0, api - ttft) * 1e9)) + elif api > 0.0: + await self._clock.sleep_ns(int(api * 1e9)) + return f"placeholder::{nid}" + + +class _ErrorBeforeFirstTokenIssuer(_TTFTIssuer): + """The named node raises ``GraphDispatchError`` after ``error_after`` virtual + seconds WITHOUT streaming a first token (the trie failure-sentinel path); all + other nodes behave like ``_TTFTIssuer``.""" + + def __init__( + self, + clock: Any, + api_by_id: dict[str, float], + ttft_by_id: dict[str, float], + error_node: str, + error_after: float, + ) -> None: + super().__init__(clock, api_by_id, ttft_by_id) + self._error_node = error_node + self._error_after = error_after + + async def dispatch( + self, + node: Any, + request: Any, + ctx: Any, + first_token_cb: Any = None, + **kwargs: Any, + ) -> str: + nid = request.node_id + if nid == self._error_node: + self.dispatched.append(nid) + self.dispatched_at[nid] = self._clock.now_ns() / 1e9 + await self._clock.sleep_ns(int(self._error_after * 1e9)) + raise GraphDispatchError(f"simulated dispatch failure for {nid!r}") + return await super().dispatch(node, request, ctx, first_token_cb, **kwargs) + + +class _DoubleStampTTFTIssuer(_TTFTIssuer): + """Invokes ``first_token_cb`` TWICE at the first-token instant for the named + node -- the second call must be a graceful no-op (guard in the stamp).""" + + def __init__( + self, + clock: Any, + api_by_id: dict[str, float], + ttft_by_id: dict[str, float], + double_node: str, + ) -> None: + super().__init__(clock, api_by_id, ttft_by_id) + self._double_node = double_node + + async def dispatch( + self, + node: Any, + request: Any, + ctx: Any, + first_token_cb: Any = None, + **kwargs: Any, + ) -> str: + nid = request.node_id + ttft = self._ttft.get(nid) + if nid == self._double_node and ttft is not None: + self.dispatched.append(nid) + self.dispatched_at[nid] = self._clock.now_ns() / 1e9 + api = self._api.get(nid, 0.0) + await self._clock.sleep_ns(int(ttft * 1e9)) + if first_token_cb is not None: + first_token_cb() + first_token_cb() # late duplicate: must be a no-op + await self._clock.sleep_ns(int(max(0.0, api - ttft) * 1e9)) + return f"placeholder::{nid}" + return await super().dispatch(node, request, ctx, first_token_cb, **kwargs) + + +async def _drive_virtual(clock: Any, task: Any) -> Any: + """Pump a virtual-clock replay: drain ready callbacks, then fast-forward sim + time to the earliest parked waiter whenever the loop goes idle.""" + loop = asyncio.get_running_loop() + ready = loop._ready # noqa: SLF001 -- idle detection for the pump + while not task.done(): + while ready and not task.done(): + await asyncio.sleep(0) + if task.done(): + break + nxt = clock.peek_min_waiter_ns() + if nxt is None: + await asyncio.sleep(0) + if not ready and clock.peek_min_waiter_ns() is None and not task.done(): + raise RuntimeError("virtual-time replay stalled") + continue + await clock.advance_to(nxt) + return task.result() + + +# --- fixture -------------------------------------------------------------- + +# P (ttft_anchor:0): streaming, t=0, ttft=2.0, api=8.0 (spawner). +# C_pre (a1:0): subagent inner request at t=1.0 -- PRE-TTFT (1.0 < 2.0), so a +# PURE dispatch anchor (start=1.0e6, ft=None). +# C_post (a2:0): subagent inner request at t=4.0 -- POST-TTFT (4.0 >= 2.0), so a +# first-token-refined anchor (start=4.0e6, ft=D-ttft=2.0e6). +# tail (ttft_anchor:1): end-anchored at t=9.0. +_TTFT_TRACE = { + "id": "ttft_anchor", "models": ["M"], "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "s", "ttft": 2.0, "api_time": 8.0, "in": 128, "out": 64, + "hash_ids": [1, 2], "stop": "tool_use", "model": "M"}, + {"t": 0.5, "type": "subagent", "agent_id": "a1", + "subagent_type": "Explore", "status": "completed", "models": ["M"], + "requests": [ + {"t": 1.0, "type": "n", "model": "M", "in": 128, "out": 16, + "hash_ids": [50, 51], "api_time": 1.0}, + ]}, + {"t": 3.5, "type": "subagent", "agent_id": "a2", + "subagent_type": "Explore", "status": "completed", "models": ["M"], + "requests": [ + {"t": 4.0, "type": "n", "model": "M", "in": 128, "out": 16, + "hash_ids": [60, 61], "api_time": 1.0}, + ]}, + {"t": 9.0, "type": "n", "model": "M", "in": 256, "out": 32, + "hash_ids": [1, 2, 3, 4], "api_time": 0.5}, + ], +} # fmt: skip + +# Node ids the trie build assigns to _TTFT_TRACE (verified against the graph): +# top-level leaves scope to the trace id (skipping subagent markers), subagent +# inner leaves scope to their recorded agent_id. +_P = "ttft_anchor:0" +_C_PRE = "a1:0" +_C_POST = "a2:0" +_TAIL = "ttft_anchor:1" + + +def _build(raw: dict) -> Any: + """Build a trie graph from a raw weka trace dict and wire a bare trace.""" + trace = WekaTrace.model_validate(raw) + parsed, pool = build_trie_graph(trace, callbacks=_STUB_CALLBACKS) + bare = TraceRecord(id=trace.id) + return msgspec.structs.replace(parsed, traces=[bare], segment_pool=pool) + + +async def _run_virtual(parsed: Any, issuer: Any) -> Any: + """Drive ``parsed``'s traces on ``issuer``'s VirtualClock and return it.""" + clock = issuer._clock + executor = TraceExecutor(parsed, credit_issuer=issuer, clock=clock) + + async def _phase() -> None: + async with asyncio.TaskGroup(): + for trace in parsed.traces: + await executor.run(trace) + + phase_task = asyncio.ensure_future(_phase()) + await _drive_virtual(clock, phase_task) + return issuer + + +# --- tests ---------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_recorded_speed_reproduces_recorded_starts(): + """At recorded speed observed==fallback: C_post fires at 4.0 (first token at + 2.0 + D'=2.0), C_pre at 1.0 (pure dispatch anchor), tail at 9.0.""" + parsed = _build(_TTFT_TRACE) + clock = VirtualClock() + issuer = _TTFTIssuer( + clock, + api_by_id={_P: 8.0, _C_PRE: 1.0, _C_POST: 1.0, _TAIL: 0.5}, + ttft_by_id={_P: 2.0}, + ) + + await _run_virtual(parsed, issuer) + + assert set(issuer.dispatched_at) == {_P, _C_PRE, _C_POST, _TAIL} + assert issuer.dispatched_at[_P] == pytest.approx(0.0, abs=1e-3) + assert issuer.dispatched_at[_C_PRE] == pytest.approx(1.0, abs=1e-3) + assert issuer.dispatched_at[_C_POST] == pytest.approx(4.0, abs=1e-3) + assert issuer.dispatched_at[_TAIL] == pytest.approx(9.0, abs=1e-3) + + +@pytest.mark.asyncio +async def test_inflated_ttft_moves_first_token_child_by_the_inflation(): + """Inflating P's ttft 2.0 -> 5.0 (api 11.0) moves the first-token-anchored + C_post to 5.0 + D'=2.0 = 7.0 (exactly the +3.0 inflation) while the pre-TTFT + dispatch-anchored C_pre stays at 1.0. No cycle RuntimeError.""" + parsed = _build(_TTFT_TRACE) + clock = VirtualClock() + issuer = _TTFTIssuer( + clock, + api_by_id={_P: 11.0, _C_PRE: 1.0, _C_POST: 1.0, _TAIL: 0.5}, + ttft_by_id={_P: 5.0}, + ) + + await _run_virtual(parsed, issuer) + + assert issuer.dispatched_at[_C_PRE] == pytest.approx(1.0, abs=1e-3) + assert issuer.dispatched_at[_C_POST] == pytest.approx(7.0, abs=1e-3) + + +@pytest.mark.asyncio +async def test_parent_errors_before_first_token_falls_back_to_dispatch(): + """P errors after 0.5s without streaming a first token. ``_finalize_node`` + latches the event with NO wall entry, so C_post gates at the dispatch + fallback P_dispatch(0.0) + D(4.0) = 4.0.""" + parsed = _build(_TTFT_TRACE) + clock = VirtualClock() + issuer = _ErrorBeforeFirstTokenIssuer( + clock, + api_by_id={_P: 8.0, _C_PRE: 1.0, _C_POST: 1.0, _TAIL: 0.5}, + ttft_by_id={_P: 2.0}, + error_node=_P, + error_after=0.5, + ) + + await _run_virtual(parsed, issuer) + + assert issuer.dispatched_at[_C_POST] == pytest.approx(4.0, abs=1e-3) + + +@pytest.mark.asyncio +async def test_ignore_edge_delays_skips_the_first_token_wait(monkeypatch): + """``AIPERF_GRAPH_IGNORE_EDGE_DELAYS`` short-circuits the first-token wait + entirely: every node dispatches at its scheduling instant, exactly once, in + causal order (P before both children).""" + monkeypatch.setattr(Environment.GRAPH, "IGNORE_EDGE_DELAYS", True, raising=False) + + parsed = _build(_TTFT_TRACE) + clock = VirtualClock() + issuer = _TTFTIssuer( + clock, + api_by_id={_P: 8.0, _C_PRE: 1.0, _C_POST: 1.0, _TAIL: 0.5}, + ttft_by_id={_P: 2.0}, + ) + + await _run_virtual(parsed, issuer) + + from collections import Counter + + counts = Counter(issuer.dispatched) + assert set(counts) == {_P, _C_PRE, _C_POST, _TAIL} + assert all(n == 1 for n in counts.values()), ( + f"every node must dispatch exactly once; got {counts}" + ) + assert issuer.dispatched_at[_P] == pytest.approx(0.0, abs=1e-3) + assert issuer.dispatched_at[_C_PRE] == pytest.approx(0.0, abs=1e-3) + assert issuer.dispatched_at[_C_POST] == pytest.approx(0.0, abs=1e-3) + order = issuer.dispatched + assert order.index(_P) < order.index(_C_PRE) + assert order.index(_P) < order.index(_C_POST) + + +@pytest.mark.asyncio +async def test_parent_resolves_without_first_token_falls_back(): + """A streaming parent whose issuer never streams a first token (no ttft) + resolves normally at api=3.0; ``_finalize_node`` latches C_post's event and + the gate falls back to P_dispatch(0.0) + D(4.0) = 4.0.""" + parsed = _build(_TTFT_TRACE) + clock = VirtualClock() + issuer = _TTFTIssuer( + clock, + api_by_id={_P: 3.0, _C_PRE: 1.0, _C_POST: 1.0, _TAIL: 0.5}, + ttft_by_id={}, # P never streams a first token + ) + + await _run_virtual(parsed, issuer) + + assert issuer.dispatched_at[_C_PRE] == pytest.approx(1.0, abs=1e-3) + assert issuer.dispatched_at[_C_POST] == pytest.approx(4.0, abs=1e-3) + + +@pytest.mark.asyncio +async def test_duplicate_stamp_during_dispatch_is_noop(): + """Invoking ``first_token_cb`` twice at the first-token instant does not + raise and does not shift C_post (recorded-speed observed anchor = 4.0).""" + parsed = _build(_TTFT_TRACE) + clock = VirtualClock() + issuer = _DoubleStampTTFTIssuer( + clock, + api_by_id={_P: 8.0, _C_PRE: 1.0, _C_POST: 1.0, _TAIL: 0.5}, + ttft_by_id={_P: 2.0}, + double_node=_P, + ) + + await _run_virtual(parsed, issuer) + + assert issuer.dispatched_at[_C_POST] == pytest.approx(4.0, abs=1e-3) + + +def test_first_token_stamp_is_idempotent_and_latches(): + """Direct unit test of the stamp closure: the first invocation records the + wall and sets the latch; a second invocation is a no-op -- it neither + re-reads the clock nor overwrites the recorded wall.""" + from aiperf.graph.dispatch.llm import _make_first_token_stamp + + class _FakeCtx: + def __init__(self) -> None: + self.node_first_token_wall_us: dict[str, float] = {} + self._events: dict[str, asyncio.Event] = {} + + def first_token_event(self, node_id: str) -> asyncio.Event: + return self._events.setdefault(node_id, asyncio.Event()) + + ctx = _FakeCtx() + walls = iter([100.0, 999.0]) + stamp = _make_first_token_stamp(ctx, "n1", lambda: next(walls)) + + stamp() + assert ctx.node_first_token_wall_us["n1"] == 100.0 + assert ctx.first_token_event("n1").is_set() + + stamp() # duplicate: guarded before the clock read, wall must not change + assert ctx.node_first_token_wall_us["n1"] == 100.0 + assert next(walls) == 999.0 # proves the 2nd clock read was NOT consumed diff --git a/tests/unit/graph/test_graph_endpoint_guard.py b/tests/unit/graph/test_graph_endpoint_guard.py new file mode 100644 index 0000000000..1f7380dccc --- /dev/null +++ b/tests/unit/graph/test_graph_endpoint_guard.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Every graph workload rejects non-chat endpoint types up front. + +The graph dispatch path materializes a chat-completions body +(``{"messages": [...], "max_completion_tokens": N, "stream": bool}``) and sends +it verbatim, bypassing ``format_payload``. Pointing that body at a non-chat +endpoint (``completions`` expects ``prompt``; ``embeddings`` expects ``input``) +makes every request 422 with no actionable up-front error. +``validate_graph_endpoint_type`` guards the workload at configure time (the +DatasetManager runs it before the graph store is built), raising +``GraphEndpointUnsupportedError`` -- the guard is format-generic and serves +every graph workload (weka / dynamo / native / dag_jsonl). Chat-compatibility +is keyed on the endpoint metadata's ``endpoint_path`` ending in +``/chat/completions`` so any future chat-completions endpoint passes without an +allowlist edit. +""" + +from __future__ import annotations + +import pytest + +from aiperf.config.flags.cli_config import CLIConfig +from aiperf.dataset.graph.workload_detect import ( + GraphEndpointUnsupportedError, + validate_graph_endpoint_type, +) +from tests.unit.conftest import make_run_from_cli + + +def _run(endpoint_type: str): + cfg = CLIConfig(model_names=["test-model"], endpoint_type=endpoint_type) + return make_run_from_cli(cfg) + + +def test_chat_endpoint_passes_guard() -> None: + """The chat endpoint (the supported shape) passes the guard silently.""" + validate_graph_endpoint_type(_run("chat")) # no raise + + +@pytest.mark.parametrize( + "endpoint_type", + ["completions", "embeddings", "nim_embeddings"], +) +def test_non_chat_endpoint_rejected_with_actionable_error(endpoint_type: str) -> None: + """Non-chat endpoint types raise a clear configure-time error. + + The message names the offending ``--endpoint-type`` and points at the chat + endpoint -- a validator-gate style up-front rejection, not a per-request 422. + """ + with pytest.raises(GraphEndpointUnsupportedError) as exc_info: + validate_graph_endpoint_type(_run(endpoint_type)) + message = str(exc_info.value) + assert endpoint_type in message + assert "chat" in message.lower() diff --git a/tests/unit/graph/test_graph_instance_stop_paths.py b/tests/unit/graph/test_graph_instance_stop_paths.py new file mode 100644 index 0000000000..cea14ea28c --- /dev/null +++ b/tests/unit/graph/test_graph_instance_stop_paths.py @@ -0,0 +1,231 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TC5 / TC7 / R3 — instance stop-path hygiene on ``GraphIRReplayStrategy``. + +* TC5: a ``CreditIssueRefusedError`` unwinding an instance is a HEALTHY stop + (request-count / duration gate closed): no ``errored_traces`` bump, no + "unwound with error" warning; genuine errors keep the error path. +* TC7: the duration-timeout branch awaiting the cancelled lane-runner must NOT + swallow an EXTERNAL cancellation of the strategy task itself. +* R3: ``--burst-phase-starts`` collapses the t*-relative leading offsets the + snapshot chop stamps on START in-edges (the node-level field is never + stamped by the trie producers), via the shared + ``aiperf.graph.scheduler.collapse_leading_start_offsets``. +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path +from typing import Any + +import pytest + +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.graph.analysis import trace_duration_us +from aiperf.graph.credit_dispatch_adapter import ( + CreditIssueRefusedError, + GraphDispatchError, +) +from aiperf.timing.strategies.graph_ir_replay import ( + GraphIRReplayStrategy, + _leaf_credit_refusal, +) + +_FIX = Path(__file__).parent / "fixtures" / "weka_min.json" + + +class _Config: + """Minimal per-phase config stub (mirrors test_lane_fanout_recycle).""" + + timing_mode = None + phase = None + + def __init__(self, **caps: Any) -> None: + self.concurrency = caps.get("concurrency") + self.expected_num_sessions = caps.get("expected_num_sessions") + self.total_expected_requests = caps.get("total_expected_requests") + self.expected_duration_sec = caps.get("expected_duration_sec") + + +def _strategy( + parsed: Any, + issuer: Any, + *, + ratio: float = 0.0, + burst: bool = False, + lifecycle: Any = None, +) -> GraphIRReplayStrategy: + return GraphIRReplayStrategy( + config=_Config(), + credit_issuer=issuer, + parsed_graph=parsed, + register_observer=lambda _obs: None, + lifecycle=lifecycle, + start_min_ratio=ratio, + start_max_ratio=ratio, + burst_phase_starts=burst, + ) + + +class _SinkIssuer: + """Base issuer stub satisfying the phase-completion surface.""" + + def mark_graph_sending_complete(self) -> None: ... + + def graph_all_returned(self) -> bool: + return True + + def set_graph_all_returned_event(self) -> None: ... + + +# --------------------------------------------------------------------------- +# TC5 — issuer refusal is a clean stop +# --------------------------------------------------------------------------- + + +class _RefusingIssuer(_SinkIssuer): + """Refuses every graph credit (the request-count / duration gate closed).""" + + async def issue_graph_credit(self, turn: Any) -> bool: + return False + + +class _ExplodingIssuer(_SinkIssuer): + """Raises a genuine (non-refusal) error on every issue.""" + + async def issue_graph_credit(self, turn: Any) -> bool: + raise RuntimeError("issuer exploded") + + +@pytest.mark.asyncio +async def test_issuer_refusal_is_clean_stop_not_error(): + """A refused instance completes without bumping ``errored_traces``.""" + parsed = from_weka_trace(str(_FIX)) + strategy = _strategy(parsed, _RefusingIssuer()) + + await strategy.execute_phase() + + assert strategy.errored_traces == 0 + assert strategy.completed_traces == strategy.admitted_traces == 1 + + +@pytest.mark.asyncio +async def test_genuine_error_still_counts_errored_trace(): + """Control: a non-refusal instance error keeps the existing error path.""" + parsed = from_weka_trace(str(_FIX)) + strategy = _strategy(parsed, _ExplodingIssuer()) + + await strategy.execute_phase() + + assert strategy.errored_traces == 1 + assert strategy.completed_traces == 1 + + +def test_leaf_credit_refusal_unwraps_groups_but_not_mixed(): + refusal = CreditIssueRefusedError("gate closed") + assert _leaf_credit_refusal(refusal) is refusal + + grouped = ExceptionGroup("trace", [refusal]) + assert _leaf_credit_refusal(grouped) is refusal + + nested = ExceptionGroup("outer", [ExceptionGroup("inner", [refusal])]) + assert isinstance(_leaf_credit_refusal(nested), CreditIssueRefusedError) + + mixed = ExceptionGroup("trace", [refusal, GraphDispatchError("real failure")]) + assert _leaf_credit_refusal(mixed) is None + + assert _leaf_credit_refusal(RuntimeError("other")) is None + + +# --------------------------------------------------------------------------- +# TC7 — external cancel not swallowed by the duration-timeout unwind +# --------------------------------------------------------------------------- + + +class _Lifecycle: + def time_left_in_seconds(self) -> float: + return 0.5 + + +@pytest.mark.asyncio +async def test_external_cancel_reraises_during_duration_timeout_unwind(monkeypatch): + """An external cancel landing on the ``await dispatch`` unwind re-raises. + + Simulates the reviewed race directly: the strategy's duration ``wait_for`` + reports a timeout while the lane runner is still unwinding, and the RUNNER + then cancels the strategy task. Pre-fix, ``suppress(CancelledError)`` + swallowed that external cancel and the coroutine returned normally. + """ + parsed = from_weka_trace(str(_FIX)) + strategy = _strategy(parsed, _SinkIssuer(), lifecycle=_Lifecycle()) + + hang = asyncio.Event() + + async def _lanes_stub(traces: list[Any]) -> None: + await hang.wait() + + strategy._run_lanes = _lanes_stub + + async def _fake_wait_for(fut: Any, timeout: float | None = None) -> None: + # Duration budget "elapses" while the lane runner keeps running, so + # the TimeoutError branch's ``await dispatch`` genuinely suspends. + raise TimeoutError + + monkeypatch.setattr(asyncio, "wait_for", _fake_wait_for) + + task = asyncio.get_running_loop().create_task( + strategy._run_traces_under_duration_budget([]) + ) + await asyncio.sleep(0) # park the task on ``await dispatch`` + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert task.cancelled() + + hang.set() # release the orphaned lane-runner task + await asyncio.sleep(0) + + +# --------------------------------------------------------------------------- +# R3 — burst collapses START-edge leading offsets on chopped trie graphs +# --------------------------------------------------------------------------- + + +def _tstar_ratio_between(parsed: Any, a: str, b: str) -> float: + arr = {nid: n.arrival_offset_us for nid, n in parsed.graph.nodes.items()} + duration = trace_duration_us(parsed, parsed.traces[0]) + return ((arr[a] + arr[b]) / 2.0) / duration + + +def test_burst_phase_starts_zeroes_start_edge_leading_offsets(): + """Burst zeroes the chop's START-edge ``min_start_delay_us``; spread keeps it. + + The trie producers stamp leading offsets on START in-edges only (the + node-level ``min_start_delay_us`` is never stamped), so the pre-fix + node-only collapse made ``--burst-phase-starts`` a no-op on trie graphs. + Inter-turn (non-START) edge delays must survive in both modes. + """ + parsed = from_weka_trace(str(_FIX)) + trace = parsed.traces[0] + ratio = _tstar_ratio_between(parsed, "trace_03_n3:0", "trace_03_n3:1") + + spread = _strategy(parsed, _SinkIssuer(), ratio=ratio, burst=False) + burst = _strategy(parsed, _SinkIssuer(), ratio=ratio, burst=True) + + spread_graph, _ = spread._graph_at_t_star(trace, spread._plans[trace.id]) + burst_graph, _ = burst._graph_at_t_star(trace, burst._plans[trace.id]) + + spread_starts = [e for e in spread_graph.graph.edges if e.source == "START"] + burst_starts = [e for e in burst_graph.graph.edges if e.source == "START"] + assert spread_starts and burst_starts + + # Spread keeps the t*-relative leading offset the chop stamped... + assert any(e.min_start_delay_us for e in spread_starts) + # ...burst collapses every leading offset to fire at phase-time 0. + assert all(not e.min_start_delay_us for e in burst_starts) + + # Recorded inter-turn pacing (non-START completion edges) is untouched. + burst_inner = [e for e in burst_graph.graph.edges if e.source != "START"] + assert burst_inner + assert any(e.delay_after_predecessor_us for e in burst_inner) diff --git a/tests/unit/graph/test_graph_parse_context.py b/tests/unit/graph/test_graph_parse_context.py new file mode 100644 index 0000000000..80c5369341 --- /dev/null +++ b/tests/unit/graph/test_graph_parse_context.py @@ -0,0 +1,378 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GraphParseContext: sentinel semantics + per-adapter ctx forwarding rules. + +Pins the adversarially-reviewed contract of the adapter protocol's ``ctx`` +parameter: + +* forward-only-when-set — a partial ctx never clobbers an adapter entry's + non-None defaults (dynamo ``prompt_corpus="coding"``, dag + ``run_streaming=True``) with ``None``; +* tri-state ``idle_gap_cap_seconds`` — explicit ``None`` (user's + ``synthesis.idle_gap_cap_seconds: null``) forwards VERBATIM (disable + warping) while ``UNSET`` forwards nothing (adapter default); +* ``ctx=None`` is byte-equal to today's protocol-default entries (literal + kwarg-dict compares); +* trust/revision publish to the loader-preload env IFF + ``tokenizer_trust_remote_code`` is set, revision passed verbatim. +""" + +from __future__ import annotations + +import copy +import inspect +import pickle +from pathlib import Path +from typing import Any + +import pytest + +import aiperf.dataset._mp_context as mp_context_mod +from aiperf.dataset.graph import parser as parser_mod +from aiperf.dataset.graph.adapters import native as native_mod +from aiperf.dataset.graph.adapters.dag_jsonl import trace as dag_trace +from aiperf.dataset.graph.adapters.dynamo import trace as dynamo_trace +from aiperf.dataset.graph.adapters.weka import trace as weka_trace +from aiperf.dataset.graph.models import GraphRecord, ParsedGraph +from aiperf.dataset.graph.parse_context import ( + UNSET, + GraphParseContext, + publish_ctx_tokenizer_env, +) + +_SpyCalls = list[tuple[tuple[Any, ...], dict[str, Any]]] + +# DynamoTraceAdapter.parse forwarded kwargs with no ctx fields set. +# ``release_replay=True`` is the production store-build opt-in the adapter +# entry always sets (freeing recorded replay hash lists), set at the adapter's +# own entry rather than threaded through the ctx. +_DYNAMO_ENV_DEFAULT_KWARGS: dict[str, Any] = { + "release_replay": True, +} + + +def _spy_entry(monkeypatch: pytest.MonkeyPatch, module: Any, name: str) -> _SpyCalls: + """Replace ``module.name`` with a spy recording (args, kwargs) per call.""" + calls: _SpyCalls = [] + + def spy(*args: Any, **kwargs: Any) -> ParsedGraph: + calls.append((args, kwargs)) + return ParsedGraph(graph=GraphRecord(), traces=[]) + + monkeypatch.setattr(module, name, spy) + return calls + + +def _spy_publish(monkeypatch: pytest.MonkeyPatch) -> _SpyCalls: + """Spy on configure_loader_tokenizer_env (no env mutation).""" + return _spy_entry(monkeypatch, mp_context_mod, "configure_loader_tokenizer_env") + + +def _full_ctx() -> GraphParseContext: + """Every forwardable field set (trust/revision left unset: publish is + covered by its own tests so forwarding compares stay env-free).""" + return GraphParseContext( + content_root_seed=1234, + content_tokenizer="tok", + prompt_corpus="sonnet", + max_osl=77, + idle_gap_cap_seconds=5.0, + default_model="fallback-model", + run_streaming=False, + delay_cap_seconds=2.5, + endpoint_extra=[("k", "v")], + ) + + +# --- UNSET sentinel semantics ------------------------------------------------- + + +def test_unset_pickle_round_trip_preserves_identity() -> None: + assert pickle.loads(pickle.dumps(UNSET)) is UNSET + + +def test_unset_deepcopy_preserves_identity() -> None: + assert copy.deepcopy(UNSET) is UNSET + + +def test_ctx_default_idle_gap_is_unset_sentinel() -> None: + assert GraphParseContext().idle_gap_cap_seconds is UNSET + + +# --- weka forwarding ------------------------------------------------------------ + + +def test_weka_parse_ctx_none_matches_protocol_default_entry( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + calls = _spy_entry(monkeypatch, weka_trace, "from_weka_trace") + p = tmp_path / "t.json" + weka_trace.WekaTraceAdapter.parse(p, ctx=None) + assert calls == [((p,), {})] + + +def test_weka_parse_full_ctx_forwards_weka_fields_only( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + calls = _spy_entry(monkeypatch, weka_trace, "from_weka_trace") + p = tmp_path / "t.json" + weka_trace.WekaTraceAdapter.parse(p, ctx=_full_ctx()) + assert calls == [ + ( + (p,), + { + "content_root_seed": 1234, + "content_tokenizer": "tok", + "prompt_corpus": "sonnet", + "max_osl": 77, + "idle_gap_cap_seconds": 5.0, + }, + ) + ] + + +def test_weka_parse_partial_ctx_forwards_only_set_fields( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + calls = _spy_entry(monkeypatch, weka_trace, "from_weka_trace") + weka_trace.WekaTraceAdapter.parse( + tmp_path / "t.json", ctx=GraphParseContext(content_root_seed=42) + ) + assert calls[0][1] == {"content_root_seed": 42} + + +def test_weka_parse_explicit_none_idle_gap_forwards_none_verbatim( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + # Explicit None = the user's `synthesis.idle_gap_cap_seconds: null` = + # warping DISABLED. Must forward, never collapse into the _USE_DEFAULT default. + calls = _spy_entry(monkeypatch, weka_trace, "from_weka_trace") + weka_trace.WekaTraceAdapter.parse( + tmp_path / "t.json", ctx=GraphParseContext(idle_gap_cap_seconds=None) + ) + assert calls[0][1] == {"idle_gap_cap_seconds": None} + + +def test_weka_parse_unset_idle_gap_forwards_nothing( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + calls = _spy_entry(monkeypatch, weka_trace, "from_weka_trace") + weka_trace.WekaTraceAdapter.parse(tmp_path / "t.json", ctx=GraphParseContext()) + assert "idle_gap_cap_seconds" not in calls[0][1] + assert calls[0][1] == {} + + +# --- dynamo forwarding ---------------------------------------------------------- + + +def test_dynamo_parse_ctx_none_matches_protocol_default_entry( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + calls = _spy_entry(monkeypatch, dynamo_trace, "from_dynamo_trace") + p = tmp_path / "t.jsonl" + dynamo_trace.DynamoTraceAdapter.parse(p, ctx=None) + assert calls == [((p,), dict(_DYNAMO_ENV_DEFAULT_KWARGS))] + + +def test_dynamo_parse_full_ctx_forwards_content_knobs( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + calls = _spy_entry(monkeypatch, dynamo_trace, "from_dynamo_trace") + dynamo_trace.DynamoTraceAdapter.parse(tmp_path / "t.jsonl", ctx=_full_ctx()) + assert calls[0][1] == { + **_DYNAMO_ENV_DEFAULT_KWARGS, + "content_root_seed": 1234, + "content_tokenizer": "tok", + "prompt_corpus": "sonnet", + "idle_gap_cap_seconds": 5.0, + } + + +def test_dynamo_parse_partial_ctx_does_not_clobber_corpus_default( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + # from_dynamo_trace defaults prompt_corpus="coding"; a partial ctx must not + # forward prompt_corpus=None over it. + calls = _spy_entry(monkeypatch, dynamo_trace, "from_dynamo_trace") + dynamo_trace.DynamoTraceAdapter.parse( + tmp_path / "t.jsonl", ctx=GraphParseContext(content_root_seed=42) + ) + assert "prompt_corpus" not in calls[0][1] + assert calls[0][1] == {**_DYNAMO_ENV_DEFAULT_KWARGS, "content_root_seed": 42} + + +def test_dynamo_parse_explicit_none_idle_gap_forwards_none_verbatim( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + # Same tri-state rule as weka: explicit None = the user's + # `synthesis.idle_gap_cap_seconds: null` = warping DISABLED. Must forward, + # never collapse into the shared 60s entry default. + calls = _spy_entry(monkeypatch, dynamo_trace, "from_dynamo_trace") + dynamo_trace.DynamoTraceAdapter.parse( + tmp_path / "t.jsonl", ctx=GraphParseContext(idle_gap_cap_seconds=None) + ) + assert calls[0][1] == {**_DYNAMO_ENV_DEFAULT_KWARGS, "idle_gap_cap_seconds": None} + + +def test_dynamo_parse_unset_idle_gap_forwards_nothing( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + calls = _spy_entry(monkeypatch, dynamo_trace, "from_dynamo_trace") + dynamo_trace.DynamoTraceAdapter.parse(tmp_path / "t.jsonl", ctx=GraphParseContext()) + assert "idle_gap_cap_seconds" not in calls[0][1] + assert calls[0][1] == dict(_DYNAMO_ENV_DEFAULT_KWARGS) + + +# --- dag_jsonl forwarding ------------------------------------------------------- + + +def test_dag_parse_ctx_none_matches_protocol_default_entry( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + calls = _spy_entry(monkeypatch, dag_trace, "from_dag_jsonl") + p = tmp_path / "t.jsonl" + dag_trace.DagJsonlGraphAdapter.parse(p, ctx=None) + assert calls == [((str(p),), {})] + + +def test_dag_parse_full_ctx_forwards_dispatch_knobs( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + calls = _spy_entry(monkeypatch, dag_trace, "from_dag_jsonl") + p = tmp_path / "t.jsonl" + dag_trace.DagJsonlGraphAdapter.parse(p, ctx=_full_ctx()) + assert calls == [ + ( + (str(p),), + { + "default_model": "fallback-model", + "run_streaming": False, + "delay_cap_seconds": 2.5, + "endpoint_extra": [("k", "v")], + }, + ) + ] + + +def test_dag_parse_partial_ctx_does_not_clobber_streaming_default( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + # from_dag_jsonl defaults run_streaming=True; a partial ctx must not + # forward run_streaming=None over it. + calls = _spy_entry(monkeypatch, dag_trace, "from_dag_jsonl") + dag_trace.DagJsonlGraphAdapter.parse( + tmp_path / "t.jsonl", ctx=GraphParseContext(content_root_seed=42) + ) + assert calls[0][1] == {} + + +# --- native --------------------------------------------------------------------- + + +def test_native_parse_accepts_and_ignores_ctx( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + calls = _spy_entry(monkeypatch, parser_mod, "parse_native") + p = tmp_path / "t.yaml" + native_mod.NativeGraphAdapter.parse(p, ctx=_full_ctx()) + native_mod.NativeGraphAdapter.parse(p, ctx=None) + assert calls == [((p,), {}), ((p,), {})] + + +# --- trust/revision publish ----------------------------------------------------- + + +def test_publish_ctx_tokenizer_env_trust_set_publishes_revision_verbatim( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = _spy_publish(monkeypatch) + publish_ctx_tokenizer_env( + GraphParseContext(tokenizer_trust_remote_code=True, tokenizer_revision="abc123") + ) + assert calls == [((), {"trust_remote_code": True, "revision": "abc123"})] + + +def test_publish_ctx_tokenizer_env_trust_false_none_revision_passes_verbatim( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # revision=None passes verbatim; configure_loader_tokenizer_env's own + # None -> "main" fallback handles it. + calls = _spy_publish(monkeypatch) + publish_ctx_tokenizer_env(GraphParseContext(tokenizer_trust_remote_code=False)) + assert calls == [((), {"trust_remote_code": False, "revision": None})] + + +def test_publish_ctx_tokenizer_env_trust_unset_skips_publish( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = _spy_publish(monkeypatch) + publish_ctx_tokenizer_env(GraphParseContext(tokenizer_revision="abc123")) + publish_ctx_tokenizer_env(None) + assert calls == [] + + +def test_weka_parse_ctx_trust_set_publishes_loader_env( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + publish_calls = _spy_publish(monkeypatch) + _spy_entry(monkeypatch, weka_trace, "from_weka_trace") + weka_trace.WekaTraceAdapter.parse( + tmp_path / "t.json", + ctx=GraphParseContext( + tokenizer_trust_remote_code=True, tokenizer_revision="pin" + ), + ) + assert publish_calls == [((), {"trust_remote_code": True, "revision": "pin"})] + + +def test_dynamo_parse_ctx_trust_set_publishes_loader_env( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + publish_calls = _spy_publish(monkeypatch) + _spy_entry(monkeypatch, dynamo_trace, "from_dynamo_trace") + dynamo_trace.DynamoTraceAdapter.parse( + tmp_path / "t.jsonl", + ctx=GraphParseContext( + tokenizer_trust_remote_code=False, tokenizer_revision="pin" + ), + ) + assert publish_calls == [((), {"trust_remote_code": False, "revision": "pin"})] + + +def test_weka_parse_ctx_none_does_not_publish( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + publish_calls = _spy_publish(monkeypatch) + _spy_entry(monkeypatch, weka_trace, "from_weka_trace") + weka_trace.WekaTraceAdapter.parse(tmp_path / "t.json", ctx=None) + assert publish_calls == [] + + +# --- parser plumbing ------------------------------------------------------------ + + +def test_parse_graph_routes_ctx_to_adapter( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + calls = _spy_entry(monkeypatch, weka_trace, "from_weka_trace") + p = tmp_path / "t.json" + parser_mod.parse_graph( + p, format="weka_trace", ctx=GraphParseContext(content_root_seed=99) + ) + assert calls[0][1] == {"content_root_seed": 99} + + +def test_parse_graph_ctx_none_routes_protocol_defaults( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + calls = _spy_entry(monkeypatch, weka_trace, "from_weka_trace") + p = tmp_path / "t.json" + parser_mod.parse_graph(p, format="weka_trace") + assert calls == [((p,), {})] + + +def test_parse_graph_signature_drops_content_root_seed_for_ctx() -> None: + params = inspect.signature(parser_mod.parse_graph).parameters + assert "ctx" in params + assert "content_root_seed" not in params + assert not hasattr(parser_mod, "_accepts_content_root_seed") diff --git a/tests/unit/graph/test_graph_path_catalog.py b/tests/unit/graph/test_graph_path_catalog.py new file mode 100644 index 0000000000..7c87ab395f --- /dev/null +++ b/tests/unit/graph/test_graph_path_catalog.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Graph-path/node-ordinal catalog over the weka segment-trie ``ParsedGraph``. + +The trie IR is a single flat top graph of ``LlmNode``s (one per recorded +request, no subgraphs/namespaces); :func:`build_graph_path_catalog` resolves each +fired node to the SAME dense ordinal the build-plane segment store was written at +(recorded ``arrival_offset_us`` order, ties broken by node id). +""" + +from pathlib import Path + +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.graph_path_catalog import build_graph_path_catalog +from aiperf.dataset.graph.models import LlmNode, resolve_trace_graph + +FIX = Path(__file__).parent / "fixtures" / "weka_min.json" + + +def test_catalog_top_only_is_stable_and_dense() -> None: + """The trie catalog assigns dense, unique, stable ordinals in arrival order.""" + parsed = from_weka_trace(str(FIX)) + cat = build_graph_path_catalog(parsed) + trace = parsed.traces[0] + assert trace.id in cat + keys = cat[trace.id] + + # One ordinal per recorded-request LlmNode; dense ``0..N-1`` and unique. + graph = resolve_trace_graph(parsed, trace) + llm_ids = {nid for nid, n in graph.nodes.items() if isinstance(n, LlmNode)} + assert set(keys) == llm_ids + assert sorted(keys.values()) == list(range(len(keys))) + + # Ordinals follow recorded arrival order (the three weka_min turns at + # t=0/1.5/3.0 fire in that order). + by_offset = sorted(llm_ids, key=lambda nid: graph.nodes[nid].arrival_offset_us or 0) + assert [keys[nid] for nid in by_offset] == list(range(len(by_offset))) + + # Stable across runs. + assert build_graph_path_catalog(parsed)[trace.id] == keys diff --git a/tests/unit/graph/test_graph_return_bridge.py b/tests/unit/graph/test_graph_return_bridge.py new file mode 100644 index 0000000000..3617a54d6f --- /dev/null +++ b/tests/unit/graph/test_graph_return_bridge.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""R2 — unconditional graph-return hook on CreditCallbackHandler. + +Graph credit returns must reach the adapter via a DEDICATED, UNCONDITIONAL +observer (not the gated ``strategy.handle_credit_return``, which is skipped +when ``can_send_any_turn()`` is False). On a graph credit the observer fires +regardless of phase-send gating, carrying ``(credit, error, cancelled)`` so the +adapter can resolve OR reject the parked Future. +""" + +from __future__ import annotations + +import pytest + +from aiperf.credit.messages import CreditReturn +from aiperf.credit.structs import Credit + +pytestmark = pytest.mark.asyncio + + +def _graph_credit( + *, x_corr: str = "x0", trace_id: str = "t0", final: bool = True +) -> Credit: + return Credit( + id=1, + phase="profiling", + conversation_id="c", + x_correlation_id=x_corr, + turn_index=0, + num_turns=1, + issued_at_ns=0, + trace_id=trace_id, + node_ordinal=0, + phase_variant="profiling", + ) + + +def _non_graph_credit() -> Credit: + return Credit( + id=2, + phase="profiling", + conversation_id="c", + x_correlation_id="x1", + turn_index=0, + num_turns=1, + issued_at_ns=0, + ) + + +def _make_handler(): + from aiperf.credit.callback_handler import CreditCallbackHandler + + class _FakeConcurrency: + def release_session_slot(self, phase) -> None: ... + def release_prefill_slot(self, phase) -> None: ... + + return CreditCallbackHandler(_FakeConcurrency()) + + +async def test_graph_return_observer_fires_unconditionally(monkeypatch): + """Even when no phase handler is registered (can_send_any_turn would be + False / the gated path is unreachable), the graph observer still fires.""" + handler = _make_handler() + seen: list[tuple[Credit, str | None, bool]] = [] + + def _observer(credit: Credit, error: str | None, cancelled: bool) -> None: + seen.append((credit, error, cancelled)) + + handler.set_graph_return_observer(_observer) + + credit = _graph_credit() + ret = CreditReturn(credit=credit, cancelled=False, error=None) + # No phase registered -> the gated strategy path no-ops, but the graph + # observer must still receive the return. + await handler.on_credit_return("w0", ret) + + assert len(seen) == 1 + got_credit, got_error, got_cancelled = seen[0] + assert got_credit.trace_id == "t0" + assert got_error is None + assert got_cancelled is False + + +async def test_graph_return_observer_forwards_error_and_cancelled(): + handler = _make_handler() + seen: list[tuple[str | None, bool]] = [] + handler.set_graph_return_observer(lambda c, e, x: seen.append((e, x))) + + await handler.on_credit_return( + "w0", + CreditReturn(credit=_graph_credit(x_corr="xe"), cancelled=False, error="boom"), + ) + await handler.on_credit_return( + "w0", + CreditReturn(credit=_graph_credit(x_corr="xc"), cancelled=True, error=None), + ) + + assert ("boom", False) in seen + assert (None, True) in seen + + +async def test_non_graph_credit_does_not_invoke_graph_observer(): + handler = _make_handler() + seen: list[Credit] = [] + handler.set_graph_return_observer(lambda c, e, x: seen.append(c)) + + await handler.on_credit_return( + "w0", CreditReturn(credit=_non_graph_credit(), cancelled=False, error=None) + ) + assert seen == [] diff --git a/tests/unit/graph/test_graph_sticky_stamp.py b/tests/unit/graph/test_graph_sticky_stamp.py new file mode 100644 index 0000000000..c58f7a4497 --- /dev/null +++ b/tests/unit/graph/test_graph_sticky_stamp.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Adapter-side trajectory-corr minting + refusal taxonomy. + +The dispatch adapter mints one ``x_correlation_id`` per TRAJECTORY +(``{conversation_id}::{nonce}`` -- the legacy per-session HTTP identity; +routing keys on the instance ``trace_id``) and stamps it on every +``TurnToSend`` -- on BOTH issuer arms: an accepted issue parks a Future and +resolves on the credit return; a refusal surfaces as the typed +``CreditIssueRefusedError`` so the executor treats it as a trace-stop +instead of sentinel-continuing. +""" + +import asyncio +from dataclasses import dataclass + +import pytest + +from aiperf.dataset.graph.graph_path_catalog import CatalogContext +from aiperf.graph.credit_dispatch_adapter import ( + CreditDispatchAdapter, + CreditIssueRefusedError, + GraphDispatchError, +) +from aiperf.graph.placement import DispatchRequest, PlacementContext + + +class _RefusingIssuer: + """Records the turn, then refuses (stop gate / caps) like the real issuer.""" + + def __init__(self) -> None: + self.sent: list[object] = [] + + async def issue_graph_credit(self, turn: object) -> bool: + self.sent.append(turn) + return False + + +class _AcceptingIssuer(_RefusingIssuer): + """Records the turn and ACCEPTS (credit placed on the wire).""" + + async def issue_graph_credit(self, turn: object) -> bool: + self.sent.append(turn) + return True + + +@dataclass +class _FakeCredit: + """Minimal Credit-like carrying the correlation identity the bridge keys on.""" + + x_correlation_id: str + turn_index: int + trace_id: str + node_ordinal: int + phase_variant: str = "profiling" + + +@dataclass +class _FakeLlmNode: + output: str = "out" + + +def _adapter(issuer: _RefusingIssuer) -> CreditDispatchAdapter: + return CreditDispatchAdapter( + credit_issuer=issuer, + catalog_context=CatalogContext( + catalog={"t-1": {"t-1:0": 0}}, + ), + trace_id="t-1", + phase_variant="profiling", + ) + + +@pytest.mark.asyncio +async def test_refusal_raises_typed_error_and_turn_carries_trajectory_corr() -> None: + issuer = _RefusingIssuer() + adapter = _adapter(issuer) + + with pytest.raises(CreditIssueRefusedError): + await adapter.dispatch( + _FakeLlmNode(), + DispatchRequest(node_id="t-1:0"), + PlacementContext(parent_trace_id="t-1", parent_node_id="t-1:0"), + ) + + assert len(issuer.sent) == 1 + turn = issuer.sent[0] + # Per-trajectory affinity: the corr is ``{conversation_id}::{nonce}``. + assert turn.x_correlation_id.startswith("t-1::") + assert adapter.phase_variant == "profiling" + + +@pytest.mark.asyncio +async def test_accepted_issue_mints_trajectory_corr_and_resolves() -> None: + """The issuing arm (issuer returns True): same corr mint, no refusal. + + The A/B counterpart to the refusal test above -- the trajectory corr must + ride the ACCEPTED turn too, and the dispatch must complete when the + correlated credit return resolves the parked Future. + """ + issuer = _AcceptingIssuer() + adapter = _adapter(issuer) + + task = asyncio.create_task( + adapter.dispatch( + _FakeLlmNode(), + DispatchRequest(node_id="t-1:0"), + PlacementContext(parent_trace_id="t-1", parent_node_id="t-1:0"), + ) + ) + await asyncio.sleep(0) + + assert len(issuer.sent) == 1 + turn = issuer.sent[0] + assert turn.x_correlation_id.startswith("t-1::") + + adapter.resolve( + _FakeCredit( + x_correlation_id=turn.x_correlation_id, + turn_index=turn.turn_index, + trace_id=turn.trace_id, + node_ordinal=turn.node_ordinal, + phase_variant=turn.phase_variant, + ), + error=None, + cancelled=False, + ) + assert isinstance(await task, str) + + +def test_refusal_error_is_a_dispatch_error_subclass() -> None: + """Existing except-GraphDispatchError sites keep catching refusals.""" + assert issubclass(CreditIssueRefusedError, GraphDispatchError) diff --git a/tests/unit/graph/test_hf_sidecar_scope_out.py b/tests/unit/graph/test_hf_sidecar_scope_out.py new file mode 100644 index 0000000000..12fe51c4db --- /dev/null +++ b/tests/unit/graph/test_hf_sidecar_scope_out.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``_looks_like_hf_dataset_id`` classification pins. + +The store build does not key on this guard (every graph workload takes +``GraphStoreBuilder._build_graph_store_streaming`` regardless of source), but +the classifier still decides whether the weka loader treats an argument as an +HF repo id (``datasets.load_dataset``) or a local file/dir. These tests pin +that classification — a known weka HF repo id returns True, and an existing +local fixture path returns False. +""" + +from __future__ import annotations + +from pathlib import Path + +from aiperf.dataset.graph.adapters.weka.trace import _looks_like_hf_dataset_id + +FIXTURES = Path(__file__).parent / "fixtures" +WEKA_MIN = FIXTURES / "weka_min.json" + + +def test_known_weka_hf_id_triggers_hf_guard(): + # A canonical published weka HF corpus id must be classified as HF so the + # loader fetches it via datasets.load_dataset instead of the filesystem. + assert _looks_like_hf_dataset_id("semianalysisai/cc-traces-weka-062126") is True + + +def test_local_fixture_path_does_not_trigger_hf_guard(): + # An existing local filesystem Path must NOT be classified as HF — the + # loader reads it as a local weka trace file. + assert _looks_like_hf_dataset_id(WEKA_MIN) is False diff --git a/tests/unit/graph/test_hf_streaming_trie_stores.py b/tests/unit/graph/test_hf_streaming_trie_stores.py new file mode 100644 index 0000000000..fee72cec1c --- /dev/null +++ b/tests/unit/graph/test_hf_streaming_trie_stores.py @@ -0,0 +1,405 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""HF-streaming trie store build path (Task T8-hf). + +The HF ``--public-dataset`` build is STREAMED per row (the parent never holds +the whole corpus's synthesized real content). Before this task the streaming +worker dropped ``pg.segment_pool`` and emitted legacy ``messages_delta`` +envelopes with NO segment store -- producing corrupt / empty prompts. + +These tests drive the streaming trie path directly through its +worker+consumer functions (``iter_item_segment_payloads`` -> +``build_unified_trie_store_from_payloads``) on a small set of synthetic weka +rows and prove: + +* the streamed payloads carry ``prompt_segment_ids`` envelopes (NOT legacy + ``messages_delta``), and the streamed catalog ordinals match + ``trie_node_ordinals``; +* worker ``materialize_graph_request_unified`` against the STREAMED unified + store reproduces the SAME prompt as the EAGER interned unified store for the + same rows + seed (byte-equal, eager store as oracle) -- and the two stores + are byte-for-byte identical on disk. + +Hermetic: the ``fake_tokenizer`` fixture pins ``Tokenizer.from_pretrained`` to +the deterministic ``FakeTokenizer`` so both the eager and streaming sides +synthesize identical content from the same seed -- no network, no real gpt2. +""" + +from __future__ import annotations + +import orjson +import pytest + +from aiperf.dataset.graph.adapters.weka.trace_parallel import ( + iter_item_segment_payloads, + parse_items, + row_work_items, +) +from aiperf.dataset.graph.models import LlmNode +from aiperf.dataset.graph.segment_ir.store_builder import ( + TraceSegmentPayload, + build_unified_trie_store_from_payloads, + build_unified_trie_store_interned, + trie_node_ordinals, +) +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) +from aiperf.graph.worker_materialize import materialize_graph_request_unified + +_SEED = 42 +_SOURCE = "synthetic/weka-trie" + +# Two small synthetic weka rows (the same dict shape a .json file / HF row +# carries). Multi-turn + a subagent so the trie pool has shared-prefix dedup +# and a deepest path; two rows exercise cross-row segment dedup in the consumer. +_ROW_A = { + "id": "trace_alpha", + "models": ["claude-opus-4-5-20251101"], + "block_size": 64, + "hash_id_scope": "local", + "requests": [ + { + "t": 0.0, + "type": "n", + "model": "claude-opus-4-5-20251101", + "in": 180, + "out": 25, + "hash_ids": [1, 2], + "input_types": ["text"], + "output_types": ["text"], + "stop": "end_turn", + "api_time": 0.8, + "think_time": 0.0, + }, # noqa: E501 + { + "t": 1.0, + "type": "n", + "model": "claude-opus-4-5-20251101", + "in": 240, + "out": 30, + "hash_ids": [1, 2, 3], + "input_types": ["text"], + "output_types": ["text"], + "stop": "end_turn", + "api_time": 0.9, + "think_time": 0.1, + }, # noqa: E501 + ], +} +_ROW_B = { + "id": "trace_beta", + "models": ["claude-opus-4-5-20251101"], + "block_size": 64, + "hash_id_scope": "local", + "requests": [ + { + "t": 0.0, + "type": "n", + "model": "claude-opus-4-5-20251101", + "in": 180, + "out": 25, + "hash_ids": [1, 2], + "input_types": ["text"], + "output_types": ["text"], + "stop": "tool_use", + "api_time": 0.8, + "think_time": 0.0, + }, # noqa: E501 + { + "t": 1.0, + "type": "subagent", + "agent_id": "agent_001", + "subagent_type": "Explore", + "duration_ms": 4000, + "total_tokens": 600, + "tool_use_count": 1, + "status": "completed", + "models": ["claude-opus-4-5-20251101"], + "tool_tokens": 0, + "system_tokens": 0, + "requests": [ + { + "t": 1.2, + "type": "n", + "model": "claude-opus-4-5-20251101", + "in": 200, + "out": 30, + "hash_ids": [10, 11], + "input_types": ["text"], + "output_types": ["text"], + "stop": "end_turn", + "api_time": 0.9, + "think_time": 0.0, + }, # noqa: E501 + ], + }, + { + "t": 6.0, + "type": "n", + "model": "claude-opus-4-5-20251101", + "in": 280, + "out": 45, + "hash_ids": [1, 2, 3, 4], + "input_types": ["text"], + "output_types": ["text"], + "stop": "end_turn", + "api_time": 1.3, + "think_time": 0.5, + }, # noqa: E501 + ], +} + +_PARSE_KWARGS = { + "tag": "weka", + "idle_gap_cap_seconds": None, + "content_root_seed": _SEED, + "max_osl": None, +} + + +def _stream_trie_payloads(rows: list[dict]) -> list[TraceSegmentPayload]: + """Materialize the streaming worker output (serial in-process path).""" + return list( + iter_item_segment_payloads( + row_work_items(rows, _SOURCE), + source_label=_SOURCE, + parse_kwargs=dict(_PARSE_KWARGS), + ) + ) + + +async def _build_eager_store(rows: list[dict], tmp_path, benchmark_id: str): + """Eager ORACLE: merged ParsedGraph -> interned unified store.""" + parsed = parse_items( + row_work_items(rows, _SOURCE), + source_label=_SOURCE, + parse_kwargs=dict(_PARSE_KWARGS), + ) + assert parsed.segment_pool is not None, "eager trie parse must surface a pool" + + store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id=benchmark_id + ) + catalog = await build_unified_trie_store_interned(parsed, store) + return parsed, catalog + + +async def _build_streaming_store(rows: list[dict], tmp_path, benchmark_id: str): + """Streaming path: per-row trie payloads -> interned unified store.""" + payloads = _stream_trie_payloads(rows) + store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id=benchmark_id + ) + return await build_unified_trie_store_from_payloads(payloads, store) + + +@pytest.mark.asyncio +async def test_streaming_emits_trie_payloads_not_legacy( + fake_tokenizer: None, # noqa: ARG001 # side-effect: deterministic tokenizer +) -> None: + """The streaming worker emits TraceSegmentPayloads w/ envelopes.""" + payloads = _stream_trie_payloads([_ROW_A, _ROW_B]) + + assert payloads, "streaming must yield at least one payload" + assert all(isinstance(p, TraceSegmentPayload) for p in payloads) + # At least one payload carries segments (the pool triples) -> a segment store + # WILL be produced; legacy runs never carry these. + assert any(p.segments for p in payloads) + + for payload in payloads: + assert payload.envelopes, f"trace {payload.trace_id} must carry envelopes" + for node_envelope in payload.envelopes: + env = orjson.loads(node_envelope.envelope_bytes) + assert "prompt_segment_ids" in env, "trie envelope, not legacy" + assert "messages_delta" not in env, "must NOT be a legacy delta envelope" + + +@pytest.mark.asyncio +async def test_streaming_catalog_matches_trie_node_ordinals( + fake_tokenizer: None, # noqa: ARG001 + tmp_path, +) -> None: + """Streamed catalog == eager catalog == per-trace ``trie_node_ordinals``.""" + parsed, eager_catalog = await _build_eager_store( + [_ROW_A, _ROW_B], tmp_path / "e", "e" + ) + stream_catalog = await _build_streaming_store([_ROW_A, _ROW_B], tmp_path / "s", "s") + + assert stream_catalog == eager_catalog, "streaming catalog must equal eager" + + for trace in parsed.traces: + llm_nodes = { + nid: n + for nid, n in parsed.graphs[trace.graph_ref].nodes.items() + if isinstance(n, LlmNode) + } + assert stream_catalog[trace.id] == trie_node_ordinals(llm_nodes) + + +@pytest.mark.asyncio +async def test_streaming_worker_materialization_byte_equal_to_eager( + fake_tokenizer: None, # noqa: ARG001 + tmp_path, +) -> None: + """CRITICAL: worker materialization over the streamed store == eager, byte-equal. + + For every trie node the worker-side ``materialize_graph_request_unified`` + reads the persisted interned manifest + content pool. The EAGER unified + store is the oracle: the streamed prompt must match it byte-for-byte (and + equal the pool ground truth), proving the HF streaming path produces the + identical materialized request the local eager path does. + """ + parsed, eager_catalog = await _build_eager_store( + [_ROW_A, _ROW_B], tmp_path / "e", "e" + ) + stream_catalog = await _build_streaming_store([_ROW_A, _ROW_B], tmp_path / "s", "s") + assert stream_catalog == eager_catalog + + pool = parsed.segment_pool + eager_client = GraphSegmentUnifiedClient(base_path=tmp_path / "e", benchmark_id="e") + stream_client = GraphSegmentUnifiedClient( + base_path=tmp_path / "s", benchmark_id="s" + ) + eager_client.open() + stream_client.open() + try: + checked = 0 + for trace in parsed.traces: + llm_nodes = { + nid: n + for nid, n in parsed.graphs[trace.graph_ref].nodes.items() + if isinstance(n, LlmNode) + } + ordinals = stream_catalog[trace.id] + for node_id, node in llm_nodes.items(): + ordinal = ordinals[node_id] + expected = pool.materialize(node.metadata["trie"]["prompt_segment_ids"]) + eager_req = materialize_graph_request_unified( + eager_client, trace.id, ordinal, "profiling" + ) + stream_req = materialize_graph_request_unified( + stream_client, trace.id, ordinal, "profiling" + ) + assert eager_req is not None and stream_req is not None + assert eager_req["messages"] == expected, node_id + # Byte-equal eager-vs-streaming: identical materialized request. + assert stream_req == eager_req, node_id + checked += 1 + assert checked == sum(len(v) for v in stream_catalog.values()) + finally: + eager_client.close() + stream_client.close() + + +@pytest.mark.asyncio +async def test_streaming_payloads_ship_content_free_structural_graph( + fake_tokenizer: None, # noqa: ARG001 # side-effect: deterministic tokenizer +) -> None: + """Each streamed payload carries a content-free structural graph for the sidecar. + + ``LlmNode.prompt`` (the dominant inline-content field), ``replay_outputs``, and + the segment pool content are all emptied -- but the pool is kept non-None so the + loaded graph keeps the trie ordinal scheme -- while topology (nodes/edges) is + preserved. This is what makes the streaming ``graph_meta`` sidecar bounded. + """ + from aiperf.dataset.graph.codecs import decode_parsed_graph_msgpack + + payloads = _stream_trie_payloads([_ROW_A, _ROW_B]) + structural = [p.structural_graph for p in payloads if p.structural_graph] + assert structural, "at least one payload must ship a structural graph" + + for blob in structural: + pg = decode_parsed_graph_msgpack(blob) + # Trie ordinal scheme preserved: pool present but content emptied. + assert pg.segment_pool is not None + assert not pg.segment_pool._by_id, "structural pool must be emptied" + # Content-free: every LlmNode.prompt emptied across every graph. + graphs = [pg.graph, *pg.graphs.values()] + llm_seen = False + for g in graphs: + for node in g.nodes.values(): + if isinstance(node, LlmNode): + llm_seen = True + assert node.prompt == [], "LlmNode.prompt must be stripped" + assert llm_seen, "structural graph must retain its LlmNode topology" + + +@pytest.mark.asyncio +async def test_streaming_structural_sink_matches_store_catalog( + fake_tokenizer: None, # noqa: ARG001 # side-effect: deterministic tokenizer + tmp_path, +) -> None: + """The merged structural sink rebuilds the SAME catalog the store was built at. + + ``catalogs_match`` is the gate that decides whether the streaming sidecar is + written; this proves it passes, so the TimingManager loads the sidecar instead + of the whole-corpus re-parse. + """ + from aiperf.dataset.graph.codecs import decode_parsed_graph_msgpack + from aiperf.dataset.graph.graph_meta_sidecar import catalogs_match + from aiperf.dataset.graph.merge import merge_parsed_graphs + + bid = "sidecar_stream_test" + payloads = _stream_trie_payloads([_ROW_A, _ROW_B]) + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id=bid) + sink: list[bytes] = [] + catalog = await build_unified_trie_store_from_payloads( + payloads, store, structural_sink=sink + ) + + assert sink, "structural sink must be populated when requested" + merged = merge_parsed_graphs(decode_parsed_graph_msgpack(b) for b in sink) + assert catalogs_match(merged, catalog), ( + "merged structural catalog must equal the store catalog (sidecar gate)" + ) + + +@pytest.mark.asyncio +async def test_streaming_unified_store_byte_matches_eager_interned( + fake_tokenizer: None, # noqa: ARG001 # side-effect: deterministic tokenizer + tmp_path, +) -> None: + """Streaming unified store (from payloads) == eager interned unified store, byte-for-byte. + + Proves ``build_unified_trie_store_from_payloads`` (the corpus-scale HF path) + produces the SAME unified store as the eager + ``build_unified_trie_store_interned``, so the unified store is the real + store on both paths. + """ + rows = [_ROW_A, _ROW_B] + + # Eager interned unified store. + eager_parsed = parse_items( + row_work_items(rows, _SOURCE), + source_label=_SOURCE, + parse_kwargs=dict(_PARSE_KWARGS), + ) + assert eager_parsed.segment_pool is not None + eager_store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id="eager" + ) + eager_catalog = await build_unified_trie_store_interned(eager_parsed, eager_store) + + # Streaming unified store from per-row payloads. + stream_store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id="stream" + ) + stream_catalog = await build_unified_trie_store_from_payloads( + _stream_trie_payloads(rows), stream_store + ) + + assert stream_catalog == eager_catalog, "catalog must match eager interned" + + eager_dir = tmp_path / "aiperf_graph_segments_eager" + stream_dir = tmp_path / "aiperf_graph_segments_stream" + eager_files = sorted(p.name for p in eager_dir.iterdir()) + stream_files = sorted(p.name for p in stream_dir.iterdir()) + assert eager_files == stream_files and eager_files, ( + f"unified store file sets differ: {eager_files} vs {stream_files}" + ) + for name in eager_files: + assert (eager_dir / name).read_bytes() == (stream_dir / name).read_bytes(), ( + f"unified store file {name!r} differs between streaming and eager" + ) diff --git a/tests/unit/graph/test_issuer_graph_path.py b/tests/unit/graph/test_issuer_graph_path.py new file mode 100644 index 0000000000..ae3f037cdb --- /dev/null +++ b/tests/unit/graph/test_issuer_graph_path.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""R4 (slot bypass) — the graph-issue path skips session-slot acquisition. + +Graph credits flow through real issuance -> router -> returns but BYPASS the +linear session-slot acquire (turn0) / release (is_final_turn && depth0) +arithmetic; the strategy owns trace concurrency. A normal credit still acquires +a session slot. Verified against the REAL ``CreditIssuer`` with recording fakes. +""" + +from __future__ import annotations + +import hashlib + +import pytest + +from aiperf.common.enums import CreditPhase +from aiperf.credit.issuer import CreditIssuer +from aiperf.credit.structs import TurnToSend + +pytestmark = pytest.mark.asyncio + + +class RecordingConcurrency: + """Records session/prefill acquisitions; always grants.""" + + def __init__(self) -> None: + self.session_acquires = 0 + self.session_releases = 0 + self.prefill_acquires = 0 + + async def acquire_session_slot(self, phase, can_proceed_fn) -> bool: + self.session_acquires += 1 + return can_proceed_fn() + + def release_session_slot(self, phase) -> None: + self.session_releases += 1 + + async def acquire_prefill_slot(self, phase, can_proceed_fn) -> bool: + self.prefill_acquires += 1 + return can_proceed_fn() + + def release_prefill_slot(self, phase) -> None: ... + + +class FakeProgress: + def increment_sent(self, turn) -> tuple[int, bool]: + return (0, False) + + def freeze_sent_counts(self) -> None: ... + + +class FakeStopChecker: + def can_start_new_session(self) -> bool: + return True + + def can_send_any_turn(self) -> bool: + return True + + def can_send_dag_child_turn(self) -> bool: + return True + + +class FakeRouter: + def __init__(self) -> None: + self.sent: list = [] + + async def send_credit(self, *, credit) -> None: + self.sent.append(credit) + + async def end_graph_trace(self, *args) -> None: ... + + +class FakeCancellation: + def next_cancellation_delay_ns(self, turn, phase) -> int | None: + return None + + +class FakeLifecycle: + started_at_ns = 0 + started_at_perf_ns = 0 + + +def _issuer( + concurrency: RecordingConcurrency, + router: FakeRouter, + url_selection_strategy=None, +) -> CreditIssuer: + return CreditIssuer( + phase=CreditPhase.PROFILING, + stop_checker=FakeStopChecker(), + progress=FakeProgress(), + concurrency_manager=concurrency, + credit_router=router, + cancellation_policy=FakeCancellation(), + lifecycle=FakeLifecycle(), + url_selection_strategy=url_selection_strategy, + ) + + +def _graph_turn() -> TurnToSend: + return TurnToSend( + conversation_id="t0", + x_correlation_id="x-t0", + turn_index=0, + num_turns=1, + trace_id="t0", + node_ordinal=0, + phase_variant="profiling", + ) + + +def _normal_turn() -> TurnToSend: + return TurnToSend( + conversation_id="c", + x_correlation_id="x", + turn_index=0, + num_turns=1, + ) + + +async def test_graph_issue_skips_session_slot_acquire(): + concurrency = RecordingConcurrency() + router = FakeRouter() + issuer = _issuer(concurrency, router) + + await issuer.issue_graph_credit(_graph_turn()) + + assert concurrency.session_acquires == 0 # bypassed + assert concurrency.prefill_acquires == 1 # still rate-limited per request + assert len(router.sent) == 1 + assert router.sent[0].trace_id == "t0" + + +async def test_graph_issue_propagates_node_ordinal_and_phase_variant(): + """TurnToSend -> ``_issue_credit_internal`` copy -> sent Credit (three-touch). + + ``node_ordinal`` and ``phase_variant`` are the worker's record-addressing + keys into the unified store; dropping either copy in + ``CreditIssuer._issue_credit_internal`` must fail HERE, not only at + runtime. Non-default values (ordinal 5, warmup variant) so a + default-passthrough cannot fake the propagation. + """ + router = FakeRouter() + issuer = _issuer(RecordingConcurrency(), router) + + turn = TurnToSend( + conversation_id="t0", + x_correlation_id="x-t0", + turn_index=0, + num_turns=1, + trace_id="t0", + node_ordinal=5, + phase_variant="warmup", + ) + assert await issuer.issue_graph_credit(turn) is True + + (credit,) = router.sent + assert credit.trace_id == "t0" + assert credit.node_ordinal == 5 + assert credit.phase_variant == "warmup" + + +async def test_normal_issue_still_acquires_session_slot(): + concurrency = RecordingConcurrency() + router = FakeRouter() + issuer = _issuer(concurrency, router) + + await issuer.issue_credit(_normal_turn()) + + assert concurrency.session_acquires == 1 # turn0 of a normal session + assert concurrency.prefill_acquires == 1 + assert len(router.sent) == 1 + + +# --------------------------------------------------------------------------- +# TC6 — URL affinity is a stable HASH of the TEMPLATE id (nonce-stripped), NOT +# a round-robin mint and NOT the per-trajectory uuid-bearing sticky key, so +# every instance/recycle of one template lands on the backend that primed its +# KV -- across separate per-phase issuers (warmup vs profiling). +# --------------------------------------------------------------------------- + +_URLS = ["http://u0", "http://u1", "http://u2"] + + +def _hash_url_index(template_id: str, n: int = len(_URLS)) -> int: + """The issuer's stable template->URL mapping, recomputed independently.""" + digest = hashlib.sha256(template_id.encode()).digest() + return int.from_bytes(digest[:8], "big") % n + + +def _sticky_graph_turn( + instance_id: str, sticky_suffix: str, phase_variant: str = "profiling" +) -> TurnToSend: + # x_correlation_id is the routing key AND the per-trajectory corr, + # ``{conversation_id}::{nonce}``; URL affinity ignores it and keys on the + # nonce-stripped ``trace_id`` template instead. + x_corr = f"{instance_id}::{sticky_suffix}" + return TurnToSend( + conversation_id=instance_id, + x_correlation_id=x_corr, + turn_index=0, + num_turns=1, + trace_id=instance_id, + node_ordinal=0, + phase_variant=phase_variant, + ) + + +def _url_issuer(router: FakeRouter) -> CreditIssuer: + from aiperf.timing.url_samplers import RoundRobinURLSampler + + return _issuer( + RecordingConcurrency(), + router, + url_selection_strategy=RoundRobinURLSampler(_URLS), + ) + + +async def test_warmup_and_profiling_issuers_pick_same_url_for_same_template(): + """Two DISTINCT instances of one TEMPLATE (warmup nonce vs profiling nonce) + map to the SAME url_index across two separate per-phase issuers, because URL + affinity is a stable hash of the nonce-stripped template id -- NOT the + per-trajectory uuid-bearing sticky key, and NOT a round-robin mint order.""" + warmup_router = FakeRouter() + profiling_router = FakeRouter() + warmup_issuer = _url_issuer(warmup_router) + profiling_issuer = _url_issuer(profiling_router) + + # Unrelated traffic on the profiling issuer: a hash-based mapping is + # order-independent, so this must not shift the shared template's url. + await profiling_issuer.issue_graph_credit(_sticky_graph_turn("t-9", "deadbeef")) + + # Same template "t-1", DIFFERENT per-instance nonces / phase variants. + await warmup_issuer.issue_graph_credit( + _sticky_graph_turn("t-1", "aaaa1111", phase_variant="warmup") + ) + await profiling_issuer.issue_graph_credit( + _sticky_graph_turn("t-1", "bbbb2222", phase_variant="profiling") + ) + + warmup_url = warmup_router.sent[-1].url_index + profiling_url = profiling_router.sent[-1].url_index + assert warmup_url is not None + assert warmup_url == profiling_url == _hash_url_index("t-1") + + +async def test_graph_url_is_stable_hash_of_template_across_lanes(): + """Distinct templates map deterministically via ``sha256(template) % urls`` + (nonce-stripped), NOT round-robin: each credit's url is a pure function of + its template id.""" + router = FakeRouter() + issuer = _url_issuer(router) + + for lane in range(4): + await issuer.issue_graph_credit(_sticky_graph_turn(f"t-{lane}", f"nonce{lane}")) + + assert [credit.url_index for credit in router.sent] == [ + _hash_url_index(f"t-{lane}") for lane in range(4) + ] + + +async def test_graph_credits_of_one_instance_share_one_url(): + """Every node credit of one instance rides the same backend (affinity), + at the template's stable hash bucket.""" + router = FakeRouter() + issuer = _url_issuer(router) + + # All three credits share template "t-2" (distinct per-turn nonces). + for turn_nonce in range(3): + await issuer.issue_graph_credit(_sticky_graph_turn("t-2", f"feed{turn_nonce}")) + + indices = {credit.url_index for credit in router.sent} + assert len(indices) == 1 + assert indices == {_hash_url_index("t-2")} + + +async def test_end_graph_trace_keeps_template_affinity_for_recycles(): + """``end_graph_trace`` closes the router session but deliberately does NOT + evict URL affinity: the entry keys on the nonce-stripped template, and a + recycle of the template must land on the backend that primed its KV.""" + router = FakeRouter() + issuer = _url_issuer(router) + + turn = _sticky_graph_turn("t-3", "0badf00d") + template = turn.trace_id + await issuer.issue_graph_credit(turn) + assert template in issuer._graph_url_affinity + first_url = router.sent[-1].url_index + + await issuer.end_graph_trace(template, "profiling") + assert issuer._graph_url_affinity.get(template) == first_url + + # A recycle (fresh nonce) reuses the primed backend. + await issuer.issue_graph_credit(_sticky_graph_turn("t-3", "5eed5eed")) + assert router.sent[-1].url_index == first_url == _hash_url_index("t-3") diff --git a/tests/unit/graph/test_lane_fanout_recycle.py b/tests/unit/graph/test_lane_fanout_recycle.py new file mode 100644 index 0000000000..1ffa0da354 --- /dev/null +++ b/tests/unit/graph/test_lane_fanout_recycle.py @@ -0,0 +1,314 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Lane fan-out + recycle + marker rotation locking tests. + +The load-generation model -- ``concurrency`` lanes wrapping the corpus +and recycling on every root final turn, gated by +``stop_checker.can_start_new_session`` -- on the dataflow +``GraphIRReplayStrategy``. These tests pin the three coupled behaviors a +single-pass model would miss: + +* **Lane fan-out**: ``concurrency`` lanes are built EVEN WHEN concurrency + exceeds the corpus size (lane ``i`` wraps onto ``traces[i % N]``), instead of + dispatching only ``N`` traces. +* **Recycle**: a freed lane re-dispatches a fresh root until the + stop-condition gate refuses, so a duration / request-count / session-count run + sustains load instead of stopping after one corpus pass. +* **Marker rotation**: each recycle pass mints a fresh ``{trace_id}#{pass}`` + instance id, so the per-instance cache-bust marker rotates across passes while + staying constant within one instance. +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path +from typing import Any + +import msgspec +import pytest + +from aiperf.common.enums import CacheBustTarget +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.timing.strategies.cache_bust import build_trace_instance_marker +from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy + +_FIX = Path(__file__).parent / "fixtures" / "weka_min.json" + + +def _corpus(n_traces: int) -> Any: + """Build a ``ParsedGraph`` whose ``traces`` list has ``n_traces`` distinct ids. + + The single-trace ``weka_min`` fixture is replicated under fresh trace ids so + the strategy sees an ``N``-trace corpus to wrap lanes over. Every replica + shares the one graph topology (single-graph workload), which is all the + lane/recycle logic needs. + """ + parsed = from_weka_trace(str(_FIX)) + base = parsed.traces[0] + traces = [msgspec.structs.replace(base, id=f"t-{i}") for i in range(n_traces)] + return msgspec.structs.replace(parsed, traces=traces) + + +class _Config: + """Minimal per-phase config the strategy reads stop thresholds from.""" + + timing_mode = None + phase = None + + def __init__( + self, + *, + concurrency: int | None = None, + expected_num_sessions: int | None = None, + total_expected_requests: int | None = None, + expected_duration_sec: float | None = None, + ) -> None: + self.concurrency = concurrency + self.expected_num_sessions = expected_num_sessions + self.total_expected_requests = total_expected_requests + self.expected_duration_sec = expected_duration_sec + + +class _StubIssuer: + """Issuer whose graph credits resolve immediately via the return observer. + + Recycle is driven by the executor completing, which only happens once every + parked dispatch Future resolves. We schedule the strategy's own + ``_on_graph_return`` on the loop so each issued credit returns successfully, + letting the per-lane recycle loop iterate. + """ + + def __init__(self) -> None: + self.observer = None + self.issued: list[Any] = [] + + def bind(self, strategy: GraphIRReplayStrategy) -> None: + self.observer = strategy._on_graph_return + + async def issue_graph_credit(self, turn: Any) -> bool: + self.issued.append(turn) + observer = self.observer + loop = asyncio.get_running_loop() + loop.call_soon(lambda: observer(turn, None, False)) + return True + + def mark_graph_sending_complete(self) -> None: ... + + def graph_all_returned(self) -> bool: + return True + + def set_graph_all_returned_event(self) -> None: ... + + +def _strategy( + parsed: Any, config: _Config, stop_checker: Any = None +) -> tuple[GraphIRReplayStrategy, _StubIssuer]: + issuer = _StubIssuer() + strategy = GraphIRReplayStrategy( + config=config, + credit_issuer=issuer, + stop_checker=stop_checker, + parsed_graph=parsed, + register_observer=lambda _obs: None, + start_min_ratio=0.0, + start_max_ratio=0.0, + ) + issuer.bind(strategy) + return strategy, issuer + + +# --------------------------------------------------------------------------- +# C2 -- lane fan-out exceeds corpus size +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_lane_count_is_concurrency_not_corpus_size(): + """concurrency=8 over a 2-trace corpus builds 8 lanes (AgentX _target_size). + + The prior model ran only ``N=2`` traces; the port wraps ``concurrency=8`` + lanes onto ``traces[i % 2]``, so the strategy sustains 8 in-flight instances. + A large session cap (>= concurrency) lets all 8 lanes start (the lane count is + ``min(concurrency, cap)``); the 8 INITIAL lanes are what we assert here. + """ + parsed = _corpus(2) + config = _Config(concurrency=8, expected_num_sessions=8) + + class _AllowAll: + def can_send_dag_child_turn(self) -> bool: + return True + + strategy, _issuer = _strategy(parsed, config, stop_checker=_AllowAll()) + + runs: list[tuple[str, int, int]] = [] + original = strategy._run_instance + + async def _spy(trace, lane_index, recycle_pass, **kwargs): + runs.append((trace.id, lane_index, recycle_pass)) + await original(trace, lane_index, recycle_pass, **kwargs) + + strategy._run_instance = _spy + await strategy.execute_phase() + + # 8 lanes, cap=8 -> exactly 8 initial instances (lanes 0..7), no recycle. + assert len(runs) == 8, f"expected 8 lanes fanned out, got {len(runs)}" + initial = [r for r in runs if r[2] == 0] + lanes = sorted(r[1] for r in initial) + assert lanes == list(range(8)), f"lanes not 0..7: {lanes}" + # Lane i wraps onto traces[i % 2]. + for trace_id, lane, _pass in initial: + assert trace_id == f"t-{lane % 2}", ( + f"lane {lane} wrapped onto {trace_id}, expected t-{lane % 2}" + ) + + +@pytest.mark.asyncio +async def test_no_stop_condition_runs_single_corpus_pass(): + """A bare run (no cap) covers the WHOLE corpus exactly once, then stops. + + corpus=3, concurrency=2: two lanes start on t-0/t-1 and the first freed + lane draws the last unclaimed template t-2 -- a SINGLE corpus pass (every + template claimed exactly once), not one-instance-per-lane truncation and + not unbounded recycle. + """ + parsed = _corpus(3) + config = _Config(concurrency=2) # concurrency < corpus, no cap + strategy, issuer = _strategy(parsed, config) + + await strategy.execute_phase() + assert strategy.admitted_traces == 3 + assert strategy.completed_traces == 3 + # Every template of the corpus ran exactly once (instance ids are + # "{template}::{nonce}"; the base template is everything before the '::'). + templates = [t.trace_id.split("::", 1)[0] for t in issuer.issued] + assert set(templates) == {"t-0", "t-1", "t-2"} + + +# --------------------------------------------------------------------------- +# C1 -- recycle until the stop condition fires +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_recycle_until_session_count_cap(): + """concurrency=2, --num-conversations=6 dispatches exactly 6 root instances. + + The v1 ``CreditCounter`` never bumps ``sent_sessions`` for graph credits, so + the strategy itself counts admitted roots and stops recycling at the cap -- + AgentX's ``can_start_new_session`` gate semantics, enforced strategy-side. + """ + parsed = _corpus(2) + config = _Config(concurrency=2, expected_num_sessions=6) + + class _AllowAll: + def can_send_dag_child_turn(self) -> bool: + return True + + strategy, _issuer = _strategy(parsed, config, stop_checker=_AllowAll()) + await strategy.execute_phase() + + assert strategy.admitted_traces == 6, ( + f"expected 6 root instances (2 lanes recycling to the cap), got " + f"{strategy.admitted_traces}" + ) + assert strategy.completed_traces == 6 + + +@pytest.mark.asyncio +async def test_recycle_stops_when_request_count_gate_closes(): + """Recycle halts once ``can_send_dag_child_turn`` (request-count) refuses. + + The gate flips False after a budget of new-root admissions; the lanes must + finish their in-flight instance and then stop recycling -- no unbounded spin. + """ + parsed = _corpus(2) + # No session cap: rely purely on the request-count-style gate to stop. + config = _Config(concurrency=2, total_expected_requests=999) + + class _BudgetGate: + def __init__(self, budget: int) -> None: + self.calls = 0 + self.budget = budget + + def can_send_dag_child_turn(self) -> bool: + self.calls += 1 + return self.calls <= self.budget + + gate = _BudgetGate(budget=3) + strategy, _issuer = _strategy(parsed, config, stop_checker=gate) + await strategy.execute_phase() + + # 2 initial lanes + recycles permitted while the gate is open (3 opens), + # then both lanes see it closed and stop. Bounded, never infinite. + assert 2 <= strategy.admitted_traces <= 2 + 3 + assert strategy.completed_traces == strategy.admitted_traces + + +# --------------------------------------------------------------------------- +# C3 -- cache-bust marker rotates across recycle passes +# --------------------------------------------------------------------------- + + +def test_marker_rotates_across_recycle_passes(): + """Distinct recycle-pass instance ids mint DISTINCT markers; same id shares one. + + The strategy stamps ``{trace_id}#{recycle_pass}`` on ``credit.trace_id``; the + worker digests that instance id (``build_trace_instance_marker``). Rotating + the pass rotates the digest, so a recycled instance cannot warm the prior + pass's prefix -- AgentX's per-recycle ``recycle_pass`` bump. + """ + bench = "seed-123" + ftp = CacheBustTarget.FIRST_TURN_PREFIX + # Instance id is ``{trace}#{lane}.{pass}``; rotate the pass on one lane. + pass0 = build_trace_instance_marker(bench, "t-0#0.0", target=ftp) + pass1 = build_trace_instance_marker(bench, "t-0#0.1", target=ftp) + pass2 = build_trace_instance_marker(bench, "t-0#0.2", target=ftp) + assert pass0 != pass1 != pass2 and pass0 != pass2, ( + "marker must rotate per recycle pass" + ) + # Two concurrent lanes wrapping the SAME template also decorrelate (AgentX's + # per-lane trajectory_index in the digest): t-0 on lane 0 vs lane 2. + lane0 = build_trace_instance_marker(bench, "t-0#0.0", target=ftp) + lane2 = build_trace_instance_marker(bench, "t-0#2.0", target=ftp) + assert lane0 != lane2, "same template on distinct lanes must mint distinct markers" + # Within one instance the marker is stable (every turn shares it). + again = build_trace_instance_marker(bench, "t-0#0.1", target=ftp) + assert again == pass1 + + +@pytest.mark.asyncio +async def test_adapter_stamps_instance_id_on_credit_trace_id(): + """The per-recycle adapter stamps a nonce-bearing INSTANCE id on + ``credit.trace_id``. + + Catalog/envelope lookups key on the BASE template id; the marker + return + de-mux key on the instance id (``{template}::{nonce}``, minted fresh per + recycle). We assert two DISTINCT instance ids (pass 0 + pass 1) are stamped, + both stripping to the template ``t-0``, while the credit's ``conversation_id`` + is the stable TEMPLATE trajectory id (nonce-free, deliberately shared across + recycles) -- NOT the per-instance trace_id. + """ + parsed = _corpus(1) + config = _Config(concurrency=1, expected_num_sessions=2) + + class _AllowAll: + def can_send_dag_child_turn(self) -> bool: + return True + + strategy, issuer = _strategy(parsed, config, stop_checker=_AllowAll()) + await strategy.execute_phase() + + instance_ids = {t.trace_id for t in issuer.issued} + # Two distinct nonce-bearing instances (pass 0 + pass 1 on lane 0). + assert len(instance_ids) == 2, f"expected 2 recycle instances: {instance_ids}" + assert all("::" in tid for tid in instance_ids), instance_ids + # The worker strips ``::{nonce}`` back to the base template id. + assert all(t.trace_id.split("::", 1)[0] == "t-0" for t in issuer.issued) + # conversation_id is the stable TEMPLATE trajectory id: nonce-free and + # identical across both recycle instances (instance identity rides trace_id). + conversation_ids = {t.conversation_id for t in issuer.issued} + assert len(conversation_ids) == 1, conversation_ids + assert conversation_ids.pop().split("::", 1)[0] == "t-0" + for turn in issuer.issued: + assert turn.conversation_id != turn.trace_id diff --git a/tests/unit/graph/test_mixed_anchor_gate.py b/tests/unit/graph/test_mixed_anchor_gate.py new file mode 100644 index 0000000000..3617047f22 --- /dev/null +++ b/tests/unit/graph/test_mixed_anchor_gate.py @@ -0,0 +1,219 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""R2/N7 -- a non-solo start-anchored in-edge is rejected loudly at graph load. + +Any fan-in involving a start-anchored in-edge was half-supported: the runtime +fired the target at the start anchor (silently ignoring the completion +predecessor's recorded ordering, or silently dropping a not-yet-dispatched +second start anchor from the firing gate), then died with a spurious "cycle +detected" when the other predecessor finished/dispatched. No shipped lowering +emits either shape (``apply_start_anchors`` replaces a node's WHOLE in-edge +set with exactly one edge), so ``Scheduler`` construction now rejects any +start-anchored in-edge that is not its target's ONLY in-edge, with a clear +``NotImplementedError`` naming the node and the offending edges. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from aiperf.dataset.graph.models import ( + GraphRecord, + LlmNode, + ParsedGraph, + StaticEdge, + TraceRecord, +) +from aiperf.graph.analysis.timeline import elaborate_trace +from aiperf.graph.scheduler import Scheduler + + +def _llm(output: str) -> LlmNode: + return LlmNode(prompt=[f"@{output}"], output=output) + + +def _mixed_graph() -> GraphRecord: + """START->a; a->d (completion); a->c (start-anchored); c->d (completion) is + fine -- the mix is on ``d``: b->d completion + a->d start-anchored.""" + return GraphRecord( + nodes={"a": _llm("a"), "b": _llm("b"), "d": _llm("d")}, + edges=[ + StaticEdge(source="START", target="a"), + StaticEdge(source="a", target="b", delay_after_predecessor_us=0.0), + StaticEdge( + source="a", target="d", delay_after_predecessor_start_us=1_000.0 + ), + StaticEdge(source="b", target="d", delay_after_predecessor_us=0.0), + ], + state={}, + ) + + +def test_scheduler_rejects_mixed_anchor_fan_in_naming_node_and_edges(): + with pytest.raises(NotImplementedError) as exc_info: + Scheduler(_mixed_graph()) + msg = str(exc_info.value) + assert msg.startswith("node 'd': ") + assert "mixed-anchor fan-in" in msg + assert "'a' -> 'd'" in msg # the start-anchored edge + assert "'b' -> 'd'" in msg # the completion edge + # The remediation must NOT steer users toward uniform start-anchored + # fan-in (equally unsupported); a start anchor must be the only in-edge. + assert "same anchor kind" not in msg + assert "ONLY in-edge" in msg + + +def test_scheduler_rejects_uniform_double_start_anchored_fan_in(): + """TWO start-anchored in-edges on one target are rejected too. + + The runtime half-supports this shape: the target fires at its FIRST anchor + parent's dispatch (`_compute_firing_gate_us` silently drops the + not-yet-dispatched second anchor), then the second anchor parent's dispatch + re-schedules the DONE target into the cycle guard (spurious "cycle + detected") and the whole trace unwinds. + """ + graph = GraphRecord( + nodes={"a": _llm("a"), "b": _llm("b"), "d": _llm("d")}, + edges=[ + StaticEdge(source="START", target="a"), + StaticEdge(source="a", target="b", delay_after_predecessor_us=0.0), + StaticEdge( + source="a", target="d", delay_after_predecessor_start_us=1_000.0 + ), + StaticEdge( + source="b", target="d", delay_after_predecessor_start_us=2_000.0 + ), + ], + state={}, + ) + with pytest.raises(NotImplementedError) as exc_info: + Scheduler(graph) + msg = str(exc_info.value) + assert msg.startswith("node 'd': ") + assert "multi-start-anchored fan-in" in msg + assert "'a' -> 'd'" in msg + assert "'b' -> 'd'" in msg + assert "ONLY in-edge" in msg + + +def test_executor_construction_rejects_mixed_anchor_fan_in(): + """The executor builds its Scheduler at construction, so a mixed-anchor + graph fails at load instead of firing early + spurious-cycling later.""" + from aiperf.graph.executor import TraceExecutor + + parsed = ParsedGraph(graph=_mixed_graph(), traces=[TraceRecord(id="t")]) + with pytest.raises(NotImplementedError, match="mixed-anchor fan-in"): + TraceExecutor(parsed) + + +def test_elaborate_trace_rejects_mixed_anchor_fan_in(): + """Static analysis shares the Scheduler, so it rejects the shape too.""" + parsed = ParsedGraph(graph=_mixed_graph(), traces=[TraceRecord(id="t")]) + with pytest.raises(NotImplementedError, match="mixed-anchor fan-in"): + elaborate_trace(parsed, parsed.traces[0]) + + +def test_start_edge_plus_start_anchored_edge_is_mixed_too(): + """A START in-edge is completion-kind (entry scheduling): mixing it with a + start-anchored in-edge re-schedules the fired entry node the same way.""" + graph = GraphRecord( + nodes={"a": _llm("a"), "c": _llm("c")}, + edges=[ + StaticEdge(source="START", target="a"), + StaticEdge(source="START", target="c"), + StaticEdge(source="a", target="c", delay_after_predecessor_start_us=500.0), + ], + state={}, + ) + with pytest.raises(NotImplementedError, match="node 'c': mixed-anchor fan-in"): + Scheduler(graph) + + +def test_supported_anchor_shapes_still_accepted(): + """All-completion fan-in and a SOLO start-anchored in-edge construct + cleanly -- the shapes shipped lowerings actually emit + (``apply_start_anchors`` gives a start-anchored node exactly one in-edge). + """ + all_completion = GraphRecord( + nodes={"a": _llm("a"), "b": _llm("b"), "d": _llm("d")}, + edges=[ + StaticEdge(source="START", target="a"), + StaticEdge(source="START", target="b"), + StaticEdge(source="a", target="d", delay_after_predecessor_us=0.0), + StaticEdge(source="b", target="d", delay_after_predecessor_us=0.0), + ], + state={}, + ) + solo_start_anchored = GraphRecord( + nodes={"a": _llm("a"), "c": _llm("c")}, + edges=[ + StaticEdge(source="START", target="a"), + StaticEdge(source="a", target="c", delay_after_predecessor_start_us=500.0), + ], + state={}, + ) + Scheduler(all_completion) + Scheduler(solo_start_anchored) + + +class _OrderedIssuer: + """Stub issuer that holds ``c``'s dispatch until ``b`` has completed.""" + + def __init__(self) -> None: + self.b_dispatched = asyncio.Event() + + async def dispatch( + self, node: object, request: object, ctx: object, **kw: object + ) -> str: + nid = request.node_id # type: ignore[attr-defined] + if nid == "c": + await self.b_dispatched.wait() + for _ in range(3): # let b's _fire task run to completion + await asyncio.sleep(0) + if nid == "b": + self.b_dispatched.set() + return f"ok::{nid}" + + +@pytest.mark.asyncio +async def test_done_node_reschedule_error_mentions_mixed_anchor_cause(): + """A still-reachable done-node re-schedule names mixed-anchor fan-in. + + ``b`` gates only on ``a``'s channel, so it fires and completes while its + second completion predecessor ``c`` is still in flight; ``c``'s completion + then re-schedules the done ``b`` into the cycle guard. The error message + must point at mixed-anchor fan-in as a likely cause of this shape. + """ + from aiperf.dataset.graph.models import ChannelRequirement, ChannelSpec + from aiperf.graph.executor import TraceExecutor + + graph = GraphRecord( + nodes={ + "a": _llm("a"), + "c": _llm("c"), + "b": LlmNode( + prompt=["@a"], + output="b", + inputs=[ChannelRequirement(channel="a", count=1)], + ), + }, + edges=[ + StaticEdge(source="START", target="a"), + StaticEdge(source="START", target="c"), + StaticEdge(source="a", target="b"), + StaticEdge(source="c", target="b"), + ], + state={"a": ChannelSpec(), "b": ChannelSpec(), "c": ChannelSpec()}, + ) + parsed = ParsedGraph(graph=graph, traces=[TraceRecord(id="t")]) + executor = TraceExecutor(parsed, credit_issuer=_OrderedIssuer()) + + with pytest.raises(ExceptionGroup) as exc_info: + await executor.run(parsed.traces[0]) + cycle_errors = [e for e in exc_info.value.exceptions if isinstance(e, RuntimeError)] + assert cycle_errors, exc_info.value.exceptions + msg = str(cycle_errors[0]) + assert "cycle detected" in msg + assert "mixed-anchor fan-in" in msg diff --git a/tests/unit/graph/test_pressure_frontier_chop.py b/tests/unit/graph/test_pressure_frontier_chop.py new file mode 100644 index 0000000000..5910cbe793 --- /dev/null +++ b/tests/unit/graph/test_pressure_frontier_chop.py @@ -0,0 +1,288 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Handoff frontier chop (``chop_trie_at_frontier``) -- pure rewrite tests. + +The extended-warmup handoff resumes PROFILING at each lane's pressure-stage +execution frontier instead of the original t* frontier. The chop drops every +node the pressure stage already executed (the server holds their KV), keeps +inter-survivor edges verbatim (recorded pacing resumes), and re-roots each +chain's frontier from START with a RESIDUAL delay: the recorded gap to the +next turn minus the wall-clock time already spent draining -- so the profiling +handoff ramps instead of bursting. +""" + +from __future__ import annotations + +import msgspec +import pytest +from pytest import param + +from aiperf.dataset.graph.models import ( + ChannelRequirement, + GraphRecord, + LlmNode, + ParsedGraph, + StaticEdge, + TraceRecord, +) +from aiperf.timing.snapshot_chop import chop_trie_at_frontier + +_S = 1_000_000.0 # one second in microseconds + + +def _llm(arrival_us: int, inputs: list[ChannelRequirement] | None = None) -> LlmNode: + return LlmNode( + prompt=["p"], + output="out", + arrival_offset_us=arrival_us, + inputs=list(inputs or []), + ) + + +def _parsed(nodes: dict[str, LlmNode], edges: list[StaticEdge]) -> ParsedGraph: + graph = GraphRecord(nodes=nodes, edges=edges) + return ParsedGraph(graph=graph, traces=[TraceRecord(id="t-1")]) + + +def _chain() -> ParsedGraph: + """START -> a(1s) -> b(2s) -> c(3s), 0.7s / 0.5s end-to-start gaps.""" + nodes = { + "a": _llm(int(1 * _S)), + "b": _llm(int(2 * _S)), + "c": _llm(int(3 * _S)), + } + edges = [ + StaticEdge(source="START", target="a", min_start_delay_us=1 * _S), + StaticEdge(source="a", target="b", delay_after_predecessor_us=0.7 * _S), + StaticEdge(source="b", target="c", delay_after_predecessor_us=0.5 * _S), + ] + return _parsed(nodes, edges) + + +def test_chop_trie_at_frontier_executed_dropped_and_frontier_rerooted(): + """executed={a}: survivors {b, c}; b re-roots from START; b->c kept verbatim.""" + out = chop_trie_at_frontier( + _chain(), + t_star_us=0.0, + executed=frozenset({"a"}), + return_wall_us={"a": 100.0}, + drain_end_wall_us=100.0, + ) + assert set(out.graph.nodes) == {"b", "c"} + pairs = {(e.source, e.target) for e in out.graph.edges} + assert pairs == {("START", "b"), ("b", "c")} + bc = next(e for e in out.graph.edges if e.source == "b") + assert bc.delay_after_predecessor_us == 0.5 * _S + + +@pytest.mark.parametrize( + "elapsed_us,expected_residual_us", + [ + param(0.3 * _S, 0.4 * _S, id="partial_credit"), + param(0.7 * _S, 0.0, id="fully_elapsed"), + param(2.0 * _S, 0.0, id="over_elapsed_clamps_to_zero"), + param(0.0, 0.7 * _S, id="no_elapsed_full_recorded_delay"), + ], +) # fmt: skip +def test_chop_trie_at_frontier_residual_credits_drain_elapsed( + elapsed_us: float, expected_residual_us: float +): + """Residual = recorded gap minus wall time already waited since pred return.""" + drain_end = 1_000_000.0 + out = chop_trie_at_frontier( + _chain(), + t_star_us=0.0, + executed=frozenset({"a"}), + return_wall_us={"a": drain_end - elapsed_us}, + drain_end_wall_us=drain_end, + ) + reroot = next(e for e in out.graph.edges if e.source == "START") + assert reroot.target == "b" + assert reroot.min_start_delay_us == pytest.approx(expected_residual_us) + + +def test_chop_trie_at_frontier_unanchored_frontier_bursts(): + """A dropped pred with no recorded return wall anchors nothing: residual 0.""" + out = chop_trie_at_frontier( + _chain(), + t_star_us=0.0, + executed=frozenset({"a"}), + return_wall_us={}, + drain_end_wall_us=500.0, + ) + reroot = next(e for e in out.graph.edges if e.source == "START") + assert reroot.min_start_delay_us == 0.0 + + +def test_chop_trie_at_frontier_pre_tstar_dropped_without_execution(): + """t* drops pre-t* nodes even when not in executed; boundary wall anchors.""" + out = chop_trie_at_frontier( + _chain(), + t_star_us=1.5 * _S, # a is pre-t* history; nothing executed in pressure + executed=frozenset(), + return_wall_us={"a": 900_000.0}, # priming return, merged by the caller + drain_end_wall_us=1_000_000.0, + ) + assert set(out.graph.nodes) == {"b", "c"} + reroot = next(e for e in out.graph.edges if e.source == "START") + # recorded 0.7s minus 0.1s elapsed since the priming return + assert reroot.min_start_delay_us == pytest.approx(0.6 * _S) + + +def test_chop_trie_at_frontier_and_fan_in_inputs_rescoped(): + """A survivor's AND-fan-in requirement on a dropped pred's channel is removed.""" + nodes = { + "a": _llm(0), + "b": _llm(int(1 * _S)), + "j": _llm( + int(2 * _S), + inputs=[ + ChannelRequirement(channel="a_out", count=1), + ChannelRequirement(channel="b_out", count=1), + ], + ), + } + edges = [ + StaticEdge(source="START", target="a"), + StaticEdge(source="START", target="b"), + StaticEdge(source="a", target="j", delay_after_predecessor_us=0.2 * _S), + StaticEdge(source="b", target="j", delay_after_predecessor_us=0.1 * _S), + ] + out = chop_trie_at_frontier( + _parsed(nodes, edges), + t_star_us=0.0, + executed=frozenset({"a"}), + return_wall_us={"a": 0.0}, + drain_end_wall_us=0.0, + ) + j = out.graph.nodes["j"] + assert [req.channel for req in j.inputs] == ["b_out"] + # j keeps its surviving pred b (no re-root)... + pairs = {(e.source, e.target) for e in out.graph.edges} + assert ("b", "j") in pairs + assert ("START", "j") not in pairs + # ...but the dropped a->j binding residual (0.2s, zero drain elapsed) is + # FOLDED into j's node-level gate instead of silently discarded. + assert j.min_start_delay_us == pytest.approx(0.2 * _S) + + +def test_chop_trie_at_frontier_kept_pred_residual_debits_and_caps(): + """The folded node-level residual uses the same debit + cap math as re-roots.""" + nodes = { + "a": _llm(0), + "b": _llm(int(1 * _S)), + "j": _llm(int(2 * _S)), + } + edges = [ + StaticEdge(source="START", target="a"), + StaticEdge(source="START", target="b"), + StaticEdge(source="a", target="j", delay_after_predecessor_us=300 * _S), + StaticEdge(source="b", target="j", delay_after_predecessor_us=0.1 * _S), + ] + out = chop_trie_at_frontier( + _parsed(nodes, edges), + t_star_us=0.0, + executed=frozenset({"a"}), + return_wall_us={"a": 100.0}, + drain_end_wall_us=100.0, + residual_cap_us=60 * _S, + ) + assert out.graph.nodes["j"].min_start_delay_us == pytest.approx(60 * _S) + + +def test_chop_trie_at_frontier_kept_pred_residual_max_combines_with_node_delay(): + """An existing node-level min_start_delay_us survives when larger than the fold.""" + j = _llm(int(2 * _S)) + nodes = { + "a": _llm(0), + "b": _llm(int(1 * _S)), + "j": msgspec.structs.replace(j, min_start_delay_us=5 * _S), + } + edges = [ + StaticEdge(source="START", target="a"), + StaticEdge(source="START", target="b"), + StaticEdge(source="a", target="j", delay_after_predecessor_us=0.2 * _S), + StaticEdge(source="b", target="j", delay_after_predecessor_us=0.1 * _S), + ] + out = chop_trie_at_frontier( + _parsed(nodes, edges), + t_star_us=0.0, + executed=frozenset({"a"}), + return_wall_us={"a": 0.0}, + drain_end_wall_us=0.0, + ) + assert out.graph.nodes["j"].min_start_delay_us == pytest.approx(5 * _S) + + +def test_chop_trie_at_frontier_full_execution_yields_empty_graph(): + out = chop_trie_at_frontier( + _chain(), + t_star_us=0.0, + executed=frozenset({"a", "b", "c"}), + return_wall_us={}, + drain_end_wall_us=0.0, + ) + assert out.graph.nodes == {} + assert out.graph.edges == [] + + +def test_chop_trie_at_frontier_nothing_executed_keeps_nodes_zeroes_lead(): + """executed empty + t*=0: full node set; chain head re-roots at 0 lead. + + Pressure fires everything ASAP, so a not-yet-started template's recorded + absolute lead is considered consumed; profiling resumes it immediately. + """ + out = chop_trie_at_frontier( + _chain(), + t_star_us=0.0, + executed=frozenset(), + return_wall_us={}, + drain_end_wall_us=0.0, + ) + assert set(out.graph.nodes) == {"a", "b", "c"} + reroot = next(e for e in out.graph.edges if e.source == "START") + assert reroot.target == "a" + assert reroot.min_start_delay_us == 0.0 + + +def test_chop_trie_at_frontier_start_anchored_delay_contributes_zero(): + """Start-anchored recorded delays never anchor a residual. + + The ledger wall is the predecessor's RETURN; debiting a dispatch-anchored + delay from a return-anchored elapsed would over-delay by the pred's live + service time (anchor mismatch), so start-anchored edges burst instead -- + the same 0.0 offset agentx gives pending handoff turns. + """ + nodes = {"a": _llm(0), "b": _llm(int(1 * _S))} + edges = [ + StaticEdge(source="START", target="a"), + StaticEdge(source="a", target="b", delay_after_predecessor_start_us=5 * _S), + ] + out = chop_trie_at_frontier( + _parsed(nodes, edges), + t_star_us=0.0, + executed=frozenset({"a"}), + return_wall_us={"a": 100.0}, + drain_end_wall_us=100.0, + ) + reroot = next(e for e in out.graph.edges if e.source == "START") + assert reroot.min_start_delay_us == 0.0 + + +def test_chop_trie_at_frontier_residual_clamped_to_cap(): + """A recorded gap beyond the cap resumes at the cap, not the full gap.""" + nodes = {"a": _llm(0), "b": _llm(int(1 * _S))} + edges = [ + StaticEdge(source="START", target="a"), + StaticEdge(source="a", target="b", delay_after_predecessor_us=300 * _S), + ] + out = chop_trie_at_frontier( + _parsed(nodes, edges), + t_star_us=0.0, + executed=frozenset({"a"}), + return_wall_us={"a": 100.0}, + drain_end_wall_us=100.0, + residual_cap_us=60 * _S, + ) + reroot = next(e for e in out.graph.edges if e.source == "START") + assert reroot.min_start_delay_us == pytest.approx(60 * _S) diff --git a/tests/unit/graph/test_pressure_stage.py b/tests/unit/graph/test_pressure_stage.py new file mode 100644 index 0000000000..13a22fb3a4 --- /dev/null +++ b/tests/unit/graph/test_pressure_stage.py @@ -0,0 +1,662 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Extended warmup (cache-pressure stage) -- config surface + strategy tests. + +AgentX v1.0 parity: ``--agentic-cache-warmup-duration N`` continues the live +replay compressed (zero idle delay, 1-token outputs) for N seconds after the +boundary-priming warmup drains, then drains and hands the execution frontier +to PROFILING. This file pins the graph-native config plumbing and the +pressure stage itself. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import msgspec +import pytest + +from aiperf.common.enums import CreditPhase +from aiperf.common.environment import Environment +from aiperf.common.scenario import TrajectoryWarmupFailedError +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.graph.credit_dispatch_adapter import CreditIssueRefusedError +from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy + +_FIX = Path(__file__).parent / "fixtures" / "weka_min.json" + + +def test_handoff_residual_cap_default(): + assert Environment.GRAPH.HANDOFF_RESIDUAL_CAP == 60.0 + + +def test_cli_flag_threads_pressure_duration_to_config(): + """--agentic-cache-warmup-duration is per-run config, not env: it lands on + cfg.agentic_cache_warmup_duration and rides TimingConfig onto the WARMUP + CreditPhaseConfig as cache_pressure_duration (no process-global writes).""" + from aiperf.config.flags.cli_config import CLIConfig + from aiperf.timing.config import TimingConfig + from tests.unit.conftest import make_run_from_cli + + cli = CLIConfig( + model_names=["test-model"], + input_file=str(_FIX), + request_count=3, + agentic_cache_warmup_duration=45.0, + ) + run = make_run_from_cli(cli) + assert run.cfg.agentic_cache_warmup_duration == 45.0 + tc = TimingConfig.from_run(run) + warmup = next(p for p in tc.phase_configs if p.phase == CreditPhase.WARMUP) + assert warmup.cache_pressure_duration == 45.0 + + +def test_cli_flag_unset_leaves_config_none(): + from aiperf.config.flags.cli_config import CLIConfig + from tests.unit.conftest import make_run_from_cli + + cli = CLIConfig( + model_names=["test-model"], + input_file=str(_FIX), + request_count=3, + ) + run = make_run_from_cli(cli) + assert run.cfg.agentic_cache_warmup_duration is None + + +def _corpus(n_traces: int) -> Any: + parsed = from_weka_trace(str(_FIX)) + base = parsed.traces[0] + traces = [msgspec.structs.replace(base, id=f"t-{i}") for i in range(n_traces)] + return msgspec.structs.replace(parsed, traces=traces) + + +class _Config: + timing_mode = None + + def __init__( + self, + *, + concurrency: int | None = None, + phase: CreditPhase = CreditPhase.WARMUP, + ) -> None: + self.phase = phase + self.concurrency = concurrency + self.expected_num_sessions = None + self.total_expected_requests = None + self.expected_duration_sec = None + + +class _ParkAfterIssuer: + """Resolve the first ``park_after`` graph credits instantly, park the rest. + + Parking stalls the pressure lanes so the ``wait_for`` duration timer fires + deterministically instead of recycling unboundedly fast under the + instant-sleep test fixtures. + + ``error_pred`` / ``cancel_pred`` optionally inject a terminal failure into + the ``observer(credit, error, cancelled)`` resolution: a credit matching + ``error_pred(credit, issuer)`` resolves with ``error_text``, one matching + ``cancel_pred(credit, issuer)`` resolves ``cancelled=True``. Both default to + ``None`` (every credit resolves cleanly, the byte-identical prior behavior). + The predicates receive the issuer so they can inspect ``issuer.strategy. + _pressure_active`` to target a pressure-stage credit -- the identity no + longer encodes priming-vs-pressure in the ``trace_id`` string. + """ + + def __init__( + self, + park_after: int | None = None, + *, + error_pred: Callable[[Any, Any], bool] | None = None, + cancel_pred: Callable[[Any, Any], bool] | None = None, + error_text: str = "boom", + ) -> None: + self.observer = None + self.strategy: GraphIRReplayStrategy | None = None + self.issued: list[Any] = [] + self._park_after = park_after + self._error_pred = error_pred + self._cancel_pred = cancel_pred + self._error_text = error_text + # Settable by the completeness-gate test: the strategy's teardown stash + # skips when not every warmup credit return landed. + self.all_returned = True + + def bind(self, strategy: GraphIRReplayStrategy) -> None: + self.observer = strategy._on_graph_return + self.strategy = strategy + + async def issue_graph_credit(self, credit: Any) -> bool: + self.issued.append(credit) + if self._park_after is not None and len(self.issued) > self._park_after: + return True # parked: never resolved + observer = self.observer + error = ( + self._error_text + if self._error_pred and self._error_pred(credit, self) + else None + ) + cancelled = ( + bool(self._cancel_pred(credit, self)) if self._cancel_pred else False + ) + asyncio.get_running_loop().call_soon(lambda: observer(credit, error, cancelled)) + return True + + def mark_graph_sending_complete(self) -> None: ... + + def graph_all_returned(self) -> bool: + return self.all_returned + + def set_graph_all_returned_event(self) -> None: ... + + +def _error_first(pressure_only: bool = False) -> Callable[[Any, Any], bool]: + """Predicate erroring the FIRST matching issued credit (fires once). + + ``pressure_only=False`` matches the first credit issued at all (a + boundary-priming credit); ``pressure_only=True`` matches the first credit + issued while the strategy's pressure stage is active + (``issuer.strategy._pressure_active``) -- the priming/pressure split is no + longer encoded in the ``trace_id`` string, so the stage flag is the signal. + """ + state = {"fired": False} + + def pred(credit: Any, issuer: Any) -> bool: + if state["fired"]: + return False + if pressure_only and not issuer.strategy._pressure_active: + return False + state["fired"] = True + return True + + return pred + + +def _strategy( + parsed: Any, + *, + duration: float | None, + park_after: int | None = None, + graph_channel: Any = None, + concurrency: int = 1, + phase: CreditPhase = CreditPhase.WARMUP, + error_pred: Callable[[Any], bool] | None = None, + cancel_pred: Callable[[Any], bool] | None = None, + error_text: str = "boom", +) -> tuple[GraphIRReplayStrategy, _ParkAfterIssuer]: + issuer = _ParkAfterIssuer( + park_after=park_after, + error_pred=error_pred, + cancel_pred=cancel_pred, + error_text=error_text, + ) + strategy = GraphIRReplayStrategy( + config=_Config(concurrency=concurrency, phase=phase), + graph_channel=graph_channel, + credit_issuer=issuer, + parsed_graph=parsed, + register_observer=lambda _obs: None, + start_min_ratio=0.5, + start_max_ratio=0.5, + t_star_random_seed=1234, + cache_pressure_duration_s=duration, + # A mis-tuned park_after must fail in seconds, not ride out the 300s + # adapter default (which collides with the global pytest --timeout). + dispatch_timeout_s=2.0, + ) + issuer.bind(strategy) + return strategy, issuer + + +def _pressure_flag_spy(strategy: GraphIRReplayStrategy, sink: list) -> None: + """Wrap ``_run_instance`` to record each call's (pressure, recycle_pass). + + The instance-id string no longer encodes the priming-vs-pressure flavor, so + the ``pressure`` kwarg on ``_run_instance`` is the observable that + distinguishes a boundary-priming instance from a pressure-stage one. + """ + original = strategy._run_instance + + async def _spy(trace, lane, recycle_pass, **kwargs): + sink.append((kwargs.get("pressure", False), recycle_pass)) + return await original(trace, lane, recycle_pass, **kwargs) + + strategy._run_instance = _spy + + +@pytest.mark.asyncio +async def test_execute_phase_without_duration_runs_no_pressure_instances(): + """Duration None: byte-identical warmup -- boundary priming only, no pressure.""" + parsed = _corpus(2) + strategy, issuer = _strategy(parsed, duration=None) + calls: list[tuple[bool, int]] = [] + _pressure_flag_spy(strategy, calls) + await strategy.execute_phase() + assert calls, "warmup must run at least one boundary-priming instance" + assert not any(pressure for pressure, _pass in calls), "no pressure instances" + assert strategy._pressure_active is False + + +@pytest.mark.asyncio +async def test_execute_phase_with_duration_runs_pressure_after_priming(): + """Pressure dispatches post-t* nodes AFTER priming drains (stage barrier). + + The priming/pressure split is no longer encoded in the ``trace_id`` string, + so we observe ``_run_instance``'s ``pressure`` flag: every priming instance + (pressure=False) runs before any pressure instance (pressure=True), the stage + flips ``_pressure_active``, and credits are issued. + """ + parsed = _corpus(2) + strategy, issuer = _strategy(parsed, duration=0.2, park_after=50) + calls: list[tuple[bool, int]] = [] + _pressure_flag_spy(strategy, calls) + await strategy.execute_phase() + + flags = [pressure for pressure, _pass in calls] + assert any(not f for f in flags), "boundary priming must still run first" + assert any(flags), "pressure stage must dispatch at least one instance" + # All priming (False) precedes all pressure (True) -- the stage barrier. + first_pressure = flags.index(True) + assert not any(flags[:first_pressure]), "priming must all precede pressure" + assert strategy._pressure_active is True + assert issuer.issued, "priming + pressure must have issued credits" + + +@pytest.mark.asyncio +async def test_pressure_ledger_records_return_walls_by_instance_and_node(): + parsed = _corpus(1) + strategy, issuer = _strategy(parsed, duration=0.2, park_after=50) + await strategy.execute_phase() + + assert strategy._return_walls, "ledger must be populated" + for instance_id, walls in strategy._return_walls.items(): + # Instance ids are ``{template}::{nonce}``. + assert "::" in instance_id + for node_id, wall in walls.items(): + assert isinstance(node_id, str) and node_id in parsed.graph.nodes + assert wall > 0.0 + + +def _ledger_key_for( + strategy: GraphIRReplayStrategy, credit: Any +) -> tuple[str, str] | None: + """Replicate ``_record_return_wall``'s (instance_id, node_id) mapping. + + Returns None when the credit does not resolve to a catalog node (unmappable + returns are never ledgered regardless of the cancelled gate). + """ + instance_id = getattr(credit, "trace_id", None) + ordinal = getattr(credit, "node_ordinal", None) + if instance_id is None or ordinal is None: + return None + template_id = instance_id.split("::", 1)[0] + inverse = { + o: nid for nid, o in strategy._catalog.catalog.get(template_id, {}).items() + } + node_id = inverse.get(ordinal) + return (instance_id, node_id) if node_id is not None else None + + +@pytest.mark.asyncio +async def test_cancelled_returns_never_enter_the_pressure_ledger(): + """Wire-cancelled turns are NOT executed: the server may never have + completed them, so the handoff must let profiling refire them (and a + successful grace-expiry cancel-drain then yields a VALID handoff). + + Keyed by (instance_id, node_id): a single-template pressure corpus recycles + the same node across passes, so a bare node-id set would collide with a + clean return of the same node on a later instance. + """ + parsed = _corpus(1) + cancelled: list[Any] = [] + + def cancel_first_pressure(credit: Any, issuer: Any) -> bool: + # Cancel exactly the FIRST pressure-stage return, identified by the + # strategy's pressure-stage flag (the flavor is no longer in trace_id). + if cancelled: + return False + if issuer.strategy._pressure_active: + cancelled.append(credit) + return True + return False + + strategy, _issuer = _strategy( + parsed, duration=0.2, park_after=50, cancel_pred=cancel_first_pressure + ) + await strategy.execute_phase() + + assert cancelled, "the stub must have cancelled at least one pressure return" + assert strategy._return_walls, "clean returns must still populate the ledger" + + cancelled_keys = { + key for c in cancelled if (key := _ledger_key_for(strategy, c)) is not None + } + assert cancelled_keys, "the cancelled credit must map to a catalog node" + recorded_keys = { + (instance_id, node_id) + for instance_id, walls in strategy._return_walls.items() + for node_id in walls + } + assert not (cancelled_keys & recorded_keys), ( + "a wire-cancelled return must be excluded from the pressure ledger" + ) + + +@pytest.mark.asyncio +async def test_pressure_recycles_lane_after_instance_completion(): + """A completed pressure instance frees its lane for a recycle (pass >= 1) draw. + + The recycle pass is no longer id-encoded, so observe ``_run_instance``'s + ``recycle_pass`` arg for pressure-stage instances directly. + """ + parsed = _corpus(1) + # park very late so at least one full instance completes and recycles + strategy, issuer = _strategy(parsed, duration=0.2, park_after=500) + calls: list[tuple[bool, int]] = [] + _pressure_flag_spy(strategy, calls) + await strategy.execute_phase() + pressure_passes = {recycle_pass for pressure, recycle_pass in calls if pressure} + assert any(p >= 1 for p in pressure_passes), ( + f"expected a recycled pressure pass (>= 1), saw {pressure_passes}" + ) + + +class _RefusingIssuer: + """``issue_graph_credit`` raises ``CreditIssueRefusedError`` (closed stop gate). + + Models the issuer stop gate refusing a fresh dispatch (request-count / + duration cap reached, or run cancelled): the adapter surfaces the refusal to + the executor, whose ``TaskGroup`` wraps it, and ``_leaf_credit_refusal`` + matches the (grouped) refusal so ``_run_instance`` reports a clean stop. + """ + + def __init__(self) -> None: + self.observer = None + self.issued: list[Any] = [] + + def bind(self, strategy: GraphIRReplayStrategy) -> None: + self.observer = strategy._on_graph_return + + async def issue_graph_credit(self, credit: Any) -> bool: + self.issued.append(credit) + raise CreditIssueRefusedError("stop gate closed for test") + + def mark_graph_sending_complete(self) -> None: ... + + def graph_all_returned(self) -> bool: + return True + + def set_graph_all_returned_event(self) -> None: ... + + +@pytest.mark.asyncio +async def test_run_instance_returns_true_on_issuer_refusal(): + """A clean issuer refusal makes ``_run_instance`` report True (lane stop signal).""" + parsed = _corpus(1) + issuer = _RefusingIssuer() + strategy = GraphIRReplayStrategy( + config=_Config(concurrency=1, phase=CreditPhase.PROFILING), + credit_issuer=issuer, + parsed_graph=parsed, + register_observer=lambda _obs: None, + start_min_ratio=0.0, + start_max_ratio=0.0, + dispatch_timeout_s=2.0, + ) + issuer.bind(strategy) + + refused = await strategy._run_instance(parsed.traces[0], 0, 0) + + assert refused is True + # A clean refusal is a healthy stop, NOT a trace error. + assert strategy._errored_traces == 0 + + +@pytest.mark.asyncio +async def test_run_instance_returns_false_on_clean_success(): + """A normally-resolving instance reports False (no lane stop).""" + parsed = _corpus(1) + strategy, _issuer = _strategy(parsed, duration=None, phase=CreditPhase.PROFILING) + + refused = await strategy._run_instance(parsed.traces[0], 0, 0) + + assert refused is False + + +class _FakeSource: + """Duck-typed graph channel (cross-phase warmup-handoff slot).""" + + def __init__(self) -> None: + self.warmup_handoff = None + + +@pytest.mark.asyncio +async def test_teardown_stashes_handoff_for_stalled_lanes(): + parsed = _corpus(1) + source = _FakeSource() + strategy, issuer = _strategy( + parsed, duration=0.2, park_after=3, graph_channel=source + ) + await strategy.execute_phase() + await strategy.teardown_phase() + + handoff = source.warmup_handoff + assert handoff is not None + assert handoff.lanes, "the parked lane must be live at drain" + for _lane, entry in handoff.lanes.items(): + assert entry.template_trace_id == "t-0" + assert entry.executed_node_ids <= set(parsed.graph.nodes) + # every executed node has a residual anchor wall <= drain end + for node_id in entry.executed_node_ids: + assert entry.return_wall_us[node_id] <= handoff.drain_end_wall_us + assert handoff.drain_end_wall_us > 0.0 + + +@pytest.mark.asyncio +async def test_teardown_skips_stash_when_returns_incomplete(): + """Grace-timeout / cancelled drains must NOT hand profiling a wrong handoff. + + With finite drain grace (this task), a grace-timeout force-complete reaches + teardown with graph_all_returned() False; the stash must skip so profiling + falls back to plain t* plans instead of refiring server-executed nodes. + """ + parsed = _corpus(1) + source = _FakeSource() + strategy, issuer = _strategy( + parsed, duration=0.2, park_after=3, graph_channel=source + ) + issuer.all_returned = False + await strategy.execute_phase() + await strategy.teardown_phase() + assert source.warmup_handoff is None + + +@pytest.mark.asyncio +async def test_teardown_merges_priming_walls_into_pass0_handoff(): + """Pass-0 pressure lanes carry their boundary-priming return walls too. + + A chain the pressure stage never advanced still needs its residual anchored + on the PRIMING return of its boundary turn (agentx baseline-return parity). + """ + parsed = _corpus(1) + source = _FakeSource() + # park immediately after priming: pressure issues nothing, lane stalls on + # its first pressure dispatch -> executed empty, priming walls present + strategy, issuer = _strategy( + parsed, duration=0.2, park_after=1, graph_channel=source + ) + await strategy.execute_phase() + await strategy.teardown_phase() + + handoff = source.warmup_handoff + assert handoff is not None and handoff.lanes + # concurrency=1 in the helper => exactly lane 0; assert on it explicitly + # (dict insertion order is pop/re-insert-sensitive across recycles). + entry = handoff.lanes[0] + # The priming instance id is ``t-0::{nonce}`` -- recover it as the ledgered + # instance that is NOT the live pressure instance. + priming_walls = next( + ( + walls + for iid, walls in strategy._return_walls.items() + if iid != entry.instance_id + ), + {}, + ) + assert priming_walls, "the priming instance must have ledgered returns" + assert set(priming_walls) <= set(entry.return_wall_us) + + +@pytest.mark.asyncio +async def test_teardown_handoff_carries_recycle_cursor_past_pass0(): + """The stashed handoff carries the pressure stage's next recycle cursor. + + A 1-template corpus resolves pass-0 at cursor 1 (one template consumed); a + pressure run that recycled at least once advances ``_pressure_next_index`` + strictly past that, and the stash persists it so profiling's bounded recycle + continues from the pressure stage's last draw (agentx shared-sampler parity). + """ + parsed = _corpus(1) + source = _FakeSource() + strategy, _issuer = _strategy( + parsed, duration=0.2, park_after=500, graph_channel=source + ) + await strategy.execute_phase() + await strategy.teardown_phase() + + handoff = source.warmup_handoff + assert handoff is not None + assert handoff.corpus_cursor >= 2, ( + "corpus_cursor must advance strictly past the pass-0 resolution cursor " + f"(1) after at least one recycle draw; saw {handoff.corpus_cursor}" + ) + + +@pytest.mark.asyncio +async def test_stash_records_pressure_lane_count(): + parsed = _corpus(1) + source = _FakeSource() + strategy, issuer = _strategy( + parsed, duration=0.2, park_after=3, graph_channel=source + ) + await strategy.execute_phase() + await strategy.teardown_phase() + assert source.warmup_handoff.pressure_lane_count == 1 + + +@pytest.mark.asyncio +async def test_teardown_without_pressure_leaves_source_untouched(): + parsed = _corpus(1) + source = _FakeSource() + strategy, _issuer = _strategy(parsed, duration=None, graph_channel=source) + await strategy.execute_phase() + await strategy.teardown_phase() + assert source.warmup_handoff is None + + +@pytest.mark.asyncio +async def test_report_warmup_failures_raises_on_errored_priming_return(): + """A terminal error on a boundary-priming return aborts the warmup (agentx parity).""" + parsed = _corpus(1) + strategy, issuer = _strategy(parsed, duration=None, error_pred=_error_first()) + await strategy.execute_phase() + + assert issuer.issued, "boundary priming must issue at least one credit" + assert strategy._warmup_failure_count == 1 + with pytest.raises(TrajectoryWarmupFailedError) as excinfo: + strategy.report_warmup_failures() + msg = str(excinfo.value) + assert "boom" in msg, f"abort message must surface the error text: {msg!r}" + assert "1 trace" in msg, f"abort message must surface the count: {msg!r}" + + +@pytest.mark.asyncio +async def test_report_warmup_failures_counts_pressure_errors(): + """A terminal error on a pressure-stage return aborts the warmup.""" + parsed = _corpus(1) + strategy, issuer = _strategy( + parsed, + duration=0.2, + park_after=50, + error_pred=_error_first(pressure_only=True), + ) + await strategy.execute_phase() + + assert strategy._pressure_active, ( + "pressure stage must have run to error a pressure-stage credit" + ) + assert strategy._warmup_failure_count >= 1 + with pytest.raises(TrajectoryWarmupFailedError): + strategy.report_warmup_failures() + + +@pytest.mark.asyncio +async def test_report_warmup_failures_ignores_cancelled_returns(): + """Cancelled returns are excluded from the abort gate even with error text. + + The pressure drain cancels executor coroutines on the duration timer, so a + cancellation surfacing at drain is self-inflicted teardown, not a server + failure -- it must not abort an otherwise-healthy warmup. + """ + parsed = _corpus(1) + strategy, issuer = _strategy( + parsed, + duration=None, + cancel_pred=lambda _c, _i: True, # every return is cancelled + error_pred=_error_first(), # the first also carries error text + error_text="err", + ) + await strategy.execute_phase() + + assert issuer.issued + assert strategy._warmup_failure_count == 0 + strategy.report_warmup_failures() # must not raise + + +def test_report_warmup_failures_message_carries_true_count_past_sample_cap(): + """More failures than the 5-sample cap still surface the TRUE total in the message.""" + + class _StubCredit: + def __init__(self, i: int) -> None: + self.trace_id = f"t-{i}#0.0" + self.node_ordinal = i + + parsed = _corpus(1) + strategy, _issuer = _strategy(parsed, duration=None) + for i in range(7): + strategy._record_warmup_failure(_StubCredit(i), f"boom-{i}") + + assert strategy._warmup_failure_count == 7 + assert len(strategy._warmup_failure_samples) == 5 + with pytest.raises(TrajectoryWarmupFailedError) as excinfo: + strategy.report_warmup_failures() + msg = str(excinfo.value) + assert "total 7" in msg, f"abort message must carry the true count: {msg!r}" + assert "boom-0" in msg, f"abort message must keep the real samples: {msg!r}" + + +@pytest.mark.asyncio +async def test_report_warmup_failures_noop_for_profiling_phase_and_clean_warmup(): + """PROFILING phases and clean warmups never abort the run.""" + parsed = _corpus(1) + + # A PROFILING phase with an errored return does NOT record or abort. + prof, _issuer = _strategy( + parsed, + duration=None, + phase=CreditPhase.PROFILING, + error_pred=_error_first(), + ) + await prof.execute_phase() + assert prof._warmup_failure_count == 0 + prof.report_warmup_failures() # profiling: no-op + + # A clean warmup (no terminal failures) does NOT abort. + warm, _issuer2 = _strategy(parsed, duration=None) + await warm.execute_phase() + assert warm._warmup_failure_count == 0 + warm.report_warmup_failures() # clean: no-op diff --git a/tests/unit/graph/test_r6_wiring.py b/tests/unit/graph/test_r6_wiring.py new file mode 100644 index 0000000000..f8b97d23c8 --- /dev/null +++ b/tests/unit/graph/test_r6_wiring.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""R6 wiring unit tests: input-type -> GRAPH_IR + graph-store build path symmetry. + +These pin the three R6 seams that connect the existing graph components into a +real run WITHOUT standing up the full multiprocess pipeline: + +* a weka input file is detected as a graph workload (``workload_detect``), +* the DatasetManager build store and the worker read client resolve the SAME + graph-store directory from ``(base_path, benchmark_id)`` -- the symmetry + the worker depends on to find what the build wrote, and +* build-plane unified-store ordinals match the dispatch-time catalog ordinals. +""" + +from __future__ import annotations + +import tempfile +from pathlib import Path + +import pytest + +from aiperf.dataset.graph.adapters.weka.trace import ( + WekaTraceAdapter, + from_weka_trace, +) +from aiperf.dataset.graph.graph_path_catalog import ( + build_catalog_context, + node_ordinal_for, +) +from aiperf.dataset.graph.segment_ir.store_builder import ( + build_unified_trie_store_interned, +) +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) + +FIX = Path(__file__).parent / "fixtures" / "weka_min.json" + + +def test_weka_fixture_detected_as_graph_workload(): + assert WekaTraceAdapter.can_load(FIX) is True + + +@pytest.mark.asyncio +async def test_build_envelope_ordinals_match_dispatch_catalog(): + # weka always parses to the segment-trie IR (``segment_pool`` set), so the + # build plane writes per-node manifests via + # ``build_unified_trie_store_interned`` and the dispatch plane resolves the + # SAME trie ordinals via ``build_catalog_context`` (which detects the trie + # and reuses the shared ordinal scheme). + parsed = from_weka_trace(str(FIX), content_root_seed=42) + assert parsed.segment_pool is not None, "weka now always builds the trie IR" + catalog = build_catalog_context(parsed) + with tempfile.TemporaryDirectory() as d: + store = GraphSegmentUnifiedBackingStore(base_path=d, benchmark_id="b1") + addr = await build_unified_trie_store_interned(parsed, store) + + client = GraphSegmentUnifiedClient(base_path=d, benchmark_id="b1").open() + trace_id = parsed.traces[0].id + node_map = addr[trace_id] + assert node_map, "expected per-node ordinals" + for node_key, ordinal in node_map.items(): + # The build ordinal is what the dispatch adapter resolves via the + # catalog -- they MUST agree or the worker reads the wrong manifest. + assert node_ordinal_for(catalog, trace_id, node_key) == ordinal + assert client.get_node_envelope(trace_id, ordinal, "profiling") is not None + client.close() diff --git a/tests/unit/graph/test_real_content_guard.py b/tests/unit/graph/test_real_content_guard.py new file mode 100644 index 0000000000..0a5e0c06f9 --- /dev/null +++ b/tests/unit/graph/test_real_content_guard.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Guard: weka ingest produces REAL-content prompts; no node holds a placeholder. + +The segment-trie IR is the only weka prompt path: ``from_weka_trace`` synthesizes +the agentx-faithful multi-turn conversation into the :class:`SegmentPool` at +INGEST time and stamps each ``LlmNode`` with a ``prompt_segment_ids`` path that +materializes to that real content. This guard asserts every materialized prompt +message carries real, non-empty text -- never an empty or placeholder prompt. +""" + +from pathlib import Path + +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.models import LlmNode + +FIX = Path(__file__).parent / "fixtures" / "weka_min.json" +FIX_SUB = Path(__file__).parent / "fixtures" / "weka_subagent.json" + +# Legacy transient placeholder the deleted delta-synthesis path stamped before +# filling real content. The trie path never emits it; assert it is absent. +_LEGACY_UNFILLED_PLACEHOLDER = "" + + +def _materialized_prompts(parsed) -> list[list[dict]]: + """Every LlmNode's materialized prompt messages on the trie top graph.""" + pool = parsed.segment_pool + assert pool is not None, "trie ingest must attach a SegmentPool" + out: list[list[dict]] = [] + for node in parsed.graph.nodes.values(): + if isinstance(node, LlmNode): + out.append(pool.materialize(node.metadata["trie"]["prompt_segment_ids"])) + return out + + +def test_default_ingest_yields_real_content_no_placeholder(): + parsed = from_weka_trace(str(FIX)) + prompts = _materialized_prompts(parsed) + assert prompts, "expected at least one LlmNode prompt" + + all_msgs = [m for prompt in prompts for m in prompt] + assert all_msgs, "expected materialized prompt messages" + assert any(m.get("role") == "user" for m in all_msgs), ( + "expected at least one user message" + ) + for m in all_msgs: + content = m.get("content") + assert isinstance(content, str) and content, "prompt content must be real text" + assert _LEGACY_UNFILLED_PLACEHOLDER not in content, ( + "real-content ingest must NOT emit the transient placeholder" + ) + + +def test_default_ingest_multi_turn_content_is_real_no_placeholder(): + # The subagent fixture exercises multi-turn + subagent-inner requests, all + # flattened into the trie top graph; every materialized prompt is real text. + parsed = from_weka_trace(str(FIX_SUB)) + prompts = _materialized_prompts(parsed) + assert len(prompts) > 1, "expected multiple turns flattened into the trie graph" + for prompt in prompts: + assert prompt, "every node must materialize a non-empty prompt" + for m in prompt: + assert isinstance(m["content"], str) and m["content"] + assert _LEGACY_UNFILLED_PLACEHOLDER not in m["content"] diff --git a/tests/unit/graph/test_recorded_mode_stream_override.py b/tests/unit/graph/test_recorded_mode_stream_override.py new file mode 100644 index 0000000000..65ab436974 --- /dev/null +++ b/tests/unit/graph/test_recorded_mode_stream_override.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request stream override for graph credits (Task 2). + +The recorded per-node wire mode (weka ``"n"``/``"s"``, dynamo ``ttft_ms``) wins +for graph credits over the global ``endpoint.streaming`` flag; a mode-less node +(override ``None``) follows the global. These tests pin the precedence at the +payload-stamp seam (``apply_run_level_payload_options``) and the ``RequestInfo`` +carrier field; the transport-side wiring is covered in +``tests/unit/transports/test_aiohttp_transport.py``. +""" + +from __future__ import annotations + +import pytest +from pytest import param + +from aiperf.common.enums import CreditPhase, ModelSelectionStrategy +from aiperf.common.models import EndpointInfo +from aiperf.common.models.model_endpoint_info import ( + ModelEndpointInfo, + ModelInfo, + ModelListInfo, +) +from aiperf.common.models.record_models import RequestInfo +from aiperf.graph.worker_materialize import apply_run_level_payload_options +from aiperf.plugin.enums import EndpointType + + +def _endpoint( + *, + streaming: bool, + use_server_token_count: bool = False, + extra: list[tuple[str, object]] | None = None, +) -> EndpointInfo: + return EndpointInfo( + type=EndpointType.CHAT, + streaming=streaming, + use_server_token_count=use_server_token_count, + extra=extra or [], + ) + + +@pytest.mark.parametrize( + "override,glob,expected", + [ + param(True, False, True, id="node-stream-wins-over-global-off"), + param(False, True, False, id="node-nonstream-wins-over-global-on"), + param(None, True, True, id="no-override-follows-global-on"), + param(None, False, False, id="no-override-follows-global-off"), + ], +) # fmt: skip +def test_apply_run_level_stream_precedence(override, glob, expected): + """Recorded per-node mode wins; ``None`` follows the global flag.""" + payload = {"stream": "stale"} + apply_run_level_payload_options( + payload, _endpoint(streaming=glob), stream_override=override + ) + assert payload["stream"] is expected + + +def test_include_usage_follows_final_stream(): + """``include_usage`` keys on the FINAL stamped ``stream``, not the global. + + ``stream_override=False`` + global True + ``use_server_token_count`` forces + the wire to non-streaming, so no ``stream_options`` is layered; the converse + (override True over global False) DOES layer it. + """ + # Override False beats global True -> final stream False -> no usage forced. + payload = {"messages": [{"role": "user", "content": "hi"}], "stream": True} + apply_run_level_payload_options( + payload, + _endpoint(streaming=True, use_server_token_count=True), + stream_override=False, + ) + assert payload["stream"] is False + assert "stream_options" not in payload + + # Override True beats global False -> final stream True -> usage forced. + payload2 = {"messages": [{"role": "user", "content": "hi"}], "stream": False} + apply_run_level_payload_options( + payload2, + _endpoint(streaming=False, use_server_token_count=True), + stream_override=True, + ) + assert payload2["stream"] is True + assert payload2["stream_options"] == {"include_usage": True} + + +def test_request_info_stream_override_defaults_none(): + """A plain (non-graph) ``RequestInfo`` construction leaves the override None.""" + model_endpoint = ModelEndpointInfo( + models=ModelListInfo( + models=[ModelInfo(name="test-model")], + model_selection_strategy=ModelSelectionStrategy.ROUND_ROBIN, + ), + endpoint=EndpointInfo( + type=EndpointType.CHAT, base_url="http://localhost:8000/v1/chat" + ), + ) + request_info = RequestInfo( + model_endpoint=model_endpoint, + turns=[], + turn_index=0, + credit_num=0, + credit_phase=CreditPhase.PROFILING, + x_request_id="rid", + x_correlation_id="cid", + conversation_id="conv", + ) + assert request_info.stream_override is None diff --git a/tests/unit/graph/test_registry_run_parse_parity.py b/tests/unit/graph/test_registry_run_parse_parity.py new file mode 100644 index 0000000000..d21df475d5 --- /dev/null +++ b/tests/unit/graph/test_registry_run_parse_parity.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Registry-dispatch-with-ctx vs run-parse parity pins (weka / dynamo / dag). + +These are parity PIN tests, NOT red tests: they were authored while the +per-format parse ladder was still alive and passed then, because +``resolve_graph_parse_context(run)`` carries exactly the knobs the ladder +threaded per-format. The ladder has since been deleted; the pins now hold +the registry dispatch (``adapter.parse(path, ctx)``) against +``parse_graph_workload(run, path)``: if the two ever diverge on any +run-derived knob, the identity compare below fails on the exact field that +drifted. + +Identity compared per format: trace ids, node ids, per-node +``prompt_segment_ids`` (the ordered SegmentPool walk), and the segment pool's +content-addressed ``_by_id`` key set -- together the run-visible parse +identity (topology + addressing + synthesized-content addresses). +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pytest +from pytest import param + +from aiperf.config.flags.cli_config import CLIConfig +from aiperf.dataset.graph.models import ParsedGraph +from aiperf.dataset.graph.segment_ir.envelope import read_prompt_segment_ids +from aiperf.dataset.graph.workload_detect import ( + parse_graph_workload, + resolve_graph_parse_context, +) +from aiperf.plugin import plugins +from aiperf.plugin.enums import PluginType +from tests.unit.conftest import make_run_from_cli + +_TESTS_DIR = Path(__file__).parents[2] + +WEKA_MIN = Path(__file__).parent / "fixtures" / "weka_min.json" +DYNAMO_NESTED = ( + _TESTS_DIR + / "unit/dataset/graph/adapters/fixtures/dynamo_nested/nested_2_level.jsonl.gz" +) +DAG_SMALL = _TESTS_DIR / "fixtures/dag/small.dag.jsonl" + + +def _run_for(path: Path, **cli_overrides: Any): + cfg = CLIConfig( + model_names=["test-model"], + input_file=str(path), + # Offline builtin tokenizer: the fake "test-model" must not trigger a + # HF load during content synthesis (weka / dynamo). + tokenizer_name="builtin", + **cli_overrides, + ) + return make_run_from_cli(cfg) + + +def _parse_identity(parsed: ParsedGraph) -> dict[str, Any]: + """The run-visible parse identity the two dispatch routes must agree on.""" + records = {"": parsed.graph, **parsed.graphs} + return { + "trace_ids": [t.id for t in parsed.traces], + "node_ids": {key: sorted(rec.nodes) for key, rec in records.items()}, + "prompt_segment_ids": { + key: {nid: read_prompt_segment_ids(n) for nid, n in rec.nodes.items()} + for key, rec in records.items() + }, + "pool_keys": sorted(parsed.segment_pool._by_id) + if parsed.segment_pool is not None + else None, + } + + +@pytest.mark.parametrize( + ("path", "fmt", "cli_overrides"), + [ + param(WEKA_MIN, "weka_trace", {"random_seed": 7}, id="weka-explicit-seed"), + param(DYNAMO_NESTED, "dynamo_trace", {}, id="dynamo"), + param(DAG_SMALL, "dag_jsonl", {"graph_format": "dag_jsonl"}, id="dag"), + ], +) # fmt: skip +def test_registry_dispatch_with_ctx_matches_run_parse( + path: Path, fmt: str, cli_overrides: dict[str, Any] +) -> None: + run = _run_for(path, **cli_overrides) + + via_run = parse_graph_workload(run, path) + adapter_cls = plugins.get_class(PluginType.GRAPH_ADAPTER, fmt) + via_registry = adapter_cls.parse(path, resolve_graph_parse_context(run)) + + run_identity = _parse_identity(via_run) + registry_identity = _parse_identity(via_registry) + assert run_identity["pool_keys"], "parse must carry real content segments" + assert run_identity == registry_identity diff --git a/tests/unit/graph/test_runtime_hardening_round3.py b/tests/unit/graph/test_runtime_hardening_round3.py new file mode 100644 index 0000000000..be2d382f42 --- /dev/null +++ b/tests/unit/graph/test_runtime_hardening_round3.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Runtime-hardening unit checks: recycle stop-condition naming + unknown-return warning. + +* ``_recycle_has_stop_condition`` returns True iff a stop condition exists, + and the ``recycle_is_bounded`` var matches its meaning. We pin the boolean + contract so a naming change stays behavior-neutral. +* ``_on_graph_return`` reports a dropped unknown-trace-id return as a + WARNING carrying the instance id + node ordinal (not a debug no-op), so an + orphaned-adapter condition is field-diagnosable. We pin the level + payload. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Any + +import pytest + + +@dataclass +class _Cfg: + phase: Any = None + concurrency: int = 1 + expected_num_sessions: int | None = None + total_expected_requests: int | None = None + expected_duration_sec: float | None = None + + +@dataclass +class _Lifecycle: + _left: float | None = None + + def time_left_in_seconds(self) -> float | None: + return self._left + + +def _minimal_strategy(config: _Cfg, lifecycle: Any = None): + """A strategy built over an EMPTY corpus -- enough to exercise the gating + helpers and the return router without any executor run.""" + from aiperf.dataset.graph.models import GraphRecord, ParsedGraph + from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy + + parsed = ParsedGraph(graph=GraphRecord(nodes={}, edges=[], state={}), traces=[]) + + class _Issuer: + async def issue_graph_credit(self, turn: Any) -> bool: + return True + + def mark_graph_sending_complete(self) -> None: ... + def graph_all_returned(self) -> bool: + return True + + def set_graph_all_returned_event(self) -> None: ... + + return GraphIRReplayStrategy( + config=config, + credit_issuer=_Issuer(), + parsed_graph=parsed, + register_observer=lambda obs: None, + lifecycle=lifecycle, + ) + + +# --- recycle stop-condition naming ------------------------------------------ + + +@pytest.mark.parametrize( + "config,lifecycle,expected", + [ + (_Cfg(), None, False), + (_Cfg(expected_num_sessions=5), None, True), + (_Cfg(total_expected_requests=10), None, True), + (_Cfg(expected_duration_sec=30.0), None, True), + (_Cfg(), _Lifecycle(_left=12.0), True), + (_Cfg(), _Lifecycle(_left=None), False), + ], +) # fmt: skip +def test_recycle_has_stop_condition_matches_name(config, lifecycle, expected): + """``_recycle_has_stop_condition`` returns True iff a stop condition EXISTS, + so the name matches the meaning (an inverse-named predicate here is a + double-negation trap for callers).""" + strategy = _minimal_strategy(config, lifecycle) + assert strategy._recycle_has_stop_condition() is expected + + +def test_resolve_lane_count_clamps_to_corpus_when_unbounded(): + """``recycle_is_bounded=False`` (no stop condition) clamps lanes to the corpus + size; True leaves the full concurrency fan-out.""" + strategy = _minimal_strategy(_Cfg(concurrency=8)) + # Unbounded: single corpus pass -> clamp to total traces. + assert strategy._resolve_lane_count(3, recycle_is_bounded=False) == 3 + # Bounded: sustain full concurrency even past the corpus size. + assert strategy._resolve_lane_count(3, recycle_is_bounded=True) == 8 + + +# --- unknown-return warning --------------------------------------------------- + + +@dataclass +class _Credit: + trace_id: str | None + node_ordinal: int | None = None + x_correlation_id: str = "x" + turn_index: int = 0 + + +def test_unknown_trace_return_logs_warning_with_instance_and_ordinal(caplog): + """A return for an instance id with no live adapter must WARN (not debug), + naming the instance id + node ordinal so the drop is field-diagnosable.""" + strategy = _minimal_strategy(_Cfg()) + with caplog.at_level(logging.WARNING): + strategy._on_graph_return( + _Credit(trace_id="t-1#0.0", node_ordinal=7), error=None, cancelled=False + ) + warnings = [r for r in caplog.records if r.levelno >= logging.WARNING] + assert warnings, "an unknown-instance return must log at WARNING" + msg = warnings[-1].getMessage() + assert "t-1#0.0" in msg + assert "7" in msg + + +def test_none_trace_id_return_is_silent_noop(caplog): + """A credit with no trace_id is not a graph return at all -> silent no-op, + NOT a warning (only an UNKNOWN graph instance id warns).""" + strategy = _minimal_strategy(_Cfg()) + with caplog.at_level(logging.WARNING): + strategy._on_graph_return(_Credit(trace_id=None), error=None, cancelled=False) + assert [r for r in caplog.records if r.levelno >= logging.WARNING] == [] diff --git a/tests/unit/graph/test_slot_graph_eager_drain.py b/tests/unit/graph/test_slot_graph_eager_drain.py new file mode 100644 index 0000000000..f572915acf --- /dev/null +++ b/tests/unit/graph/test_slot_graph_eager_drain.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Every non-weka graph takes the in-process interned drain. + +``GraphStoreBuilder._build_graph_store_streaming`` routes only ``weka_trace`` through the +worker-pool payload/trie drain; every other format (here: ``native``) parses +once in-process and drains that SAME parse through +``build_unified_trie_store_interned`` -- the interned drain. That is true for +slot-carrying ``@channel`` graphs (whose assembly items/capture the streaming +payload envelope cannot carry) AND for plain slot-free graphs (for which the +payload round trip is pure overhead in-process). These tests pin that route: +both shapes return the FULL parse and never touch the weka trie drain, and the +persisted store carries the real (including dynamic-slot) envelopes. +``graph_carries_assembly_slots`` is retained for the ``workload_detect`` +t*-gate, so its detection is pinned here too. +""" + +from __future__ import annotations + +from pathlib import Path +from types import MethodType, SimpleNamespace +from typing import Any + +import orjson +import pytest + +from aiperf.dataset.graph.parser import parse_native +from aiperf.dataset.graph.segment_ir.store_builder import ( + graph_carries_assembly_slots, +) +from aiperf.dataset.graph.store_build import GraphStoreBuilder +from aiperf.dataset.graph_segment_unified_store import GraphSegmentUnifiedClient + +SLOT_GRAPH = """ +graph: + nodes: + plan: + prompt: [{role: user, content: "Make a plan."}] + output: plan_out + review: + prompt: + - role: user + content: ["Review this plan: ", "@plan_out"] + output: review_out + edges: + - {source: START, target: plan} + - {source: plan, target: review} + - {source: review, target: END} +traces: + - id: t1 +""" + +STATIC_GRAPH = """ +graph: + nodes: + plan: + prompt: [{role: user, content: "Make a plan."}] + output: plan_out + edges: + - {source: START, target: plan} + - {source: plan, target: END} +traces: + - id: t1 +""" + + +def _parse(tmp_path: Path, yaml_text: str): + p = tmp_path / "workload.yaml" + p.write_text(yaml_text) + return p, parse_native(p) + + +def _interned_stub() -> SimpleNamespace: + """Just the attributes the non-weka interned branch reads from self, with + the REAL interned-drain helpers bound and a weka trie drain that fails + loudly if reached (no non-weka format may ever take it).""" + stub = SimpleNamespace( + run=SimpleNamespace(benchmark_id="bench"), + info=lambda *a, **k: None, + warning=lambda *a, **k: None, + _sidecar_path=None, + ) + stub._build_interned_unified_store = MethodType( + GraphStoreBuilder._build_interned_unified_store, stub + ) + stub._write_graph_sidecar = MethodType(GraphStoreBuilder._write_graph_sidecar, stub) + + async def _fail_trie(payloads: Any, base_path: Path) -> None: + raise AssertionError("non-weka graph must not take the weka trie payload drain") + + stub._build_graph_store_streaming_trie = _fail_trie + return stub + + +def test_graph_carries_assembly_slots_true_for_slot_graph(tmp_path: Path) -> None: + _, parsed = _parse(tmp_path, SLOT_GRAPH) + assert graph_carries_assembly_slots(parsed) is True + + +def test_graph_carries_assembly_slots_false_for_static_graph(tmp_path: Path) -> None: + _, parsed = _parse(tmp_path, STATIC_GRAPH) + assert graph_carries_assembly_slots(parsed) is False + + +@pytest.mark.asyncio +async def test_build_graph_store_streaming_slot_graph_takes_interned_drain( + tmp_path: Path, monkeypatch +) -> None: + """A slot-carrying native parse drains through the interned builder (slot + envelope persisted) and the SAME parse is returned -- never the trie drain.""" + from aiperf.dataset.graph import workload_detect + + graph_path, parsed = _parse(tmp_path, SLOT_GRAPH) + monkeypatch.setattr( + workload_detect, "parse_graph_workload", lambda run, path: parsed + ) + stub = _interned_stub() + + catalog, returned = await GraphStoreBuilder._build_graph_store_streaming( + stub, graph_path, tmp_path, "native" + ) + + assert returned is parsed + assert set(catalog) == {"t1"} + with GraphSegmentUnifiedClient(tmp_path, "bench").open() as client: + review_bytes = client.get_node_envelope( + "t1", catalog["t1"]["review"], "profiling" + ) + plan_bytes = client.get_node_envelope("t1", catalog["t1"]["plan"], "profiling") + assert review_bytes is not None and plan_bytes is not None + # The dynamic-slot envelopes only the interned drain persists: the reader's + # composed message with the plan node's response slot, and the producer's + # capture flag. + review_envelope = orjson.loads(review_bytes) + assert review_envelope["items"] == [ + { + "m": { + "role": "user", + "parts": [ + {"t": "Review this plan: "}, + {"sv": catalog["t1"]["plan"]}, + ], + } + } + ] + assert orjson.loads(plan_bytes).get("capture") is True + + +@pytest.mark.asyncio +async def test_build_graph_store_streaming_static_graph_takes_interned_drain( + tmp_path: Path, monkeypatch +) -> None: + """A plain (slot-free) native parse ALSO takes the interned drain now: the + weka trie payload drain is weka-only, so the flip routes every non-weka + format -- slot-carrying or not -- through the interned builder and returns + the FULL parse.""" + from aiperf.dataset.graph import workload_detect + + graph_path, parsed = _parse(tmp_path, STATIC_GRAPH) + monkeypatch.setattr( + workload_detect, "parse_graph_workload", lambda run, path: parsed + ) + stub = _interned_stub() + + catalog, returned = await GraphStoreBuilder._build_graph_store_streaming( + stub, graph_path, tmp_path, "native" + ) + + assert returned is parsed + assert set(catalog) == {"t1"} + # The interned store the worker opens carries the plain node's envelope. + with GraphSegmentUnifiedClient(tmp_path, "bench").open() as client: + plan_bytes = client.get_node_envelope("t1", catalog["t1"]["plan"], "profiling") + assert plan_bytes is not None diff --git a/tests/unit/graph/test_start_anchor_edges.py b/tests/unit/graph/test_start_anchor_edges.py new file mode 100644 index 0000000000..2ac37d992f --- /dev/null +++ b/tests/unit/graph/test_start_anchor_edges.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Start-anchor post-pass over interval-order edges. + +``apply_start_anchors`` replaces an overlapped node's interval-order edges with a +single start-anchored edge (``delay_after_predecessor_start_us``) when the node's +``causal_parent_id`` names a parent that is still IN FLIGHT at the node's recorded +start. Nodes whose causal parent had already finished keep their interval-order +edges, and nodes with no / unknown causal parent are untouched. Also covers the +``with_fan_in_inputs`` exclusion of start-anchored edges. +""" + +import pytest + +from aiperf.dataset.graph.models import LlmNode, StaticEdge +from aiperf.dataset.graph.segment_ir.interval_order import ( + apply_start_anchors, + build_interval_edges, + compute_ranks, +) +from aiperf.dataset.graph.segment_ir.trie_content import ( + TrieNode, + TrieRequest, + with_fan_in_inputs, +) + + +def _node(nid, t, api, causal=None): + n = TrieNode( + node_id=nid, + request=TrieRequest( + hash_ids=[], input_length=64, output_length=8, t=t, api_time=api + ), + order=0, + causal_parent_id=causal, + ) + n.warped_start = t + return n + + +def test_overlapped_child_gets_single_start_anchored_edge(): + p = _node("p", 0.0, 8.0) + c = _node("c", 2.5, 1.0, causal="p") + nodes = [p, c] + compute_ranks(nodes) + edges = build_interval_edges(nodes) + apply_start_anchors(nodes, edges) + (e,) = edges["c"] + assert e.source == "p" + assert e.delay_after_predecessor_start_us == pytest.approx(2.5e6) + assert e.delay_after_predecessor_us is None + + +def test_non_overlapped_child_keeps_interval_edges(): + p = _node("p", 0.0, 1.0) + c = _node("c", 2.5, 1.0, causal="p") # p ended at 1.0 < 2.5: no overlap + nodes = [p, c] + compute_ranks(nodes) + edges = build_interval_edges(nodes) + before = list(edges["c"]) + apply_start_anchors(nodes, edges) + assert edges["c"] == before + assert edges["c"][0].delay_after_predecessor_start_us is None + + +def test_missing_or_unknown_causal_parent_untouched(): + a = _node("a", 0.0, 1.0) + b = _node("b", 0.5, 1.0, causal="ghost") + nodes = [a, b] + compute_ranks(nodes) + edges = build_interval_edges(nodes) + before = {k: list(v) for k, v in edges.items()} + apply_start_anchors(nodes, edges) + assert edges == before + + +def test_delay_computed_on_warped_clock(): + p = _node("p", 100.0, 8.0) + c = _node("c", 103.0, 1.0, causal="p") + p.warped_start, c.warped_start = 10.0, 13.0 # warp compressed by 90s + nodes = [p, c] + compute_ranks(nodes) + edges = build_interval_edges(nodes) + apply_start_anchors(nodes, edges) + assert edges["c"][0].delay_after_predecessor_start_us == pytest.approx(3.0e6) + + +def test_fan_in_inputs_skip_start_anchored_edges(): + llm = LlmNode(prompt=[], output="c_out") + edges = [ + StaticEdge(source="p", target="c", delay_after_predecessor_start_us=1e6), + StaticEdge(source="q", target="c", delay_after_predecessor_us=0.0), + ] + out = with_fan_in_inputs(llm, edges) + assert [r.channel for r in out.inputs] == ["q_out"] diff --git a/tests/unit/graph/test_start_anchor_fidelity_tool.py b/tests/unit/graph/test_start_anchor_fidelity_tool.py new file mode 100644 index 0000000000..b8db9b84ff --- /dev/null +++ b/tests/unit/graph/test_start_anchor_fidelity_tool.py @@ -0,0 +1,264 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The offline fidelity proof validates START-ANCHORED nodes as dispatch-to-dispatch +offsets, not as zero end-to-start delays. + +Task 3 taught the weka trie builder to collapse a mid-flight spawn / chain-overlap +node's incoming edges into a single ``StaticEdge.delay_after_predecessor_start_us`` +(the warped start-to-start gap). This locks the matching change to +``tools/weka_trace_fidelity.py``: :func:`build_recorded_trace` must route those +edges into ``_RecordedNode.start_anchor`` (NOT into the end-to-start +``predecessors`` / ``pred_delay_us``), and +:func:`causality_timing_vs_real_trace` must compare each start-anchored child's +OBSERVED ``request_start_ns`` gap from its PARENT'S dispatch against that warped +delay -- so a run that places the child at parent-dispatch + delay passes and one +that shifts it off does not. + +The fixture geometry mirrors ``_OVERLAP_TRACE`` in +``tests/unit/graph/test_start_anchor_runtime.py``: +``START->start_anchor:0``; ``start_anchor:0->a1:0`` start-delay 2.5s; +``start_anchor:0->start_anchor:1`` start-delay 5.0s; +``start_anchor:0->start_anchor:2`` end-delay 1.0s ++ ``start_anchor:1->start_anchor:2`` end-delay 0. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from tools.weka_trace_fidelity import ( + build_recorded_trace, + causality_timing_vs_real_trace, +) +from tools.weka_trie_timing_sim import main as timing_sim_main +from tools.weka_trie_timing_sim import simulate_trace + +# P: t=0 api=8.0 (long, spawner); C: subagent first at t=2.5 (P in flight); +# Q: chain-overlap at t=5.0 (P in flight); R: t=9.0 (after P ends, end-anchored). +# Byte-identical to ``test_start_anchor_runtime._OVERLAP_TRACE``. +_OVERLAP_TRACE = { + "id": "start_anchor", "models": ["M"], "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "n", "model": "M", "in": 128, "out": 64, + "hash_ids": [1, 2], "api_time": 8.0, "stop": "tool_use"}, + {"t": 2.0, "type": "subagent", "agent_id": "a1", + "subagent_type": "Explore", "status": "completed", "models": ["M"], + "requests": [ + {"t": 2.5, "type": "n", "model": "M", "in": 128, "out": 16, + "hash_ids": [50, 51], "api_time": 1.0}, + ]}, + {"t": 5.0, "type": "n", "model": "M", "in": 192, "out": 32, + "hash_ids": [1, 2, 3], "api_time": 1.0}, + {"t": 9.0, "type": "n", "model": "M", "in": 256, "out": 32, + "hash_ids": [1, 2, 3, 4], "api_time": 0.5}, + ], +} # fmt: skip + +_P = "start_anchor:0" +_C = "a1:0" +_Q = "start_anchor:1" +_R = "start_anchor:2" + +_S = 1_000_000_000 # 1e9 ns per second +_ORIGIN_NS = 1 * _S # arbitrary fresh run-origin (absolute wall-clock differs) + + +def _write_trace(tmp_path: Path) -> Path: + """Materialize ``_OVERLAP_TRACE`` to a JSON file the tool can rebuild from.""" + path = tmp_path / "start_anchor.json" + path.write_text(json.dumps(_OVERLAP_TRACE)) + return path + + +def _record(node_id: str, request_start_ns: int) -> dict: + """One raw-export JSONL line for a profiling dispatch of ``node_id``. + + The tool recovers the node id from ``x_request_id`` + (``{node_id}::{nonce}``, worker-minted) and selects the trace by the + ``conversation_id`` base. + """ + return { + "metadata": { + "conversation_id": "start_anchor#0", + "x_request_id": f"{node_id}::deadbeefdeadbeefdeadbeefdeadbeef", + "benchmark_phase": "profiling", + "request_start_ns": request_start_ns, + "credit_issued_ns": None, + }, + "payload": {"messages": []}, + } + + +def _write_raw(tmp_path: Path, dispatch_ns: dict[str, int]) -> Path: + """Write a raw-export JSONL placing each node at its given absolute dispatch ns.""" + path = tmp_path / "profile_export_raw.jsonl" + lines = [json.dumps(_record(nid, ns)) for nid, ns in dispatch_ns.items()] + path.write_text("\n".join(lines) + "\n") + return path + + +# A faithful replay against the zero-latency mock: P at the origin, the two +# start-anchored children at parent-dispatch + their warped start-delay, and the +# end-anchored R at its binding pred's (start_anchor:1) observed dispatch + delay 0. +_FAITHFUL_NS = { + _P: _ORIGIN_NS, + _C: _ORIGIN_NS + int(2.5 * _S), + _Q: _ORIGIN_NS + int(5.0 * _S), + _R: _ORIGIN_NS + int(5.0 * _S), +} + + +def test_build_recorded_trace_captures_start_anchor(tmp_path: Path) -> None: + """Start-anchored edges land in ``start_anchor``, NOT in end-to-start preds. + + The overlap children ``a1:0`` / ``start_anchor:1`` each carry one + ``delay_after_predecessor_start_us`` edge off ``start_anchor:0``; the builder + must record ``(start_anchor:0, warped_delay_us)`` and keep them out of + ``predecessors`` so they are never mistimed as a zero end-to-start wait. + """ + trace_file = _write_trace(tmp_path) + recorded = build_recorded_trace(trace_file, idle_gap_cap_seconds=60.0) + + # start_anchor:0 is non-streaming (type "n"), so both overlap children are no-ttft-parent + # start anchors -- the third (first-token delay) slot is None. + assert recorded.nodes[_C].start_anchor == (_P, 2.5e6, None) + assert recorded.nodes[_C].predecessors == [] + assert recorded.nodes[_Q].start_anchor == (_P, 5.0e6, None) + assert recorded.nodes[_Q].predecessors == [] + # R stays a normal end-anchored AND-join (its edges are NOT start-anchored). + assert recorded.nodes[_R].start_anchor is None + assert sorted(recorded.nodes[_R].predecessors) == [_P, _Q] + + +def test_faithful_start_anchored_replay_passes(tmp_path: Path) -> None: + """A replay placing C exactly 2.5s and Q 5.0s after P's dispatch PASSES. + + The proof compares each start-anchored child's observed dispatch gap from its + PARENT'S dispatch against the warped start-to-start delay, and counts the two + edges as ``exact``. + """ + trace_file = _write_trace(tmp_path) + raw = _write_raw(tmp_path, _FAITHFUL_NS) + + report = causality_timing_vs_real_trace(raw, trace_file, idle_gap_cap_seconds=60.0) + + assert report.passed, report.render() + # The two start-anchored edges (C, Q) are validated as exact start-to-start. + assert report.exact_edges >= 2 + + +def test_shifted_start_anchored_child_fails(tmp_path: Path) -> None: + """Shifting C's dispatch +1.0s (to 3.5s after P) FAILS -- the proof actually + checks start-anchored timing rather than skipping those nodes. + + +1.0s exceeds the abs tolerance (0.75s) and the relative tolerance + (0.15 * 2.5s = 0.375s), so the start-anchor comparison must flag it. + """ + trace_file = _write_trace(tmp_path) + shifted = dict(_FAITHFUL_NS) + shifted[_C] = _FAITHFUL_NS[_C] + int(1.0 * _S) + raw = _write_raw(tmp_path, shifted) + + report = causality_timing_vs_real_trace(raw, trace_file, idle_gap_cap_seconds=60.0) + + assert not report.passed + assert any(_C in m.where for m in report.mismatches), report.render() + + +def test_timing_sim_tool_reconstructs_overlap_timeline(tmp_path: Path) -> None: + """The sibling ``weka_trie_timing_sim`` simulator gates start-anchored edges off + the predecessor's DISPATCH, so it reconstructs the overlap timeline byte-exact. + + Before Task 7 the simulator only knew end-to-start / START edges, so a + start-anchored child fell through to ``sim_end(parent) + 0`` and landed at the + parent's completion (a1:0 at 8.0s, not 2.5s) -- a divergence. With the + dispatch-anchored branch every node reconstructs its recorded warped start + (0.0 / 2.5 / 5.0 / 9.0), so all four check exact with zero divergences. + """ + trace_file = _write_trace(tmp_path) + + checked, exact, diverged, first_token_edges = simulate_trace( + trace_file, cap=60.0, tol=1e-3 + ) + + assert diverged == [] + assert checked == exact == 4 + # _OVERLAP_TRACE's spawner (start_anchor:0) is non-streaming, so no first-token edges. + assert first_token_edges == 0 + + +# The start-anchored child (``a:0``) ties its parent's (``tie_order:2``) recorded +# start, and its node-id STRING sorts before the parent's ("a:0" < "tie_order:2"): +# the old ``(recorded_start, node_id)`` processing order simulated the child first +# and zero-defaulted the parent's dispatch, landing the child at t=0 instead of +# t=2.0 (a false divergence). Under the ``{scope}:{turn}`` id scheme the child's id +# derives from its subagent ``agent_id`` ("a"), so the empty padding markers no +# longer affect the id -- they are kept inert only to preserve the fixture shape. +_TIE_TRACE = { + "id": "tie_order", "models": ["M"], "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "n", "model": "M", "in": 64, "out": 8, + "hash_ids": [1], "api_time": 0.5}, + {"t": 1.0, "type": "n", "model": "M", "in": 128, "out": 8, + "hash_ids": [1, 2], "api_time": 0.5}, + {"t": 2.0, "type": "n", "model": "M", "in": 192, "out": 64, + "hash_ids": [1, 2, 3], "api_time": 10.0, "stop": "tool_use"}, + *( + {"t": 2.0, "type": "subagent", "agent_id": f"pad{i}", + "subagent_type": "X", "status": "completed", "models": ["M"], + "requests": []} + for i in range(3, 10) + ), + {"t": 2.0, "type": "subagent", "agent_id": "a", + "subagent_type": "X", "status": "completed", "models": ["M"], + "requests": [ + {"t": 2.0, "type": "n", "model": "M", "in": 64, "out": 8, + "hash_ids": [50], "api_time": 1.0}, + ]}, + ], +} # fmt: skip + + +def test_timing_sim_dependency_order_beats_lexicographic_tie( + tmp_path: Path, +) -> None: + """The simulator processes a start-anchored parent BEFORE its child even when + the child's node-id string sorts first at an identical recorded start. + + ``a:0`` start-anchors to ``tie_order:2`` (in flight, delay 0) and ties its + recorded start; dependency (topological) order must gate it at the parent's + simulated dispatch (2.0s), not at a zero-defaulted 0.0s -- all four nodes + reconstruct the recorded warped timeline exactly. + """ + trace_file = tmp_path / "tie_order.json" + trace_file.write_text(json.dumps(_TIE_TRACE)) + + checked, exact, diverged, first_token_edges = simulate_trace( + trace_file, cap=60.0, tol=1e-3 + ) + + assert diverged == [], diverged + assert checked == exact == 4 + assert first_token_edges == 0 + + +def test_timing_sim_main_vacuous_empty_dir_exits_nonzero( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + """A dir with no ``*.json`` trace files is a VACUOUS simulation -> exit 1. + + A simulation that checked nothing must not exit green (the exit code is the + acceptance signal in CI pipelines). + """ + empty = tmp_path / "empty" + empty.mkdir() + + rc = timing_sim_main([str(empty)]) + + assert rc == 1 + assert "VACUOUS: nothing checked" in capsys.readouterr().out diff --git a/tests/unit/graph/test_start_anchor_runtime.py b/tests/unit/graph/test_start_anchor_runtime.py new file mode 100644 index 0000000000..8cade4ae30 --- /dev/null +++ b/tests/unit/graph/test_start_anchor_runtime.py @@ -0,0 +1,291 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Runtime proof: start-anchored edges dispatch children at the predecessor's +DISPATCH instant (firing-gate clear), not its completion. + +Drives the REAL ``TraceExecutor`` on a ``VirtualClock`` over weka trie graphs +whose overlap geometry produces ``StaticEdge.delay_after_predecessor_start_us`` +edges (Task 3 stamping). The harness (``_STUB_CALLBACKS``, ``_VTimeIssuer``, +``_drive_virtual``) mirrors ``test_executor_runs_weka.py``. + +Contract: +1. Recorded-speed replay dispatches at exactly the recorded start instants. +2. A slowed parent keeps DISPATCH-anchored children fixed and moves the + end-anchored child to the parent's (later) finish, with NO cycle RuntimeError + even though a dispatch-anchored child finished long before the parent. +3. A delayed parent DISPATCH shifts a dispatch-anchored child by exactly the + same amount (dispatch-relative delay preserved). +4. ``AIPERF_GRAPH_IGNORE_EDGE_DELAYS`` short-circuits the gate uniformly while + dispatch-time scheduling still fires every child exactly once. +""" + +import asyncio +from typing import Any + +import msgspec +import pytest + +from aiperf.common.clock import VirtualClock +from aiperf.common.environment import Environment +from aiperf.dataset.graph.adapters.weka.trace_models import WekaTrace +from aiperf.dataset.graph.adapters.weka.trie_build import ( + ReconCallbacks, + build_trie_graph, +) +from aiperf.dataset.graph.models import TraceRecord +from aiperf.graph.executor import TraceExecutor + +_BLOCK_SIZE = 64 + + +def _stub_decode_block_tokens(hash_ids: list[int]) -> list[int]: + out: list[int] = [] + for h in hash_ids: + out.extend(range(h * 100, h * 100 + _BLOCK_SIZE)) + return out + + +def _stub_partial_tail_tokens(n_tokens: int, seed: str) -> list[int]: + base = sum(ord(c) for c in seed) * 1000 + return list(range(base, base + n_tokens)) + + +def _stub_decode_tokens_to_text(tokens: list[int]) -> str: + return "|".join(str(t) for t in tokens) + + +_STUB_CALLBACKS = ReconCallbacks( + decode_block_tokens=_stub_decode_block_tokens, + sample_partial_tail_tokens=_stub_partial_tail_tokens, + decode_tokens_to_text=_stub_decode_tokens_to_text, +) + + +class _VTimeIssuer: + """Records each node's VIRTUAL dispatch time, then consumes the node's + recorded ``api_time`` in virtual time. + + ``api_by_id`` maps ``request.node_id`` -> recorded processing seconds. The + executor records the node's finish as ``dispatch_start + api_time`` (the + issuer's virtual sleep), so an end-anchored successor's firing gate clears at + the recorded end-to-start instant while a start-anchored successor clears at + ``dispatch + delay``. + """ + + def __init__(self, clock: Any, api_by_id: dict[str, float]) -> None: + self._clock = clock + self._api = api_by_id + self.dispatched_at: dict[str, float] = {} + self.dispatched: list[str] = [] + + async def dispatch( + self, + node: Any, + request: Any, + ctx: Any, + **kwargs: Any, + ) -> str: + nid = request.node_id + self.dispatched.append(nid) + self.dispatched_at[nid] = self._clock.now_ns() / 1e9 + api_s = self._api.get(nid, 0.0) + if api_s > 0.0: + await self._clock.sleep_ns(int(api_s * 1e9)) + return f"placeholder::{nid}" + + +async def _drive_virtual(clock: Any, task: Any) -> Any: + """Pump a virtual-clock replay: drain ready callbacks, then fast-forward sim + time to the earliest parked waiter whenever the loop goes idle.""" + loop = asyncio.get_running_loop() + ready = loop._ready # noqa: SLF001 -- idle detection for the pump + while not task.done(): + while ready and not task.done(): + await asyncio.sleep(0) + if task.done(): + break + nxt = clock.peek_min_waiter_ns() + if nxt is None: + await asyncio.sleep(0) + if not ready and clock.peek_min_waiter_ns() is None and not task.done(): + raise RuntimeError("virtual-time replay stalled") + continue + await clock.advance_to(nxt) + return task.result() + + +# --- fixtures ------------------------------------------------------------- + +# P: t=0 api=8.0 (long, spawner); C: subagent first at t=2.5 (P in flight); +# Q: chain-overlap at t=5.0 (P in flight); R: t=9.0 (after P ends, end-anchored) +_OVERLAP_TRACE = { + "id": "start_anchor", "models": ["M"], "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "n", "model": "M", "in": 128, "out": 64, + "hash_ids": [1, 2], "api_time": 8.0, "stop": "tool_use"}, + {"t": 2.0, "type": "subagent", "agent_id": "a1", + "subagent_type": "Explore", "status": "completed", "models": ["M"], + "requests": [ + {"t": 2.5, "type": "n", "model": "M", "in": 128, "out": 16, + "hash_ids": [50, 51], "api_time": 1.0}, + ]}, + {"t": 5.0, "type": "n", "model": "M", "in": 192, "out": 32, + "hash_ids": [1, 2, 3], "api_time": 1.0}, + {"t": 9.0, "type": "n", "model": "M", "in": 256, "out": 32, + "hash_ids": [1, 2, 3, 4], "api_time": 0.5}, + ], +} # fmt: skip + +# Node ids the trie build assigns to _OVERLAP_TRACE (verified against the graph). +_P = "start_anchor:0" +_C = "a1:0" +_Q = "start_anchor:1" +_R = "start_anchor:2" + +# P1 (t=0, api=1.0), P2 (t=2.0, api=8.0, tool_use), subagent at t=3.0 with C +# (t=4.0, api=1.0): C anchors to P2's DISPATCH with D=2.0; P2 anchors to P1's +# finish with recorded end-to-start delay 1.0. +_CHAIN_TRACE = { + "id": "start_anchor_chain", "models": ["M"], "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "n", "model": "M", "in": 128, "out": 64, + "hash_ids": [1, 2], "api_time": 1.0}, + {"t": 2.0, "type": "n", "model": "M", "in": 192, "out": 32, + "hash_ids": [1, 2, 3], "api_time": 8.0, "stop": "tool_use"}, + {"t": 3.0, "type": "subagent", "agent_id": "a1", + "subagent_type": "Explore", "status": "completed", "models": ["M"], + "requests": [ + {"t": 4.0, "type": "n", "model": "M", "in": 128, "out": 16, + "hash_ids": [50, 51], "api_time": 1.0}, + ]}, + ], +} # fmt: skip + +# Node ids the trie build assigns to _CHAIN_TRACE (verified against the graph). +_P1 = "start_anchor_chain:0" +_P2 = "start_anchor_chain:1" +_CHAIN_C = "a1:0" + + +def _build(raw: dict) -> Any: + """Build a trie graph from a raw weka trace dict and wire a bare trace.""" + trace = WekaTrace.model_validate(raw) + parsed, pool = build_trie_graph(trace, callbacks=_STUB_CALLBACKS) + bare = TraceRecord(id=trace.id) + return msgspec.structs.replace(parsed, traces=[bare], segment_pool=pool) + + +async def _run_virtual(parsed: Any, api_by_id: dict[str, float]) -> _VTimeIssuer: + """Drive ``parsed``'s traces on a VirtualClock and return the issuer.""" + clock = VirtualClock() + issuer = _VTimeIssuer(clock, api_by_id) + executor = TraceExecutor(parsed, credit_issuer=issuer, clock=clock) + + async def _phase() -> None: + async with asyncio.TaskGroup(): + for trace in parsed.traces: + await executor.run(trace) + + phase_task = asyncio.ensure_future(_phase()) + await _drive_virtual(clock, phase_task) + return issuer + + +# --- tests ---------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_recorded_speed_virtual_replay_matches_recorded_starts(): + """At recorded speed the four nodes dispatch at exactly 0.0/2.5/5.0/9.0s. + + C and Q are start-anchored to P's dispatch (0.0) at +2.5 / +5.0. R is + end-anchored: it waits for P's finish (0.0 + api 8.0) + recorded delay 1.0. + """ + parsed = _build(_OVERLAP_TRACE) + api_by_id = {_P: 8.0, _C: 1.0, _Q: 1.0, _R: 0.5} + + issuer = await _run_virtual(parsed, api_by_id) + + assert set(issuer.dispatched_at) == {_P, _C, _Q, _R} + assert issuer.dispatched_at[_P] == pytest.approx(0.0, abs=1e-3) + assert issuer.dispatched_at[_C] == pytest.approx(2.5, abs=1e-3) + assert issuer.dispatched_at[_Q] == pytest.approx(5.0, abs=1e-3) + assert issuer.dispatched_at[_R] == pytest.approx(9.0, abs=1e-3) + + +@pytest.mark.asyncio +async def test_slowed_parent_moves_children_with_dispatch_not_wall_clock(): + """Doubling P's server time keeps DISPATCH-anchored children fixed. + + P inflated to api=16.0. C and Q are anchored to P's DISPATCH (0.0), so they + stay at 2.5 / 5.0. R is end-anchored: it fires at P's finish (16.0) + + recorded end-to-start delay (1.0) = 17.0. + + Crucially C finishes at 3.5 -- long before P finishes at 16.0. Start-anchored + successors are excluded from ``successors_after``, so P's completion does NOT + re-schedule C into the cycle guard. Reaching the asserts proves no + RuntimeError was raised. + """ + parsed = _build(_OVERLAP_TRACE) + api_by_id = {_P: 16.0, _C: 1.0, _Q: 1.0, _R: 0.5} + + issuer = await _run_virtual(parsed, api_by_id) + + assert issuer.dispatched_at[_C] == pytest.approx(2.5, abs=1e-3) + assert issuer.dispatched_at[_Q] == pytest.approx(5.0, abs=1e-3) + assert issuer.dispatched_at[_R] == pytest.approx(17.0, abs=1e-3) + + +@pytest.mark.asyncio +async def test_delayed_parent_dispatch_shifts_children(): + """Slowing P1 pushes P2's DISPATCH later; the start-anchored child moves with + it by exactly the same amount (dispatch-relative delay preserved). + + Recorded: P1=0->finish 1.0; P2 dispatch = 1.0 + delay 1.0 = 2.0; C = + P2 dispatch 2.0 + D 2.0 = 4.0. Inflating P1 api to 5.0 moves P2's finish to + 5.0, so P2 dispatches at 5.0 + 1.0 = 6.0 and C at 6.0 + 2.0 = 8.0 -- C shifts + by the same +4.0 as P2. + """ + parsed = _build(_CHAIN_TRACE) + api_by_id = {_P1: 5.0, _P2: 8.0, _CHAIN_C: 1.0} + + issuer = await _run_virtual(parsed, api_by_id) + + assert set(issuer.dispatched_at) == {_P1, _P2, _CHAIN_C} + assert issuer.dispatched_at[_P1] == pytest.approx(0.0, abs=1e-3) + assert issuer.dispatched_at[_P2] == pytest.approx(6.0, abs=1e-3) + assert issuer.dispatched_at[_CHAIN_C] == pytest.approx(8.0, abs=1e-3) + + +@pytest.mark.asyncio +async def test_ignore_edge_delays_fires_children_at_parent_dispatch(monkeypatch): + """``AIPERF_GRAPH_IGNORE_EDGE_DELAYS`` collapses the gate but keeps + dispatch-time scheduling. + + With delays ignored, C and Q dispatch at the same virtual instant P + dispatches (0.0) -- scheduling still happens at P's dispatch, so P precedes + both in dispatch order -- and all four nodes dispatch exactly once. + """ + monkeypatch.setattr(Environment.GRAPH, "IGNORE_EDGE_DELAYS", True, raising=False) + + parsed = _build(_OVERLAP_TRACE) + api_by_id = {_P: 8.0, _C: 1.0, _Q: 1.0, _R: 0.5} + + issuer = await _run_virtual(parsed, api_by_id) + + from collections import Counter + + counts = Counter(issuer.dispatched) + assert set(counts) == {_P, _C, _Q, _R} + assert all(n == 1 for n in counts.values()), ( + f"every node must dispatch exactly once; got {counts}" + ) + assert issuer.dispatched_at[_P] == pytest.approx(0.0, abs=1e-3) + assert issuer.dispatched_at[_C] == pytest.approx(0.0, abs=1e-3) + assert issuer.dispatched_at[_Q] == pytest.approx(0.0, abs=1e-3) + # Scheduling still happens at P's dispatch, so P precedes both children. + order = issuer.dispatched + assert order.index(_P) < order.index(_C) + assert order.index(_P) < order.index(_Q) diff --git a/tests/unit/graph/test_start_anchor_snapshot.py b/tests/unit/graph/test_start_anchor_snapshot.py new file mode 100644 index 0000000000..8db01c55f9 --- /dev/null +++ b/tests/unit/graph/test_start_anchor_snapshot.py @@ -0,0 +1,184 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Start-anchored edges survive the snapshot + pacing consumers. + +Two consumers read ``StaticEdge.delay_after_predecessor_start_us`` (Task 3 +stamping) besides the executor (Task 5): + +1. ``chop_trie_at_tstar`` (``timing.snapshot_chop``, the segment-trie t* + frontier chop) keys on which SOURCES survive, not on the delay KIND -- so it + drops a start-anchored parent (re-rooting the child from START) and keeps a + surviving start-anchored edge VERBATIM without adding a competing synthetic + START anchor (no double-counted ``min_start_delay_us``). The ``test_chop_*`` + cases LOCK that behavior through the real chop path. +2. ``GraphIRReplayStrategy._max_inter_turn_gap_seconds`` scans every edge delay + for the idle-gap advisory; a start-anchored delay is an inter-turn gap too. + +Uses the overlap fixture from ``test_start_anchor_runtime.py`` (Task 5), copied +here so the module is self-contained. Geometry: ``START->start_anchor:0``; +``start_anchor:0->a1:0`` start-delay 2.5e6; ``start_anchor:0->start_anchor:1`` +start-delay 5.0e6; ``start_anchor:0->start_anchor:2`` end-delay 1.0e6 ++ ``start_anchor:1->start_anchor:2`` end-delay 0. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +from aiperf.dataset.graph.adapters.weka.trace_models import WekaTrace +from aiperf.dataset.graph.adapters.weka.trie_build import ( + ReconCallbacks, + build_trie_graph, +) +from aiperf.dataset.graph.models import ( + GraphRecord, + LlmNode, + ParsedGraph, + StaticEdge, + TraceRecord, +) +from aiperf.timing.snapshot_chop import chop_trie_at_tstar +from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy + +_BLOCK_SIZE = 64 + + +def _stub_decode_block_tokens(hash_ids: list[int]) -> list[int]: + out: list[int] = [] + for h in hash_ids: + out.extend(range(h * 100, h * 100 + _BLOCK_SIZE)) + return out + + +def _stub_partial_tail_tokens(n_tokens: int, seed: str) -> list[int]: + base = sum(ord(c) for c in seed) * 1000 + return list(range(base, base + n_tokens)) + + +def _stub_decode_tokens_to_text(tokens: list[int]) -> str: + return "|".join(str(t) for t in tokens) + + +_STUB_CALLBACKS = ReconCallbacks( + decode_block_tokens=_stub_decode_block_tokens, + sample_partial_tail_tokens=_stub_partial_tail_tokens, + decode_tokens_to_text=_stub_decode_tokens_to_text, +) + +# P: t=0 spawner; C: subagent first at t=2.5 (P in flight); Q: chain-overlap at +# t=5.0 (P in flight); R: t=9.0 (after P ends, end-anchored). Builds +# ``start_anchor:0 -> a1:0`` / ``start_anchor:0 -> start_anchor:1`` start-anchored +# edges + a ``start_anchor:1 -> start_anchor:2`` end-anchored survivor edge. +_OVERLAP_TRACE = { + "id": "start_anchor", "models": ["M"], "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "n", "model": "M", "in": 128, "out": 64, + "hash_ids": [1, 2], "api_time": 8.0, "stop": "tool_use"}, + {"t": 2.0, "type": "subagent", "agent_id": "a1", + "subagent_type": "Explore", "status": "completed", "models": ["M"], + "requests": [ + {"t": 2.5, "type": "n", "model": "M", "in": 128, "out": 16, + "hash_ids": [50, 51], "api_time": 1.0}, + ]}, + {"t": 5.0, "type": "n", "model": "M", "in": 192, "out": 32, + "hash_ids": [1, 2, 3], "api_time": 1.0}, + {"t": 9.0, "type": "n", "model": "M", "in": 256, "out": 32, + "hash_ids": [1, 2, 3, 4], "api_time": 0.5}, + ], +} # fmt: skip + + +# --- chop: LOCK the already-correct behavior (SOURCE survival, not delay kind) - + + +def test_chop_drops_parent_rewrites_child_to_absolute_start(): + """A dropped start-anchored parent re-roots its child from START. + + t* = 3.0s: P (arrival 0) and C (2.5) are warmup-dropped; Q (5.0) survives. + Q's only predecessor (P) is chopped, so Q re-roots from START at its + t*-relative absolute offset -- the synthetic START edge carries NO + start-anchored delay. + """ + parsed, _pool = build_trie_graph( + WekaTrace.model_validate(_OVERLAP_TRACE), callbacks=_STUB_CALLBACKS + ) + chopped = chop_trie_at_tstar(parsed, t_star_us=3_000_000) + (q_edge,) = [e for e in chopped.graph.edges if e.target == "start_anchor:1"] + assert q_edge.source == "START" + assert q_edge.min_start_delay_us == pytest.approx(2_000_000) + assert q_edge.delay_after_predecessor_start_us is None + + +def test_chop_t_star_zero_short_circuits_graph_unchanged(): + """t* <= 0 returns the parsed graph object UNCHANGED (documented contract).""" + parsed, _pool = build_trie_graph( + WekaTrace.model_validate(_OVERLAP_TRACE), callbacks=_STUB_CALLBACKS + ) + assert chop_trie_at_tstar(parsed, t_star_us=0) is parsed + + +def _shifted_requests( + requests: list[dict[str, Any]], dt: float +) -> list[dict[str, Any]]: + """Shift every recorded ``t`` (including nested subagent bodies) by ``dt``.""" + shifted: list[dict[str, Any]] = [] + for req in requests: + req = dict(req) + req["t"] = req["t"] + dt + if "requests" in req: + req["requests"] = _shifted_requests(req["requests"], dt) + shifted.append(req) + return shifted + + +def test_chop_keeps_surviving_start_anchored_edge_verbatim(): + """A surviving start-anchored edge is kept verbatim THROUGH the chop path. + + ``t* = 0`` short-circuits ``chop_trie_at_tstar`` entirely (graph returned + unchanged), which would leave the keep-verbatim branch unexecuted. Shifting + the trace +2s and chopping at t*=1s keeps every node while forcing the real + edge-recompute path to run: the start-anchored survivor edge must come out + unchanged, with no competing synthetic START anchor on its target (no + double-counted ``min_start_delay_us``). + """ + trace = dict( + _OVERLAP_TRACE, requests=_shifted_requests(_OVERLAP_TRACE["requests"], 2.0) + ) + parsed, _pool = build_trie_graph( + WekaTrace.model_validate(trace), callbacks=_STUB_CALLBACKS + ) + chopped = chop_trie_at_tstar(parsed, t_star_us=1_000_000) + assert chopped is not parsed, "t*>0 must run the real chop path" + assert set(chopped.graph.nodes) == set(parsed.graph.nodes), "nothing dropped" + + (c_edge,) = [e for e in chopped.graph.edges if e.target == "a1:0"] + assert c_edge.source == "start_anchor:0" + assert c_edge.delay_after_predecessor_start_us == pytest.approx(2.5e6) + assert c_edge.min_start_delay_us is None, ( + "surviving start-anchored edge must not gain a START-style offset" + ) + + +def test_max_inter_turn_gap_includes_start_anchored_delays(): + """A corpus whose only delay is start-anchored reports it as the max gap.""" + nodes = { + "A": LlmNode(prompt=["@a"], output="a"), + "B": LlmNode(prompt=["@b"], output="b"), + } + edges = [ + StaticEdge(source="START", target="A"), + StaticEdge(source="A", target="B", delay_after_predecessor_start_us=7_000_000), + ] + parsed = ParsedGraph( + graph=GraphRecord(nodes=nodes, edges=edges, state={}), + traces=[TraceRecord(id="t")], + ) + stub = SimpleNamespace(_parsed=parsed) + + gap_s = GraphIRReplayStrategy._max_inter_turn_gap_seconds(stub) + + assert gap_s == pytest.approx(7.0) diff --git a/tests/unit/graph/test_start_anchor_weka_stamping.py b/tests/unit/graph/test_start_anchor_weka_stamping.py new file mode 100644 index 0000000000..eddbf11e01 --- /dev/null +++ b/tests/unit/graph/test_start_anchor_weka_stamping.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weka ``_flatten_requests`` stamps each TrieNode's ``causal_parent_id``. + +The causal parent of a recorded leaf is the previous n/s leaf in its OWN list +(chain-prev), else the nearest preceding n/s leaf in an enclosing list (the +spawner), else ``None``. That stamping is what feeds +:func:`~aiperf.dataset.graph.segment_ir.interval_order.apply_start_anchors` +(already wired inside ``build_trie_ir``): when the causal parent is still IN +FLIGHT at a node's recorded start, the node's incoming edges collapse to one +start-anchored edge (``delay_after_predecessor_start_us``). The end-to-end case +also runs the graph validator to defeat the adapter-tests-skip-validator trap. +""" + +from __future__ import annotations + +from collections import defaultdict + +import msgspec +import pytest + +from aiperf.dataset.graph.adapters.weka.trace_models import WekaTrace +from aiperf.dataset.graph.adapters.weka.trie_build import ( + ReconCallbacks, + _flatten_requests, + build_trie_graph, +) +from aiperf.dataset.graph.models import TraceRecord +from aiperf.dataset.graph.validator import ValidationSeverity, validate + +_REQ = {"type": "n", "model": "M", "in": 128, "out": 16, "hash_ids": [1, 2]} + +_BLOCK_SIZE = 64 + + +def _stub_decode_block_tokens(hash_ids: list[int]) -> list[int]: + out: list[int] = [] + for h in hash_ids: + out.extend(range(h * 100, h * 100 + _BLOCK_SIZE)) + return out + + +def _stub_partial_tail_tokens(n_tokens: int, seed: str) -> list[int]: + base = sum(ord(c) for c in seed) * 1000 + return list(range(base, base + n_tokens)) + + +def _stub_decode_tokens_to_text(tokens: list[int]) -> str: + return "|".join(str(t) for t in tokens) + + +_STUB_CALLBACKS = ReconCallbacks( + decode_block_tokens=_stub_decode_block_tokens, + sample_partial_tail_tokens=_stub_partial_tail_tokens, + decode_tokens_to_text=_stub_decode_tokens_to_text, +) + + +def _trace(requests): + return WekaTrace.model_validate( + { + "id": "t", + "models": ["M"], + "block_size": 64, + "hash_id_scope": "local", + "requests": requests, + } + ) + + +def test_chain_prev_stamped(): + tr = _trace( + [ + {**_REQ, "t": 0.0, "api_time": 1.0}, + {**_REQ, "t": 2.0, "api_time": 1.0, "hash_ids": [1, 2, 3], "in": 192}, + ] + ) + nodes = _flatten_requests(tr.requests, root_scope="t") + assert nodes[0].causal_parent_id is None + assert nodes[1].causal_parent_id == "t:0" + + +def test_subagent_first_leaf_gets_spawner(): + tr = _trace( + [ + {**_REQ, "t": 0.0, "api_time": 8.0, "stop": "tool_use"}, + { + "t": 2.0, + "type": "subagent", + "agent_id": "a1", + "subagent_type": "Explore", + "status": "completed", + "models": ["M"], + "requests": [ + {**_REQ, "t": 2.5, "api_time": 1.0, "hash_ids": [50, 51]}, + { + **_REQ, + "t": 4.0, + "api_time": 1.0, + "hash_ids": [50, 51, 52], + "in": 192, + }, + ], + }, + ] + ) + nodes = _flatten_requests(tr.requests, root_scope="t") + by_id = {n.node_id: n for n in nodes} + assert by_id["a1:0"].causal_parent_id == "t:0" # spawner + assert by_id["a1:1"].causal_parent_id == "a1:0" # inner chain-prev + + +def test_marker_first_in_list_inherits_outer_spawner(): + inner = { + "t": 1.0, + "type": "subagent", + "agent_id": "a2", + "subagent_type": "Explore", + "status": "completed", + "models": ["M"], + "requests": [{**_REQ, "t": 1.2, "api_time": 0.5, "hash_ids": [60, 61]}], + } + tr = _trace( + [ + {**_REQ, "t": 0.0, "api_time": 4.0, "stop": "tool_use"}, + { + "t": 0.5, + "type": "subagent", + "agent_id": "a1", + "subagent_type": "Explore", + "status": "completed", + "models": ["M"], + "requests": [inner], # nested marker is FIRST entry of its list + }, + ] + ) + nodes = _flatten_requests(tr.requests, root_scope="t") + by_id = {n.node_id: n for n in nodes} + # a2:0's list has no preceding leaf; nor does its parent list; the + # spawner is the top-level t:0. + assert by_id["a2:0"].causal_parent_id == "t:0" + + +def test_end_to_end_graph_has_start_anchored_edges_and_validates(): + tr = _trace( + [ + {**_REQ, "t": 0.0, "api_time": 8.0, "stop": "tool_use"}, + { + "t": 2.0, + "type": "subagent", + "agent_id": "a1", + "subagent_type": "Explore", + "status": "completed", + "models": ["M"], + "requests": [{**_REQ, "t": 2.5, "api_time": 1.0, "hash_ids": [50, 51]}], + }, + {**_REQ, "t": 5.0, "api_time": 1.0, "hash_ids": [1, 2, 3], "in": 192}, + ] + ) + parsed, _pool = build_trie_graph(tr, callbacks=_STUB_CALLBACKS) + incoming = defaultdict(list) + for e in parsed.graph.edges: + incoming[e.target].append(e) + (spawn_edge,) = incoming["a1:0"] + assert spawn_edge.source == "t:0" + assert spawn_edge.delay_after_predecessor_start_us == pytest.approx(2.5e6) + (chain_edge,) = incoming["t:1"] + assert chain_edge.source == "t:0" + assert chain_edge.delay_after_predecessor_start_us == pytest.approx(5.0e6) + # "No completion wait" needs BOTH halves: an empty AND-fan-in alone is + # ambiguous (a node with no gating at all also has empty inputs). The sole + # incoming edge is the start-anchored one (destructured above) and it must + # carry NO completion delay, so the start anchor is provably the only gate. + assert parsed.graph.nodes["a1:0"].inputs == [] + assert spawn_edge.delay_after_predecessor_us is None + # validator must accept the graph (adapter-tests-skip-validator trap). Attach + # a TraceRecord so validation runs on the same trace wrap the production + # ingest path (``_parse_trace_dict``) applies. + validated = msgspec.structs.replace( + parsed, traces=[TraceRecord(id=tr.id)], segment_pool=_pool + ) + blocking = [ + i for i in validate(validated) if i.severity is ValidationSeverity.ERROR + ] + assert blocking == [], blocking diff --git a/tests/unit/graph/test_trie_prefix_reuse_oracle.py b/tests/unit/graph/test_trie_prefix_reuse_oracle.py new file mode 100644 index 0000000000..e6130c6130 --- /dev/null +++ b/tests/unit/graph/test_trie_prefix_reuse_oracle.py @@ -0,0 +1,669 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Differential oracle for ``build_trie_ir``'s prefix-path reuse. + +``build_trie_ir`` splices the +content-parent's already-emitted sid chain for whole messages strictly inside +the inherited block prefix, instead of re-decoding + re-hashing every covered +message of every node -- avoiding a quadratic prefix re-emission at corpus +scale. + +This file pins that optimization byte-for-byte against the straightforward +emission loop: +:func:`_oracle_build` is the naive ``build_trie_ir`` algorithm (the +quadratic per-node emission with fully inlined message assembly), kept here +as the reference ORACLE. Every generated and every ``@example``-pinned node set +must produce an IDENTICAL result from :func:`build_trie_ir` and the oracle: + +* pool key INSERTION ORDER, +* full :class:`Segment` tuples (role, content, parent_id, wire_json), +* every node's ``prompt_path`` / ``response_id`` / ``small_prompt`` flag. + +The oracle depends only on helpers whose behavior production shares +(``resolve_content_parents``, ``compute_asst_caps``, ``assign_block_tags``, +``add_message_chain``, ``assert_covered_isl``, ``emit_response_segment``), so a +reuse/splice bug shows up as a differential mismatch rather than being masked. + +One deliberate non-difference: production INTERLEAVES callback calls +differently (the oracle decodes ALL of a node's messages, then adds the chain; +``_assemble_messages_from`` decodes and ``pool.add``s per message). The +``pool.add`` call sequence -- the byte-ordering authority -- is identical, and +the callbacks are pure per-build functions of their arguments (the +``ReconCallbacks`` determinism contract), so reordering decode calls between +adds is byte-invisible; this file asserts outputs, never callback interleaving. +""" + +from __future__ import annotations + +import pytest +from hypothesis import example, given, settings +from hypothesis import strategies as st + +from aiperf.dataset.graph.adapters.dynamo.store_backed_pool import InterningSegmentPool +from aiperf.dataset.graph.segment_ir.interval_order import ( + apply_start_anchors, + build_interval_edges, + compute_ranks, +) +from aiperf.dataset.graph.segment_ir.pool import SegmentPool +from aiperf.dataset.graph.segment_ir.trie_content import ( + ReconCallbacks, + TrieBuild, + TrieISLMismatchError, + TrieNode, + TrieNodeBuild, + TrieRequest, + add_message_chain, + apply_idle_gap_warp, + assert_covered_isl, + assign_block_tags, + build_trie_ir, + compute_asst_caps, + emit_response_segment, + resolve_content_parents, +) + +BS = 4 + + +# --------------------------------------------------------------------------- # +# The verbatim PRE-change emission loop, kept as the reference oracle. +# --------------------------------------------------------------------------- # +def _oracle_build( + nodes: list[TrieNode], + *, + block_size: int, + callbacks: ReconCallbacks, + pool: SegmentPool, + small_prompt_fallback: bool = False, +) -> TrieBuild: + """Pre-rewrite ``build_trie_ir``: re-decode + re-emit EVERY covered message. + + Byte-for-byte the emission loop as it stood before the prefix-path-reuse + rewrite -- the quadratic path that re-grouped every node's full covered + prefix and re-``pool.add``ed each message (dedup made repeats no-ops). The + non-emission pipeline (parent resolution, warp, ranks, edges, anchors, caps, + tags) is the shared helper chain both paths run; only the per-node emission + below differs from the rewrite. + """ + resolve_content_parents(nodes) + apply_idle_gap_warp(nodes, None) + compute_ranks(nodes) + edges_by_node = build_interval_edges(nodes) + apply_start_anchors(nodes, edges_by_node) + caps = compute_asst_caps(nodes, block_size) + tags = assign_block_tags(nodes, block_size, caps) + + builds: dict[str, TrieNodeBuild] = {} + for node in nodes: + node_tags = tags[node.node_id] + covered = len(node_tags) + small = False + if covered == 0 and small_prompt_fallback and node.request.input_length > 0: + toks = callbacks.sample_partial_tail_tokens( + node.request.input_length, f"{node.node_id}:tiny" + ) + prompt_path = [ + pool.add( + role="user", + content=callbacks.decode_tokens_to_text(toks), + tokens=toks, + parent_id=None, + ) + ] + small = True + else: + # Inlined verbatim pre-change ``assemble_messages``: group the full + # covered tag list and emit every message fresh, root->tip. + groups: list[tuple[str, list[int]]] = [] + for j, (role, starts) in enumerate(node_tags): + if starts or not groups: + groups.append((role, [j])) + else: + groups[-1][1].append(j) + hash_ids = node.request.hash_ids + messages: list[tuple[str, str, list[int]]] = [] + assembled_tokens = 0 + for role, idxs in groups: + toks: list[int] = [] + for j in idxs: + toks.extend(callbacks.decode_block_tokens([hash_ids[j]])) + assembled_tokens += len(toks) + messages.append((role, callbacks.decode_tokens_to_text(toks), toks)) + prompt_path = add_message_chain(pool, messages) + if callbacks.block_exact: + assert_covered_isl(node, assembled_tokens, block_size) + response_id = emit_response_segment( + node, + pool=pool, + parent_id=prompt_path[-1] if prompt_path else None, + callbacks=callbacks, + ) + builds[node.node_id] = TrieNodeBuild( + prompt_path=prompt_path, response_id=response_id, small_prompt=small + ) + return TrieBuild(builds=builds, edges_by_node=edges_by_node) + + +# --------------------------------------------------------------------------- # +# Collision-free stub callbacks + node construction. +# --------------------------------------------------------------------------- # +def _stub_callbacks(block_size: int = BS, block_exact: bool = True) -> ReconCallbacks: + """Deterministic, injective stubs: block tokens = [hash_id] * block_size.""" + return ReconCallbacks( + decode_block_tokens=lambda hids: [t for h in hids for t in [h] * block_size], + # Keyed on the seed string so response/tiny tokens stay node-unique and + # never alias a prompt block's tokens. + sample_partial_tail_tokens=lambda n, seed: [abs(hash(seed)) % 100003] * n, + decode_tokens_to_text=lambda toks: " ".join(str(t) for t in toks), + block_exact=block_exact, + ) + + +def _node( + nid: str, + order: int, + hashes: list[int], + in_tok: int, + out_tok: int = 4, + t: float | None = None, +) -> TrieNode: + return TrieNode( + node_id=nid, + request=TrieRequest( + hash_ids=list(hashes), + input_length=in_tok, + output_length=out_tok, + t=float(order) if t is None else t, + api_time=0.5, + ), + order=order, + ) + + +def _fresh(spec: list[tuple[str, int, list[int], int, int]]) -> list[TrieNode]: + """Rebuild nodes from a (nid, order, hashes, in_tok, out_tok) spec. + + A fresh node list per leg is mandatory: the trie pipeline mutates nodes in + place (content_parent, warped_start, rank), so oracle and rewrite must each + get their own copies. + """ + return [ + _node(nid, order, hashes, in_tok, out_tok) + for nid, order, hashes, in_tok, out_tok in spec + ] + + +# --------------------------------------------------------------------------- # +# Differential harness. +# --------------------------------------------------------------------------- # +def _assert_identical( + spec: list[tuple[str, int, list[int], int, int]], + *, + block_size: int = BS, + small_prompt_fallback: bool = False, + block_exact: bool = True, + interning: bool = False, +) -> tuple[SegmentPool, TrieBuild]: + """Build with the oracle and the rewrite; assert byte-identical results.""" + + def _pool() -> SegmentPool: + return InterningSegmentPool() if interning else SegmentPool() + + pool_o = _pool() + res_o = _oracle_build( + _fresh(spec), + block_size=block_size, + callbacks=_stub_callbacks(block_size, block_exact), + pool=pool_o, + small_prompt_fallback=small_prompt_fallback, + ) + pool_n = _pool() + res_n = build_trie_ir( + _fresh(spec), + block_size=block_size, + callbacks=_stub_callbacks(block_size, block_exact), + pool=pool_n, + idle_gap_cap_seconds=None, + small_prompt_fallback=small_prompt_fallback, + ) + + # (1) pool insertion ORDER is byte-identical. + assert list(pool_o.by_id.keys()) == list(pool_n.by_id.keys()), ( + "pool insertion order diverged" + ) + # (2) full Segment tuples are byte-identical. + for sid in pool_o.by_id: + so = pool_o.by_id[sid] + sn = pool_n.by_id[sid] + assert (so.role, so.content, so.parent_id, so.wire_json) == ( + sn.role, + sn.content, + sn.parent_id, + sn.wire_json, + ), f"segment {sid} tuple diverged" + # (3) per-node prompt_path / response_id / small_prompt. + assert res_o.builds.keys() == res_n.builds.keys() + for nid, bo in res_o.builds.items(): + bn = res_n.builds[nid] + assert bo.prompt_path == bn.prompt_path, f"node {nid} prompt_path diverged" + assert bo.response_id == bn.response_id, f"node {nid} response_id diverged" + assert bo.small_prompt == bn.small_prompt, f"node {nid} small_prompt diverged" + return pool_n, res_n + + +# --------------------------------------------------------------------------- # +# Named adversarial legs (the geometries from /tmp/adv_reuse_diff.py, adopted +# as the mandated @example pins per both adversarial reviews). +# --------------------------------------------------------------------------- # +# depth >= 8 extending chain -- the deep-chain gate-outcome pin (Review 0 rev 2): +# a cumulative-count bookkeeping slip on a long reuse chain trips the ISL gate. +_EXTEND_CHAIN_DEEP = [ + (f"x{i}", i, list(range(50, 50 + i + 1)), BS * (i + 1), 4) for i in range(9) +] +# child == parent: whole prefix reused, empty fresh region. +_DUPLICATE = [ + ("a", 0, [1, 2, 3], 12, 4), + ("b", 1, [1, 2, 3, 4, 5], 20, 4), + ("c", 2, [1, 2, 3, 4, 5], 20, 4), + ("c2", 3, [1, 2, 3, 4, 5], 20, 4), +] +# branch point: BOTH grandchild record-composition cases live here. 'q' +# truncates p's message into a fragment ([0,2) of p's [0,4)); 's' then splices +# straight THROUGH q's fragment-bearing record (fragment-then-reuse), while +# 'r' branches off q past the fragment (reuse-after-fragment). +_BRANCH_POINT_FRAGMENT = [ + ("p", 0, [60, 61, 62, 63], 16, 4), + ("q", 1, [60, 61, 70, 71], 16, 4), + ("r", 2, [60, 61, 70, 99], 16, 4), + ("s", 3, [60, 61, 70, 71, 80, 81], 24, 4), +] +# over-share: 'e' declares in//bs = 2 < lcp = 6 -> covered clamp bounds reuse. +_OVER_SHARE = [ + ("a", 0, [1, 2, 3], 12, 4), + ("b", 1, [1, 2, 3, 4, 5, 6], 24, 4), + ("e", 2, [1, 2, 3, 4, 5, 6], 8, 4), + ("e2", 3, [1, 2, 3, 4, 5, 6, 7], 28, 4), +] +# under-covering parent: 'p' covers 2 blocks < child lcp 4 -> len(parent_tags) clamp. +_UNDER_COVER_PARENT = [ + ("p", 0, [20, 21, 22, 23], 8, 4), + ("q", 1, [20, 21, 22, 23, 24], 20, 4), + ("r", 2, [20, 21, 22, 23, 24, 25], 24, 4), +] +# degenerate pull-back caps: r4 has new_n 0 (empty fresh region, whole-splice +# prompt_path) and caps its owner's frozen tags via compute_asst_caps. +_DEGENERATE_PULLBACK = [ + ("r1", 0, [30], 4, 4), + ("r2", 1, [30, 31], 8, 4), + ("r3", 2, [30, 31, 32], 12, 4), + ("r4", 3, [30, 31, 32], 8, 4), + ("r5", 4, [30, 31, 32, 33], 16, 4), +] +# covered==0 parent (no record published) with covered>0 children reusing NOTHING. +_ZERO_COVERED_PARENT = [ + ("h", 0, [1, 2, 3], 2, 4), + ("i", 1, [1, 2, 3, 4], 16, 4), + ("j", 2, [1, 2, 3, 4], 16, 4), +] +# boundary landing exactly on a message end. +_BOUNDARY_EXACT_MESSAGE_END = [ + ("a", 0, [1], 4, 8), + ("b", 1, [1, 2, 3], 12, 4), + ("c", 2, [1, 2, 3, 9], 16, 4), + ("d", 3, [1, 2, 9, 9], 12, 4), +] +# tags coincide beyond the lcp but hash ids differ -- the geometry-not-tags trap. +_TAGS_COINCIDE_BEYOND_LCP = [ + ("a", 0, [1, 2, 3, 4], 16, 4), + ("b", 1, [1, 2, 8, 9], 16, 4), + ("c", 2, [1, 2, 8, 9, 10], 20, 4), +] +# j>=1 splice PLUS a straddling fragment in ONE node: 'c' over-shares b +# (covered 4 < lcp 6) while c's degenerate pull-back caps b's assistant run to +# 0, merging b's new region into one message [3,6) -- so c splices j=1 whole +# message [0,3), resumes at block 3 (a PARENT-COPIED tag, exercising that branch +# of the production resume-point assert), and re-emits the fragment [3,4) fresh. +# 'd' then splices straight through c's fragment-bearing record (j=2). +_SPLICE_PLUS_FRAGMENT = [ + ("a", 0, [1, 2, 3], 12, 4), + ("b", 1, [1, 2, 3, 4, 5, 6], 24, 4), + ("c", 2, [1, 2, 3, 4, 5, 6], 16, 4), + ("d", 3, [1, 2, 3, 4, 5, 6, 7], 28, 4), +] +# dynamo-shaped: negative virtual hash ids, extending chain (block_size=16 leg +# supplied via the parametrize below). +_DYNAMO_NEGATIVE_IDS = [ + ("d0", 0, [-1, -2], 32, 8), + ("d1", 1, [-1, -2, -3], 48, 8), + ("d2", 2, [-1, -2, -3, -4], 64, 8), + ("d3", 3, [-1, -2, -3, -4, -5], 80, 8), +] + +_NAMED_CASES = { + "extend_chain_deep": _EXTEND_CHAIN_DEEP, + "duplicate": _DUPLICATE, + "branch_point_fragment": _BRANCH_POINT_FRAGMENT, + "over_share": _OVER_SHARE, + "under_cover_parent": _UNDER_COVER_PARENT, + "degenerate_pullback_caps": _DEGENERATE_PULLBACK, + "zero_covered_parent": _ZERO_COVERED_PARENT, + "boundary_exact_message_end": _BOUNDARY_EXACT_MESSAGE_END, + "tags_coincide_beyond_lcp": _TAGS_COINCIDE_BEYOND_LCP, + "splice_plus_fragment": _SPLICE_PLUS_FRAGMENT, +} + + +@pytest.mark.parametrize("name", list(_NAMED_CASES)) +@pytest.mark.parametrize("fallback", [False, True], ids=["no_fallback", "fallback"]) +def test_named_cases_byte_identical(name: str, fallback: bool) -> None: + """Every adversarial geometry emits byte-identically under both fallbacks.""" + _assert_identical(_NAMED_CASES[name], small_prompt_fallback=fallback) + + +@pytest.mark.parametrize("name", list(_NAMED_CASES)) +def test_named_cases_block_exact_false(name: str) -> None: + """block_exact=False (scheduling-only placeholder parse) stays byte-identical.""" + _assert_identical(_NAMED_CASES[name], block_exact=False, small_prompt_fallback=True) + + +@pytest.mark.parametrize("name", list(_NAMED_CASES)) +def test_named_cases_interning_pool(name: str) -> None: + """Every geometry stays byte-identical on the eager InterningSegmentPool route.""" + _assert_identical(_NAMED_CASES[name], small_prompt_fallback=True, interning=True) + + +def test_dynamo_shaped_negative_ids_block16() -> None: + """Dynamo-shaped leg: negative virtual hash ids, block_size=16, fallback on.""" + spec = [ + ("d0", 0, [-1, -2], 32, 8), + ("d1", 1, [-1, -2, -3], 48, 8), + ("d2", 2, [-1, -2, -3, -4], 64, 8), + ("d3", 3, [-1, -2, -3, -4, -5], 80, 8), + ] + _assert_identical(spec, block_size=16, small_prompt_fallback=True) + + +def test_small_prompt_parent_with_covered_child() -> None: + """A covered>0 CHILD of a small-prompt (covered==0) parent must reuse NOTHING. + + The small-prompt parent publishes no emission record; the child's inherited + count is structurally 0 (its parent tags are empty), so the missing-parent- + record guard (records.get() AND inherited > 0) must hold. Both fallback and + non-fallback are exercised: with fallback the parent emits a single tiny user + message (still no record), without it the parent emits an empty prompt. + """ + # 'tiny' covers 0 blocks (in=2 < bs=4); 'child' shares its hash prefix but + # declares enough input to cover blocks. + spec = [ + ("tiny", 0, [1, 2, 3], 2, 4), + ("child", 1, [1, 2, 3, 4], 16, 4), + ("grandchild", 2, [1, 2, 3, 4, 5], 20, 4), + ] + _assert_identical(spec, small_prompt_fallback=True) + _assert_identical(spec, small_prompt_fallback=False) + + +def test_interning_pool_splice_is_identical_to_canonical_first_born() -> None: + """On InterningSegmentPool spliced prompt_path entries ARE the parent's + canonical first-born str objects (claim vi -- identity, not just equality). + + Every value re-listed across nodes must resolve to ONE canonical object, so + the reuse splice must hand back the exact object the first emitter interned, + never a fresh-but-equal str. + """ + spec = _EXTEND_CHAIN_DEEP + pool = InterningSegmentPool() + result = build_trie_ir( + _fresh(spec), + block_size=BS, + callbacks=_stub_callbacks(), + pool=pool, + idle_gap_cap_seconds=None, + ) + canonical: dict[str, str] = {} + for build in result.builds.values(): + for sid in build.prompt_path: + if sid in canonical: + # Same value re-listed downstream must be the SAME object. + assert sid is canonical[sid], "interning identity broken by splice" + else: + canonical[sid] = sid + # The deep chain must actually re-list shared prefixes (else the assert is vacuous). + total = sum(len(b.prompt_path) for b in result.builds.values()) + assert total > len(canonical), "expected shared prefixes to be re-listed" + + +def test_resume_point_lands_on_message_start() -> None: + """The reuse resume block must open a message (the production-assert invariant). + + Re-derives, per node, the geometric ``inherited`` and the parent's message-end + boundaries the way the rewrite does, and checks the resume tag is a message + start -- the cheapest tripwire against grouping drift in the frozen file. + """ + import bisect + + from aiperf.dataset.graph.segment_ir.trie_content import ( + compute_turn_block_geometry, + ) + + spec = _EXTEND_CHAIN_DEEP + [("z", 9, list(range(50, 60)), 40, 4)] + nodes = _fresh(spec) + resolve_content_parents(nodes) + caps = compute_asst_caps(nodes, BS) + tags = assign_block_tags(nodes, BS, caps) + # message end blocks per node (exclusive), from the frozen tags. + ends: dict[str, list[int]] = {} + for node in nodes: + node_tags = tags[node.node_id] + me: list[int] = [] + for k, (_role, starts) in enumerate(node_tags): + if starts and k > 0: + me.append(k) + me.append(len(node_tags)) + ends[node.node_id] = me + parent = node.content_parent + if parent is None: + continue + geo = compute_turn_block_geometry( + parent.request.hash_ids, + node.request.hash_ids, + node.request.input_length, + BS, + ) + inherited = min(geo.lcp, len(tags[parent.node_id]), geo.m_curr_covered) + if inherited <= 0: + continue + j = bisect.bisect_right(ends[parent.node_id], inherited) + if j <= 0: + continue + start_block = ends[parent.node_id][j - 1] + if start_block < len(node_tags): + assert node_tags[start_block][1] is True, ( + f"resume block {start_block} of {node.node_id} is not a message start" + ) + + +def test_reused_prefix_decoded_exactly_once_kills_quadratic() -> None: + """A deep chain decodes each shared block ONCE -- the quadratic is dead. + + Counting stub: every ``decode_block_tokens`` call is recorded. Under the + rewrite each block is materialized exactly once (at first occurrence) and + reused by splice thereafter, so the total decode count equals the number of + DISTINCT covered blocks across the chain (linear), not the per-node + re-decode sum (quadratic). This simultaneously pins that a reused prefix is + never re-decoded -- so a callback that would drift on a repeat emission + cannot even be called there (repeat-emission drift is excluded by the + ReconCallbacks determinism contract, not re-checked at the child). + """ + n = 12 + spec = [ + (f"x{i}", i, list(range(100, 100 + i + 1)), BS * (i + 1), 4) for i in range(n) + ] + calls: list[int] = [] + + def _counting_decode(hids: list[int]) -> list[int]: + calls.append(len(hids)) + return [t for h in hids for t in [h] * BS] + + cb = ReconCallbacks( + decode_block_tokens=_counting_decode, + sample_partial_tail_tokens=lambda n, seed: [abs(hash(seed)) % 100003] * n, + decode_tokens_to_text=lambda toks: " ".join(str(t) for t in toks), + ) + build_trie_ir( + _fresh(spec), + block_size=BS, + callbacks=cb, + pool=SegmentPool(), + idle_gap_cap_seconds=None, + ) + # Node i covers i+1 blocks; distinct covered blocks across the chain = the + # LAST node's covered count = n (blocks 100..100+n-1). The quadratic path + # decoded sum(i+1 for i) = n(n+1)/2 blocks instead. + distinct_blocks = n + total_decoded = sum(calls) + assert total_decoded == distinct_blocks, ( + f"decoded {total_decoded} blocks; linear reuse must decode {distinct_blocks} " + f"(quadratic would decode {n * (n + 1) // 2})" + ) + + +def test_first_occurrence_decode_drift_still_aborts() -> None: + """First-occurrence (fresh materialization) decode drift still hard-aborts. + + The reuse narrowing does NOT weaken the gate on freshly materialized blocks: + a decode that emits 2x block_size tokens on the FIRST occurrence of a block + must still trip the ISL gate (at the parent, before any child could reuse). + """ + drifted = ReconCallbacks( + decode_block_tokens=lambda hids: [t for h in hids for t in [h] * (BS * 2)], + sample_partial_tail_tokens=lambda n, seed: [7] * n, + decode_tokens_to_text=lambda toks: " ".join(str(t) for t in toks), + ) + spec = [ + ("a", 0, [1, 2], 8, 4), + ("b", 1, [1, 2, 3], 12, 4), + ] + with pytest.raises(TrieISLMismatchError): + build_trie_ir( + _fresh(spec), + block_size=BS, + callbacks=drifted, + pool=SegmentPool(), + idle_gap_cap_seconds=None, + ) + + +# --------------------------------------------------------------------------- # +# Hypothesis differential over arbitrary node sets. +# --------------------------------------------------------------------------- # +# Small alphabet -> dense shared prefixes, branch points, and exact duplicates. +_HASH = st.integers(min_value=-3, max_value=6) +_ROW = st.tuples(st.lists(_HASH, max_size=8), st.integers(min_value=0, max_value=40)) +_NODE_SETS = st.lists(_ROW, min_size=1, max_size=12) + + +def _spec_from_rows( + rows: list[tuple[list[int], int]], +) -> list[tuple[str, int, list[int], int, int]]: + return [(f"r{i}", i, hashes, in_tok, 4) for i, (hashes, in_tok) in enumerate(rows)] + + +@settings(deadline=None, max_examples=300) +@given(rows=_NODE_SETS) +@example(rows=[(h, in_tok) for _, _, h, in_tok, _ in _EXTEND_CHAIN_DEEP]) +@example(rows=[(h, in_tok) for _, _, h, in_tok, _ in _BRANCH_POINT_FRAGMENT]) +@example(rows=[(h, in_tok) for _, _, h, in_tok, _ in _OVER_SHARE]) +@example(rows=[(h, in_tok) for _, _, h, in_tok, _ in _UNDER_COVER_PARENT]) +@example(rows=[(h, in_tok) for _, _, h, in_tok, _ in _TAGS_COINCIDE_BEYOND_LCP]) +@example(rows=[(h, in_tok) for _, _, h, in_tok, _ in _SPLICE_PLUS_FRAGMENT]) +def test_prefix_reuse_matches_oracle_property( + rows: list[tuple[list[int], int]], +) -> None: + """The rewrite == the verbatim pre-change oracle on arbitrary node sets. + + Both fallback modes are exercised so the small-prompt / record-publication + branches are differenced too (record composition is the only inductive state + the rewrite introduces). + """ + spec = _spec_from_rows(rows) + _assert_identical(spec, small_prompt_fallback=False) + _assert_identical(spec, small_prompt_fallback=True) + + +# --------------------------------------------------------------------------- # +# Real-fixture leg: a committed weka trace through the REAL weka lowering. +# --------------------------------------------------------------------------- # +def test_real_weka_fixture_byte_matches_oracle() -> None: + """A real weka fixture trace, real production callbacks: oracle == rewrite. + + Stub-free: the trace is parsed through the real weka models + flattening and + both legs run the real builtin-tokenizer/coding-corpus callbacks + (fresh per leg, so per-trace decode caches are cold), byte-comparing pool + key INSERTION ORDER, full Segment tuples, and every node's build artifacts. + The fixture is an extending chain, so the reuse path genuinely fires. + """ + from pathlib import Path + + import orjson + + from aiperf.dataset.graph.adapters.weka.trace_models import WekaTrace + from aiperf.dataset.graph.adapters.weka.trie_build import ( + _default_callbacks, + _flatten_requests, + ) + + fixture = Path(__file__).parent / "fixtures" / "weka_min.json" + trace = WekaTrace.model_validate(orjson.loads(fixture.read_bytes())) + bs = trace.block_size + + def _real_callbacks() -> ReconCallbacks: + return _default_callbacks( + "builtin", + "coding", + 1234, + trace_id=trace.id, + block_size=bs, + hash_scope=trace.hash_id_scope, + ) + + pool_o = SegmentPool() + res_o = _oracle_build( + _flatten_requests(trace.requests, root_scope=trace.id), + block_size=bs, + callbacks=_real_callbacks(), + pool=pool_o, + ) + pool_n = SegmentPool() + res_n = build_trie_ir( + _flatten_requests(trace.requests, root_scope=trace.id), + block_size=bs, + callbacks=_real_callbacks(), + pool=pool_n, + idle_gap_cap_seconds=None, + ) + + assert list(pool_o.by_id.keys()) == list(pool_n.by_id.keys()), ( + "pool insertion order diverged on the real fixture" + ) + for sid, so in pool_o.by_id.items(): + sn = pool_n.by_id[sid] + assert (so.role, so.content, so.parent_id, so.wire_json) == ( + sn.role, + sn.content, + sn.parent_id, + sn.wire_json, + ), f"segment {sid} tuple diverged on the real fixture" + assert res_o.builds.keys() == res_n.builds.keys() + for nid, bo in res_o.builds.items(): + bn = res_n.builds[nid] + assert bo.prompt_path == bn.prompt_path + assert bo.response_id == bn.response_id + assert bo.small_prompt == bn.small_prompt + # The fixture must actually exercise reuse: its requests extend one chain + # ([1,2] -> [1,2,3] -> [1,2,3,4]), so later prompt paths share a prefix. + paths = [res_n.builds[nid].prompt_path for nid in sorted(res_n.builds)] + assert len(paths) >= 3 + assert paths[1][: len(paths[0])] == paths[0] + assert paths[2][: len(paths[1])] == paths[1] diff --git a/tests/unit/graph/test_tstar_activation.py b/tests/unit/graph/test_tstar_activation.py new file mode 100644 index 0000000000..cffa1adeb9 --- /dev/null +++ b/tests/unit/graph/test_tstar_activation.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""t* activation: the config window + env seed thread through to GraphIRReplayStrategy. + +The prior unit built the t* machinery (``GraphIRConversationSource`` sampling, +warmup variants, snapshot rewrite) but left it INERT: nothing passed real ratios +into the strategy, so every run sampled t*=0 (identity replay). These tests pin +the activation seam: + +* ``BenchmarkConfig`` exposes ``trajectory_start_min/max_ratio`` + (``--trajectory-start-min/max-ratio``); the sampling seed is the run's + resolved ``--random-seed`` carried on the phase config, +* ``PhaseRunner._build_graph_ir_strategy`` threads the phase config's window + + the env seed into the strategy, so a graph run with a positive window + samples a non-trivial t* (NOT t*=0 identity). +""" + +from __future__ import annotations + +from pathlib import Path + +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.plugin.enums import TimingMode +from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy + +_FIX = Path(__file__).parent / "fixtures" / "weka_min.json" + + +def test_config_tstar_window_defaults_off() -> None: + # Bare default is full replay (t* window OFF); the AgentX 0.0..1.0 + # window is scenario-applied (--scenario inferencex-agentx-mvp). + from aiperf.config import BenchmarkConfig + + assert BenchmarkConfig.model_fields["trajectory_start_min_ratio"].default is None + assert BenchmarkConfig.model_fields["trajectory_start_max_ratio"].default is None + + +def test_strategy_with_positive_window_samples_nonzero_tstar() -> None: + """A deterministic 0.5..0.5 window engages t* (non-identity partition).""" + parsed = from_weka_trace(str(_FIX)) + strategy = GraphIRReplayStrategy( + credit_issuer=object(), + parsed_graph=parsed, + register_observer=lambda _obs: None, + start_min_ratio=0.5, + start_max_ratio=0.5, + t_star_random_seed=42, + ) + plans = strategy._plans + assert plans, "expected a per-trace t* plan" + gt = next(iter(plans.values())) + assert gt.t_star_us > 0, "positive window must engage t* (not the inert t*=0)" + + +def test_strategy_zero_window_is_inert_identity() -> None: + """Explicit [0, 0] strategy kwargs keep the byte-identical t*=0 path.""" + parsed = from_weka_trace(str(_FIX)) + strategy = GraphIRReplayStrategy( + credit_issuer=object(), + parsed_graph=parsed, + register_observer=lambda _obs: None, + start_min_ratio=0.0, + start_max_ratio=0.0, + ) + gt = next(iter(strategy._plans.values())) + assert gt.t_star_us == 0 + + +def test_strategy_default_window_is_inert() -> None: + """No ratio kwargs => the 0.0..0.0 constructor default keeps t*=0 (full replay).""" + parsed = from_weka_trace(str(_FIX)) + strategy = GraphIRReplayStrategy( + credit_issuer=object(), + parsed_graph=parsed, + register_observer=lambda _obs: None, + t_star_random_seed=42, + ) + gt = next(iter(strategy._plans.values())) + assert gt.t_star_us == 0 + + +def test_build_graph_ir_strategy_threads_config_ratios() -> None: + """``PhaseRunner._build_graph_ir_strategy`` sources the phase-config window + plus the env seed. + + Build the strategy via the same code path the runner uses (without standing + up a full PhaseRunner) and confirm the config-driven positive window engages t*. + """ + parsed = from_weka_trace(str(_FIX)) + + captured: dict = {} + + class _CapturingStrategy(GraphIRReplayStrategy): + def __init__(self, **kw): + captured.update(kw) + super().__init__(**kw) + + class _Handler: + def set_graph_return_observer(self, obs) -> None: + self._obs = obs + + def set_graph_first_token_observer(self, obs) -> None: + self._ft_obs = obs + + class _Config: + timing_mode = TimingMode.GRAPH_IR + phase = None + concurrency = None + expected_num_sessions = None + trajectory_start_min_ratio = 0.5 + trajectory_start_max_ratio = 0.5 + burst_phase_starts = None + random_seed = 7 + + from aiperf.timing.graph_channel import GraphPhaseChannel + from aiperf.timing.phase.runner import PhaseRunner + + runner = PhaseRunner.__new__(PhaseRunner) + runner._config = _Config() + runner._conversation_source = None + runner._graph_channel = GraphPhaseChannel(parsed_graph=parsed) + runner._scheduler = None + runner._stop_checker = None + runner._credit_issuer = object() + runner._lifecycle = None + runner._callback_handler = _Handler() + + strategy = runner._build_graph_ir_strategy(_CapturingStrategy) + + assert captured["start_min_ratio"] == 0.5 + assert captured["start_max_ratio"] == 0.5 + assert captured["t_star_random_seed"] == 7 + gt = next(iter(strategy._plans.values())) + assert gt.t_star_us > 0, "config window must engage t* through the runner seam" diff --git a/tests/unit/graph/test_tstar_selection.py b/tests/unit/graph/test_tstar_selection.py new file mode 100644 index 0000000000..e0919d0872 --- /dev/null +++ b/tests/unit/graph/test_tstar_selection.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""t* snapshot selection in the graph IR source (Task 12). + +Per trace the source samples a wall-clock instant ``t*`` in +``[start_min_ratio, start_max_ratio] * trace_duration`` (deterministic under a +seed, mirroring agentx ``TrajectorySource``). ``t*==0`` (the default +full-replay ratio) means full native replay with no warmup history. +""" + +from __future__ import annotations + +from pathlib import Path + +import orjson + +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.timing.graph_ir_source import GraphIRConversationSource + +FIX = Path(__file__).parent / "fixtures" / "weka_min.json" + + +def _multi_trace_corpus(tmp_path: Path, n_traces: int = 3) -> Path: + """A directory of ``n_traces`` weka traces with DISTINCT durations. + + Derived from the committed single-trace fixture by re-id'ing and scaling + the recorded timestamps, so per-trace t* windows differ and cross-trace / + cross-seed variation is actually observable. + """ + base = orjson.loads(FIX.read_bytes()) + corpus = tmp_path / "corpus" + corpus.mkdir() + for i in range(n_traces): + doc = dict(base) + doc["id"] = f"trace_{i}" + scale = float(i + 1) + doc["requests"] = [{**req, "t": req["t"] * scale} for req in base["requests"]] + (corpus / f"trace_{i}.json").write_bytes(orjson.dumps(doc)) + return corpus + + +def test_full_replay_default_yields_tstar_zero() -> None: + parsed = from_weka_trace(str(FIX)) + src = GraphIRConversationSource( + parsed=parsed, start_min_ratio=0.0, start_max_ratio=0.0 + ) + gt = next(iter(src.iter_traces())) + # Ratios [0, 0] => t*=0 => full native replay, no warmup history. + assert gt.t_star_us == 0 + + +def test_positive_window_samples_nonzero_tstar() -> None: + parsed = from_weka_trace(str(FIX)) + src = GraphIRConversationSource( + parsed=parsed, + start_min_ratio=0.5, + start_max_ratio=0.5, + random_seed=42, + ) + gt = next(iter(src.iter_traces())) + # weka_min: 3 root turns at offsets 0, 1.5s, 3.0s. A 50% window over the + # ~3s duration engages a positive t* (not the inert t*=0 identity replay). + assert gt.t_star_us > 0 + + +def test_tstar_is_deterministic_under_seed() -> None: + parsed = from_weka_trace(str(FIX)) + kw = dict(start_min_ratio=0.0, start_max_ratio=0.7, random_seed=123) + a = next(iter(GraphIRConversationSource(parsed=parsed, **kw).iter_traces())) + b = next(iter(GraphIRConversationSource(parsed=parsed, **kw).iter_traces())) + assert a.t_star_us == b.t_star_us + + +def test_tstar_varies_across_traces_and_seeds(tmp_path: Path) -> None: + """Distinct seeds pick distinct t* profiles over a REAL multi-trace corpus. + + A single-trace fixture would make any per-trace-id disjunction vacuous; + three traces with distinct durations force both axes to be observable. + """ + parsed = from_weka_trace(str(_multi_trace_corpus(tmp_path))) + assert len(parsed.traces) == 3 + kw = dict(start_min_ratio=0.0, start_max_ratio=0.9) + + seed_1 = { + gt.trace_id: gt.t_star_us + for gt in GraphIRConversationSource( + parsed=parsed, random_seed=1, **kw + ).iter_traces() + } + seed_2 = { + gt.trace_id: gt.t_star_us + for gt in GraphIRConversationSource( + parsed=parsed, random_seed=2, **kw + ).iter_traces() + } + + assert set(seed_1) == set(seed_2) and len(seed_1) == 3 + # Across seeds: the sampled t* profile must differ. + assert seed_1 != seed_2 + # Across traces (within one seed): distinct-duration traces must not all + # collapse to one t*. + assert len(set(seed_1.values())) > 1 + + +def test_source_yields_one_graphtrace_per_trace() -> None: + parsed = from_weka_trace(str(FIX)) + src = GraphIRConversationSource(parsed=parsed) + traces = list(src.iter_traces()) + assert len(traces) == len(parsed.traces) + assert traces[0].parsed_graph is parsed diff --git a/tests/unit/graph/test_typed_channel_placeholders.py b/tests/unit/graph/test_typed_channel_placeholders.py new file mode 100644 index 0000000000..5cce1e70d4 --- /dev/null +++ b/tests/unit/graph/test_typed_channel_placeholders.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Type-correct channel placeholders on the LlmNode credit path. + +The credit path resolves a content-free placeholder string (`""`) and the +failure path writes a gate sentinel; both flow into the channel store, whose +global `snapshot_at_seq` reduces EVERY channel. A messages-typed output channel +therefore needs a type-correct empty (`[]`) — `add_messages` rejects non-list +values — on both the success path (`dispatch/llm.py`) and the dispatch-failure +sentinel branch (`executor._handle_node_exception`). These tests run a two-node +chain where node `a` writes a messages-typed channel and node `b`'s snapshot +reduces it: with untyped placeholders either path raises `TypeError` and the +trace errors. +""" + +import asyncio +from typing import Any + +import pytest + +from aiperf.dataset.graph.models import ( + ChannelRequirement, + ChannelSpec, + ChannelType, + GraphRecord, + LlmNode, + ParsedGraph, + ReducerName, + StaticEdge, + TraceRecord, +) +from aiperf.dataset.graph.segment_ir.pool import SegmentPool +from aiperf.graph.credit_dispatch_adapter import GraphDispatchError +from aiperf.graph.executor import TraceExecutor + + +def _messages_chain_parsed(*, with_pool: bool) -> ParsedGraph: + graph = GraphRecord( + state={ + "a_out": ChannelSpec( + type=ChannelType.MESSAGES, reducer=ReducerName.ADD_MESSAGES + ), + "b_out": ChannelSpec(), + }, + nodes={ + "a": LlmNode( + prompt=[{"role": "user", "content": "q1"}], + output="a_out", + ), + "b": LlmNode( + prompt=[{"role": "user", "content": "q2"}], + output="b_out", + inputs=[ChannelRequirement(channel="a_out", count=1)], + ), + }, + edges=[ + StaticEdge(source="START", target="a"), + StaticEdge(source="a", target="b"), + StaticEdge(source="b", target="END"), + ], + ) + return ParsedGraph( + graph=graph, + traces=[TraceRecord(id="t1")], + segment_pool=SegmentPool() if with_pool else None, + ) + + +class _EchoIssuer: + """Resolves every dispatch with the credit path's placeholder string.""" + + async def dispatch( + self, node: Any, request: Any, ctx: Any, first_token_cb: Any = None + ) -> str: + return "" + + +class _FailFirstIssuer(_EchoIssuer): + """Raises GraphDispatchError for node 'a', succeeds for everything else.""" + + async def dispatch( + self, node: Any, request: Any, ctx: Any, first_token_cb: Any = None + ) -> str: + if request.node_id == "a": + raise GraphDispatchError("simulated dispatch failure") + return "" + + +@pytest.mark.asyncio +async def test_success_placeholder_is_list_for_messages_channel() -> None: + parsed = _messages_chain_parsed(with_pool=False) + executor = TraceExecutor(parsed, credit_issuer=_EchoIssuer()) + async with asyncio.TaskGroup(): + result = await executor.run(parsed.traces[0]) + assert result.trace_id == "t1" + + +@pytest.mark.asyncio +async def test_failure_sentinel_is_list_for_messages_channel() -> None: + parsed = _messages_chain_parsed(with_pool=True) + executor = TraceExecutor(parsed, credit_issuer=_FailFirstIssuer()) + async with asyncio.TaskGroup(): + result = await executor.run(parsed.traces[0]) + assert result.trace_id == "t1" + + +@pytest.mark.asyncio +async def test_failure_without_pool_keeps_orphan_semantics() -> None: + """Non-segment-store parses keep the historical orphan-then-error behavior.""" + parsed = _messages_chain_parsed(with_pool=False) + executor = TraceExecutor(parsed, credit_issuer=_FailFirstIssuer()) + with pytest.raises(ExceptionGroup) as excinfo: + async with asyncio.TaskGroup(): + await executor.run(parsed.traces[0]) + from aiperf.graph.channel_store import ChannelOrphanedError + + # The dispatch failure is contained ("continuing past the failed turn"); + # what surfaces is the ORPHANED downstream channel -- the exact historical + # orphan-then-error semantics this test locks. + assert excinfo.group_contains(ChannelOrphanedError) + + +class _RefuseFirstIssuer(_EchoIssuer): + """Refuses node 'a' the way the issuer stop gate does at duration end.""" + + async def dispatch( + self, node: Any, request: Any, ctx: Any, first_token_cb: Any = None + ) -> str: + if request.node_id == "a": + from aiperf.graph.credit_dispatch_adapter import CreditIssueRefusedError + + raise CreditIssueRefusedError("refused by issuer stop gate") + return "" + + +@pytest.mark.asyncio +async def test_issuer_refusal_stops_trace_even_with_pool() -> None: + """Refusal is a trace-stop, never sentinel-continued. + + Contrast with test_failure_sentinel_is_list_for_messages_channel: a plain + dispatch failure on the same parse continues with omission; a refusal must + unwind so duration-end traces stop instead of sentinel-spinning through + their remaining nodes. + """ + parsed = _messages_chain_parsed(with_pool=True) + executor = TraceExecutor(parsed, credit_issuer=_RefuseFirstIssuer()) + with pytest.raises(ExceptionGroup) as excinfo: + async with asyncio.TaskGroup(): + await executor.run(parsed.traces[0]) + from aiperf.graph.credit_dispatch_adapter import CreditIssueRefusedError + + # The unwind must carry the REFUSAL, proving the stop was the issuer gate + # (not a sentinel-continued dispatch failure or unrelated error). + assert excinfo.group_contains(CreditIssueRefusedError) diff --git a/tests/unit/graph/test_unified_interned_materialize.py b/tests/unit/graph/test_unified_interned_materialize.py new file mode 100644 index 0000000000..8c40e1cb98 --- /dev/null +++ b/tests/unit/graph/test_unified_interned_materialize.py @@ -0,0 +1,139 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from pathlib import Path + +import orjson +import pytest + +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) + + +@pytest.mark.asyncio +async def test_unified_interned_bytes_match_dict(tmp_path): + store = GraphSegmentUnifiedBackingStore(tmp_path, "b") + ha = store.put_segment("a", "system", "SYS") + hb = store.put_segment("b", "user", "hi") + store.add_node_manifest_interned( + "t0", 0, "profiling", [ha, hb], {"model": "m"}, False + ) + await store.finalize() + + from aiperf.graph.worker_materialize import ( + materialize_graph_request_unified, + materialize_graph_request_unified_bytes, + ) + + class _Ep: + streaming = False + extra = [] + use_server_token_count = False + use_legacy_max_tokens = False + + with GraphSegmentUnifiedClient(tmp_path, "b").open() as c: + payload = materialize_graph_request_unified( + c, "t0", 0, "profiling", use_legacy_max_tokens=False + ) + built = materialize_graph_request_unified_bytes( + c, "t0", 0, "profiling", use_legacy_max_tokens=False, endpoint=_Ep() + ) + assert payload["messages"] == [ + {"role": "system", "content": "SYS"}, + {"role": "user", "content": "hi"}, + ] + body, model, effective_stream = built + assert model == "m" + # Node stream False + global (_Ep) streaming False -> effective wire False. + assert effective_stream is False + assert orjson.loads(body)["messages"] == payload["messages"] + + +@pytest.mark.asyncio +async def test_interned_bytes_self_consistent_on_real_fixture(tmp_path, monkeypatch): + """Dict path == parsed bytes path == pool-derived messages on a real fixture. + + Grounds the interned bytes fast path three ways on a real weka parse: the + materialized ``messages`` equal the segment pool's own ``(role, content)`` + walk of the node's hex path (not just fn-vs-fn agreement), and the bytes + body's parsed JSON equals the dict path + run-level options EXACTLY (the + worker's parity contract, cache-bust NONE = the verbatim-bytes fast path). + """ + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + from aiperf.common.enums import CacheBustTarget + from aiperf.common.models import EndpointInfo + from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace + from aiperf.dataset.graph.segment_ir.store_builder import ( + _prompt_segment_ids, + _trie_llm_nodes, + build_unified_trie_store_interned, + trie_node_ordinals, + ) + from aiperf.graph.worker_materialize import ( + apply_run_level_payload_options, + materialize_graph_request_unified, + materialize_graph_request_unified_bytes, + ) + from aiperf.plugin.enums import EndpointType + + fx = Path(__file__).parent / "fixtures" / "weka_subagent.json" + parsed = from_weka_trace(fx) + trace = parsed.traces[0] + nodes = _trie_llm_nodes(parsed, trace) + ordinals = trie_node_ordinals(nodes) + node_id = next(nid for nid, n in nodes.items() if _prompt_segment_ids(n)) + ordinal = ordinals[node_id] + path = _prompt_segment_ids(nodes[node_id]) + + # Pool-derived ground truth: the node's hex path walked directly against + # the parse's content-addressed pool. + pool = parsed.segment_pool._by_id + expected_messages = [ + {"role": pool[sid].role, "content": pool[sid].content} for sid in path + ] + + # A real EndpointInfo (NONE cache-bust = the verbatim-bytes fast path). + # base_url is a read-only property on EndpointInfo -- the settable field is + # base_urls. + endpoint = EndpointInfo( + type=EndpointType.CHAT, + base_urls=["http://localhost:8000/v1/chat"], + streaming=True, + extra=[], + use_server_token_count=False, + use_legacy_max_tokens=False, + cache_bust=CacheBustTarget.NONE, + ) + + us = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id="b") + await build_unified_trie_store_interned(parsed, us) + + with GraphSegmentUnifiedClient(tmp_path, "b").open() as uc: + payload = materialize_graph_request_unified(uc, trace.id, ordinal, "profiling") + built = materialize_graph_request_unified_bytes( + uc, + trace.id, + ordinal, + "profiling", + use_legacy_max_tokens=False, + endpoint=endpoint, + ) + + assert payload is not None + assert payload["messages"] == expected_messages + # Mirror the worker's dict path: extract the recorded per-node override from + # the materialized payload BEFORE applying run-level options, so the dict + # baseline resolves ``stream`` the same way the bytes path does. + env_stream = payload.get("stream") + stream_override = bool(env_stream) if env_stream is not None else None + apply_run_level_payload_options(payload, endpoint, stream_override=stream_override) + + assert built is not None + body, model, effective_stream = built + assert orjson.loads(body) == payload + assert model == payload.get("model") + assert effective_stream == payload["stream"] diff --git a/tests/unit/graph/test_warmup_boundary_rewrite.py b/tests/unit/graph/test_warmup_boundary_rewrite.py new file mode 100644 index 0000000000..0ab563956c --- /dev/null +++ b/tests/unit/graph/test_warmup_boundary_rewrite.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TC1 — boundary-turn warmup (``rewrite_for_warmup``) + warmup phase routing. + +The AgentX-parity auto-warmup contract (``timing.config``): the WARMUP phase +dispatches exactly ONE priming credit per chain live at t* -- the chain's +boundary turn (the last node recorded before t*) -- instead of replaying the +ENTIRE post-t* workload at ``max_tokens=1``. Chains are the per-session linear +paths the trie node ids encode (root chain + each subagent chain). These tests +pin: + +* the pure rewrite on a two-chain (root + subagent) trace: exactly the two + boundary nodes, START-rooted, zero leading offsets, fan-in inputs cleared; +* liveness: a chain entirely pre-t* is skipped; t*<=0 yields an EMPTY graph; +* the strategy's WARMUP phase variant dispatches ONLY the boundary nodes while + the PROFILING variant at the same t* dispatches the post-t* chop set; +* a zero t* window warmup issues no credits and finalizes immediately. +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path +from typing import Any + +import pytest + +from aiperf.common.enums import CreditPhase +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.graph.analysis import trace_duration_us +from aiperf.timing.strategies.graph_ir_replay import ( + GraphIRReplayStrategy, + rewrite_for_warmup, +) + +_SUBAGENT_FIX = Path(__file__).parent / "fixtures" / "weka_subagent.json" + + +def _arrivals(parsed: Any) -> dict[str, int]: + return {nid: node.arrival_offset_us for nid, node in parsed.graph.nodes.items()} + + +def _midpoint(parsed: Any, a: str, b: str) -> float: + arr = _arrivals(parsed) + return (arr[a] + arr[b]) / 2.0 + + +# --------------------------------------------------------------------------- +# rewrite_for_warmup — pure rewrite +# --------------------------------------------------------------------------- + + +def test_two_live_chains_yield_exactly_the_two_boundary_nodes(): + """t* between the subagent's turns: root boundary trace_sub_n2s1:0 + subagent boundary agent_001:0. + + Chains on the subagent fixture: root = [trace_sub_n2s1:0, trace_sub_n2s1:1], + subagent = [agent_001:0, agent_001:1]. With t* between agent_001:0 and + agent_001:1 both chains are live (each has a pre-t* and a post-t* node), so + warmup primes each chain's LAST pre-t* turn -- and nothing else. + """ + parsed = from_weka_trace(str(_SUBAGENT_FIX)) + t_star_us = _midpoint(parsed, "agent_001:0", "agent_001:1") + + warmup = rewrite_for_warmup(parsed, t_star_us) + + assert set(warmup.graph.nodes) == {"trace_sub_n2s1:0", "agent_001:0"} + # START-rooted with ZERO leading offsets (warmup bursts, never replays + # recorded gaps): one plain START edge per boundary node, no delay fields. + assert {(e.source, e.target) for e in warmup.graph.edges} == { + ("START", "trace_sub_n2s1:0"), + ("START", "agent_001:0"), + } + for edge in warmup.graph.edges: + assert not edge.min_start_delay_us + assert not edge.delay_after_predecessor_us + assert not edge.delay_after_predecessor_start_us + assert not edge.delay_after_predecessor_first_token_us + for node in warmup.graph.nodes.values(): + assert node.inputs == [] + assert not node.min_start_delay_us + + +def test_chain_entirely_pre_tstar_is_not_live(): + """t* past the whole subagent chain: only the root chain is primed.""" + parsed = from_weka_trace(str(_SUBAGENT_FIX)) + t_star_us = _midpoint(parsed, "agent_001:1", "trace_sub_n2s1:1") + + warmup = rewrite_for_warmup(parsed, t_star_us) + + # Subagent chain [agent_001:0, agent_001:1] is entirely pre-t* (nothing of it + # is profiled) -> no priming; root chain boundary is trace_sub_n2s1:0 + # (trace_sub_n2s1:1 is post-t*). + assert set(warmup.graph.nodes) == {"trace_sub_n2s1:0"} + + +def test_boundary_nodes_keep_identity_and_dispatch_payload(): + """Boundary nodes keep ids, trie envelope, and dispatch body fields verbatim. + + The worker resolves the unmodified catalog ordinal from the node id and + materializes the recorded prompt from the store; only ``inputs`` and the + leading offset are cleared by the rewrite. + """ + parsed = from_weka_trace(str(_SUBAGENT_FIX)) + t_star_us = _midpoint(parsed, "agent_001:0", "agent_001:1") + + warmup = rewrite_for_warmup(parsed, t_star_us) + + for nid, node in warmup.graph.nodes.items(): + original = parsed.graph.nodes[nid] + assert node.metadata == original.metadata + assert node.extra_body == original.extra_body + assert node.model == original.model + assert node.max_tokens == original.max_tokens + assert node.arrival_offset_us == original.arrival_offset_us + + +@pytest.mark.parametrize("t_star_us", [0, -1.0]) # fmt: skip +def test_tstar_zero_yields_empty_warmup_graph(t_star_us): + """t*<=0 (full native replay / zero-duration trace) => EMPTY warmup graph.""" + parsed = from_weka_trace(str(_SUBAGENT_FIX)) + warmup = rewrite_for_warmup(parsed, t_star_us) + assert warmup.graph.nodes == {} + assert warmup.graph.edges == [] + + +# --------------------------------------------------------------------------- +# strategy warmup phase routing — dispatch sets +# --------------------------------------------------------------------------- + + +class _Config: + """Minimal per-phase config stub (mirrors test_lane_fanout_recycle).""" + + timing_mode = None + + def __init__(self, *, phase: CreditPhase | None = None) -> None: + self.phase = phase + self.concurrency = None + self.expected_num_sessions = None + self.total_expected_requests = None + self.expected_duration_sec = None + + +class _StubIssuer: + """Issuer whose graph credits resolve immediately via the return observer.""" + + def __init__(self) -> None: + self.observer = None + self.issued: list[Any] = [] + + def bind(self, strategy: GraphIRReplayStrategy) -> None: + self.observer = strategy._on_graph_return + + async def issue_graph_credit(self, turn: Any) -> bool: + self.issued.append(turn) + observer = self.observer + loop = asyncio.get_running_loop() + loop.call_soon(lambda: observer(turn, None, False)) + return True + + def mark_graph_sending_complete(self) -> None: ... + + def graph_all_returned(self) -> bool: + return True + + def set_graph_all_returned_event(self) -> None: ... + + +def _strategy( + parsed: Any, phase: CreditPhase | None, ratio: float +) -> tuple[GraphIRReplayStrategy, _StubIssuer]: + issuer = _StubIssuer() + strategy = GraphIRReplayStrategy( + config=_Config(phase=phase), + credit_issuer=issuer, + parsed_graph=parsed, + register_observer=lambda _obs: None, + start_min_ratio=ratio, + start_max_ratio=ratio, + ) + issuer.bind(strategy) + return strategy, issuer + + +def _issued_node_ids(issuer: _StubIssuer) -> set[str]: + # Node id ``{scope}:{turn}`` is recovered from the credit's own legacy-shaped + # identity: conversation_id (``{trace}`` root / ``{trace}::{scope}`` child) + # + turn_index (the node's 0-based turn). + node_ids = set() + for turn in issuer.issued: + trace, sep, scope = turn.conversation_id.partition("::") + node_ids.add(f"{scope if sep else trace}:{turn.turn_index}") + return node_ids + + +@pytest.mark.asyncio +async def test_warmup_phase_dispatches_only_boundary_turns(): + """WARMUP at a two-live-chain t* dispatches exactly {trace_sub_n2s1:0, agent_001:0} at 'warmup'.""" + parsed = from_weka_trace(str(_SUBAGENT_FIX)) + duration = trace_duration_us(parsed, parsed.traces[0]) + ratio = _midpoint(parsed, "agent_001:0", "agent_001:1") / duration + + strategy, issuer = _strategy(parsed, CreditPhase.WARMUP, ratio) + await strategy.execute_phase() + + assert _issued_node_ids(issuer) == {"trace_sub_n2s1:0", "agent_001:0"} + assert all(turn.phase_variant == "warmup" for turn in issuer.issued) + assert strategy.completed_traces == 1 + assert strategy.errored_traces == 0 + + +@pytest.mark.asyncio +async def test_profiling_phase_at_same_tstar_dispatches_post_tstar_set(): + """Control: PROFILING at the same t* runs the chop survivors, not boundaries.""" + parsed = from_weka_trace(str(_SUBAGENT_FIX)) + duration = trace_duration_us(parsed, parsed.traces[0]) + ratio = _midpoint(parsed, "agent_001:0", "agent_001:1") / duration + + strategy, issuer = _strategy(parsed, CreditPhase.PROFILING, ratio) + await strategy.execute_phase() + + assert _issued_node_ids(issuer) == {"agent_001:1", "trace_sub_n2s1:1"} + assert all(turn.phase_variant == "profiling" for turn in issuer.issued) + + +@pytest.mark.asyncio +async def test_warmup_with_zero_window_issues_nothing_and_finalizes(): + """t*=0: empty warmup graph -- no credit issued, phase completes cleanly.""" + parsed = from_weka_trace(str(_SUBAGENT_FIX)) + strategy, issuer = _strategy(parsed, CreditPhase.WARMUP, 0.0) + + await asyncio.wait_for(strategy.execute_phase(), timeout=5.0) + + assert issuer.issued == [] + assert strategy.completed_traces == strategy.admitted_traces + assert strategy.errored_traces == 0 diff --git a/tests/unit/graph/test_warmup_handoff_consume.py b/tests/unit/graph/test_warmup_handoff_consume.py new file mode 100644 index 0000000000..ea8eaea1fc --- /dev/null +++ b/tests/unit/graph/test_warmup_handoff_consume.py @@ -0,0 +1,448 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""PROFILING consumes the extended-warmup handoff -- frontier resume tests. + +Pins the re-cut contract: a profiling pass-0 lane with a handoff entry resumes +the handed-off template at its execution frontier (executed nodes never +re-dispatch; residual re-roots), recycle passes are untouched, and the full +warmup -> teardown -> profiling transition works end to end on real strategy +objects sharing one conversation source (regression guard for the +adapter-tests-skip-validator gap: this test path exercises the REAL phase +transition, not just node-level fields). +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path +from typing import Any + +import msgspec +import pytest + +from aiperf.common.enums import CacheBustTarget, CreditPhase +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.timing.graph_warmup_handoff import GraphWarmupHandoff, LaneHandoff +from aiperf.timing.strategies.cache_bust import build_trace_instance_marker +from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy + +_FIX = Path(__file__).parent / "fixtures" / "weka_min.json" + + +def _corpus(n_traces: int) -> Any: + parsed = from_weka_trace(str(_FIX)) + base = parsed.traces[0] + traces = [msgspec.structs.replace(base, id=f"t-{i}") for i in range(n_traces)] + return msgspec.structs.replace(parsed, traces=traces) + + +class _Config: + timing_mode = None + phase = CreditPhase.WARMUP + + def __init__(self, *, concurrency: int | None = None) -> None: + self.concurrency = concurrency + self.expected_num_sessions = None + self.total_expected_requests = None + self.expected_duration_sec = None + + +class _StubIssuer: + """Resolve every graph credit instantly on the next loop tick.""" + + def __init__(self) -> None: + self.observer = None + self.issued: list[Any] = [] + + def bind(self, strategy: GraphIRReplayStrategy) -> None: + self.observer = strategy._on_graph_return + + async def issue_graph_credit(self, credit: Any) -> bool: + self.issued.append(credit) + observer = self.observer + asyncio.get_running_loop().call_soon(lambda: observer(credit, None, False)) + return True + + def mark_graph_sending_complete(self) -> None: ... + + def graph_all_returned(self) -> bool: + return True + + def set_graph_all_returned_event(self) -> None: ... + + +class _ParkAfterIssuer: + """Resolve the first ``park_after`` graph credits instantly, park the rest. + + Parking stalls the pressure lanes so the ``wait_for`` duration timer fires + deterministically instead of recycling unboundedly fast under the + instant-sleep test fixtures. + """ + + def __init__(self, park_after: int | None = None) -> None: + self.observer = None + self.issued: list[Any] = [] + self._park_after = park_after + + def bind(self, strategy: GraphIRReplayStrategy) -> None: + self.observer = strategy._on_graph_return + + async def issue_graph_credit(self, credit: Any) -> bool: + self.issued.append(credit) + if self._park_after is not None and len(self.issued) > self._park_after: + return True # parked: never resolved + observer = self.observer + asyncio.get_running_loop().call_soon(lambda: observer(credit, None, False)) + return True + + def mark_graph_sending_complete(self) -> None: ... + + def graph_all_returned(self) -> bool: + return True + + def set_graph_all_returned_event(self) -> None: ... + + +class _FakeSource: + """Duck-typed graph channel (cross-phase warmup-handoff slot).""" + + def __init__(self) -> None: + self.warmup_handoff = None + + +def _warmup_strategy( + parsed: Any, + *, + duration: float | None, + park_after: int | None = None, + graph_channel: Any = None, + concurrency: int = 1, +) -> tuple[GraphIRReplayStrategy, _ParkAfterIssuer]: + issuer = _ParkAfterIssuer(park_after=park_after) + strategy = GraphIRReplayStrategy( + config=_Config(concurrency=concurrency), + graph_channel=graph_channel, + credit_issuer=issuer, + parsed_graph=parsed, + register_observer=lambda _obs: None, + start_min_ratio=0.5, + start_max_ratio=0.5, + t_star_random_seed=1234, + cache_pressure_duration_s=duration, + dispatch_timeout_s=2.0, + ) + issuer.bind(strategy) + return strategy, issuer + + +def _profiling_strategy(parsed, *, handoff, session_cap=1, concurrency=1): + issuer = _StubIssuer() + config = _Config(concurrency=concurrency) + config.phase = CreditPhase.PROFILING + config.expected_num_sessions = session_cap + strategy = GraphIRReplayStrategy( + config=config, + credit_issuer=issuer, + parsed_graph=parsed, + register_observer=lambda _obs: None, + start_min_ratio=0.5, + start_max_ratio=0.5, + t_star_random_seed=1234, + warmup_handoff=handoff, + ) + issuer.bind(strategy) + return strategy, issuer + + +@pytest.mark.asyncio +async def test_profiling_pass0_skips_executed_nodes(): + parsed = _corpus(1) + all_nodes = set(parsed.graph.nodes) + executed = {sorted(all_nodes)[0]} # first node id, deterministic + handoff = GraphWarmupHandoff( + lanes={ + 0: LaneHandoff( + template_trace_id="t-0", + instance_id="t-0#0.p0", + t_star_us=0.0, + executed_node_ids=frozenset(executed), + return_wall_us={nid: 0.0 for nid in executed}, + ) + }, + drain_end_wall_us=1.0, + corpus_cursor=1, + pressure_lane_count=1, + ) + strategy, issuer = _profiling_strategy(parsed, handoff=handoff) + await strategy.execute_phase() + + pass0 = [c for c in issuer.issued if c.trace_id == handoff.lanes[0].instance_id] + dispatched = {getattr(c, "node_ordinal", None) for c in pass0} + executed_ordinals = {strategy._catalog.catalog["t-0"][nid] for nid in executed} + assert pass0, "pass-0 must dispatch the surviving frontier" + assert not (dispatched & executed_ordinals), ( + "executed nodes must NOT re-dispatch in the handoff resume" + ) + + +@pytest.mark.asyncio +async def test_profiling_lane_runs_handoff_template_not_lane_assignment(): + """Pressure recycled the lane onto t-1; profiling pass-0 must resume t-1.""" + parsed = _corpus(2) # lane 0's pass-0 assignment would be t-0 + handoff = GraphWarmupHandoff( + lanes={ + 0: LaneHandoff( + template_trace_id="t-1", + instance_id="t-1#0.p0", + t_star_us=0.0, + executed_node_ids=frozenset(), + return_wall_us={}, + ) + }, + drain_end_wall_us=1.0, + corpus_cursor=1, + pressure_lane_count=1, + ) + strategy, issuer = _profiling_strategy(parsed, handoff=handoff, session_cap=1) + await strategy.execute_phase() + assert any(c.trace_id.startswith("t-1#0.") for c in issuer.issued) + + +@pytest.mark.asyncio +async def test_profiling_without_handoff_is_unchanged(): + parsed = _corpus(1) + strategy, issuer = _profiling_strategy(parsed, handoff=None) + await strategy.execute_phase() + # One single-pass instance, ``t-0::{nonce}`` (nonce minted fresh per instance). + ids = {c.trace_id for c in issuer.issued} + assert len(ids) == 1 + assert ids.pop().split("::", 1)[0] == "t-0" + + +@pytest.mark.asyncio +async def test_warmup_pressure_to_profiling_end_to_end_no_refire(): + """Full transition on real strategies sharing one source. + + Warmup (priming + pressure, park-stalled) -> teardown stashes handoff -> + profiling built with the popped handoff -> pass 0 refires NO node the + pressure stage executed on the live instance. + + ``park_after=2`` is tuned to keep the pass-0 pressure instance LIVE at drain + with one PRESSURE credit resolved: priming issues one credit (resolves), the + first pressure credit resolves, the second parks. + """ + parsed = _corpus(1) + source = _FakeSource() + warmup, warmup_issuer = _warmup_strategy( + parsed, duration=0.2, park_after=2, graph_channel=source + ) + await warmup.execute_phase() + await warmup.teardown_phase() + handoff = source.warmup_handoff + assert handoff is not None and handoff.lanes + + entry = handoff.lanes[0] + assert entry.executed_node_ids, "E2E must exercise a non-empty executed set" + strategy, issuer = _profiling_strategy(parsed, handoff=handoff) + await strategy.execute_phase() + + executed_ordinals = { + strategy._catalog.catalog[entry.template_trace_id][nid] + for nid in entry.executed_node_ids + } + # Resumed pass-0 credits carry the stashed live instance id (marker + # continuity), not a freshly minted ``.0`` id. + pass0 = [c for c in issuer.issued if c.trace_id == entry.instance_id] + assert pass0, "profiling pass-0 must dispatch the surviving frontier" + assert not ({getattr(c, "node_ordinal", None) for c in pass0} & executed_ordinals) + + +@pytest.mark.asyncio +async def test_resumed_pass0_reuses_pressure_instance_id_for_marker_continuity(): + """Resumed pass-0 credits reuse the stashed live instance id verbatim. + + The per-instance cache-bust marker digests ``credit.trace_id`` + (``build_trace_instance_marker``). Reusing the pressure instance's id at the + handoff resume keeps that marker byte-identical across the WARMUP -> PROFILING + boundary, so the KV the pressure stage built at the id transfers instead of + cold-prefilling behind a fresh ``.0`` marker. + """ + parsed = _corpus(1) + source = _FakeSource() + warmup, _ = _warmup_strategy( + parsed, duration=0.2, park_after=2, graph_channel=source + ) + await warmup.execute_phase() + await warmup.teardown_phase() + handoff = source.warmup_handoff + assert handoff is not None and handoff.lanes + + entry = handoff.lanes[0] + strategy, issuer = _profiling_strategy(parsed, handoff=handoff) + await strategy.execute_phase() + + resumed = [c for c in issuer.issued if c.trace_id == entry.instance_id] + assert resumed, "resumed pass-0 credits must carry the stashed live instance id" + + bid = "seed-1234" + assert build_trace_instance_marker( + bid, resumed[0].trace_id, target=CacheBustTarget.FIRST_TURN_PREFIX + ) == build_trace_instance_marker( + bid, entry.instance_id, target=CacheBustTarget.FIRST_TURN_PREFIX + ) + + +@pytest.mark.asyncio +async def test_bounded_profiling_recycle_continues_from_handoff_cursor(): + """A bounded profiling recycle resumes the corpus draw at the handoff cursor. + + Corpus t-0,t-1,t-2; lane 0 resumes the handoff template t-0, then recycles. + ``corpus_cursor=2`` means the pressure stage last drew position 1, so the + freed lane continues at position 2 (t-2) -- NOT the pass-0 default cursor of + 1 (t-1) a fresh profiling run would use (agentx shared-sampler parity). + """ + parsed = _corpus(3) + handoff = GraphWarmupHandoff( + lanes={ + 0: LaneHandoff( + template_trace_id="t-0", + instance_id="t-0::press0", + t_star_us=0.0, + executed_node_ids=frozenset(), + return_wall_us={}, + ) + }, + drain_end_wall_us=1.0, + corpus_cursor=2, + pressure_lane_count=1, + ) + strategy, issuer = _profiling_strategy(parsed, handoff=handoff, session_cap=2) + await strategy.execute_phase() + + templates = {c.trace_id.split("::", 1)[0] for c in issuer.issued} + assert "t-2" in templates, ( + "the freed lane must recycle from the handoff cursor onto t-2" + ) + assert "t-1" not in templates, ( + "t-1 (the pass-0 default cursor draw) must NOT be served" + ) + + +@pytest.mark.asyncio +async def test_single_pass_profiling_ignores_handoff_cursor(): + """Single-pass profiling covers the whole corpus, ignoring the handoff cursor. + + With no stop conditions the lanes do ONE corpus pass whose termination check + (``next_index >= len(traces)``) encodes cover-the-corpus-once. Carrying the + handoff's ``corpus_cursor=999`` would skip past the corpus and drop templates, + so single-pass deliberately keeps its own cursor. + """ + parsed = _corpus(2) + handoff = GraphWarmupHandoff( + lanes={}, drain_end_wall_us=0.0, corpus_cursor=999, pressure_lane_count=0 + ) + strategy, issuer = _profiling_strategy(parsed, handoff=handoff, session_cap=None) + await strategy.execute_phase() + + templates = {c.trace_id.split("::", 1)[0] for c in issuer.issued} + assert {"t-0", "t-1"} <= templates, ( + f"single-pass must claim every corpus position once; saw {templates}" + ) + + +@pytest.mark.asyncio +async def test_single_pass_lane_keeps_pass0_plan_despite_pressure_lane_count(): + """Single-pass mode NEVER fresh-starts: cover-the-corpus-once wins. + + A drained-empty pressure lane (pressure_lane_count=1, no entry) in a + single-pass profiling run (no stop conditions) keeps its normal pass-0 + assignment -- no .f0 instance, no cursor draw. + """ + parsed = _corpus(2) + handoff = GraphWarmupHandoff( + lanes={}, # lane 0 was a pressure lane but drained empty + drain_end_wall_us=1.0, + corpus_cursor=999, # a fresh-start draw would jump here; single-pass ignores it + pressure_lane_count=1, + ) + strategy, issuer = _profiling_strategy(parsed, handoff=handoff, session_cap=None) + await strategy.execute_phase() + + ids = {c.trace_id for c in issuer.issued} + templates = {trace_id.split("::", 1)[0] for trace_id in ids} + # The fresh-start gate is BOUNDED-mode-only: single-pass keeps each lane's + # normal pass-0 assignment (lane 0 -> t-0), ignoring the handoff cursor. + assert "t-0" in templates, ( + f"single-pass lane must resume its plain pass-0 assignment; saw {ids}" + ) + # cover-the-corpus-once still holds despite the nonzero pressure lane count + assert {"t-0", "t-1"} <= templates, ( + f"single-pass must claim every corpus position once; saw {templates}" + ) + + +@pytest.mark.asyncio +async def test_fresh_start_lane_runs_full_replay_of_cursor_template(): + """A pressure lane with no live instance at drain fresh-starts in profiling. + + AgentX parity (_build_handoff_trajectories: empty lanes get a fresh + recycle conversation at turn 0): the lane must NOT re-run its t* resume + (pressure already replayed it -- measuring it again inflates cache-hit + stats); it draws the next cursor template and replays it in full. + """ + parsed = _corpus(3) + handoff = GraphWarmupHandoff( + lanes={}, # lane 0 was a pressure lane but completed at drain + drain_end_wall_us=1.0, + corpus_cursor=2, + pressure_lane_count=1, + ) + strategy, issuer = _profiling_strategy(parsed, handoff=handoff, session_cap=1) + await strategy.execute_phase() + # Fresh template drawn from the cursor (t-2), full t*=0 replay under a fresh + # ``t-2::{nonce}`` instance (the fresh-start flavor is logged, not id-encoded; + # the per-instance nonce already precludes any boundary-priming id collision). + pass0 = [c for c in issuer.issued if c.trace_id.split("::", 1)[0] == "t-2"] + all_ordinals = set(strategy._catalog.catalog["t-2"].values()) + assert pass0, "the fresh-started lane must dispatch the cursor template t-2" + assert {getattr(c, "node_ordinal", None) for c in pass0} == all_ordinals + # and the lane-assigned template's t* resume did NOT run + assert not any(c.trace_id.split("::", 1)[0] == "t-0" for c in issuer.issued) + + +@pytest.mark.asyncio +async def test_profiling_honors_handoff_lanes_past_session_clamp(): + """--num-conversations < concurrency must not drop drained lanes. + + ALL handoff trajectories dispatch; the session cap gates only + recycles. + """ + parsed = _corpus(2) + handoff = GraphWarmupHandoff( + lanes={ + 0: LaneHandoff( + template_trace_id="t-0", + instance_id="t-0#0.p0", + t_star_us=0.0, + executed_node_ids=frozenset(), + return_wall_us={}, + ), + 1: LaneHandoff( + template_trace_id="t-1", + instance_id="t-1#1.p0", + t_star_us=0.0, + executed_node_ids=frozenset(), + return_wall_us={}, + ), + }, + drain_end_wall_us=1.0, + corpus_cursor=2, + pressure_lane_count=2, + ) + # session_cap=1 would normally clamp profiling to ONE lane + strategy, issuer = _profiling_strategy( + parsed, handoff=handoff, session_cap=1, concurrency=2 + ) + await strategy.execute_phase() + ids = {c.trace_id for c in issuer.issued} + assert "t-0#0.p0" in ids and "t-1#1.p0" in ids diff --git a/tests/unit/graph/test_warmup_routing_and_seed.py b/tests/unit/graph/test_warmup_routing_and_seed.py new file mode 100644 index 0000000000..2292cc9489 --- /dev/null +++ b/tests/unit/graph/test_warmup_routing_and_seed.py @@ -0,0 +1,372 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Warmup routing + seed-determinism + auto-detect precedence (Task F2 / M-seed / F3). + +Covers the ``timing.config`` deliverables: + +* a graph workload's WARMUP phase routes through ``TimingMode.GRAPH_IR`` (not the + linear ``REQUEST_RATE`` path); +* ``resolve_graph_content_seed`` is the run ``--random-seed`` verbatim (None + when unset -- no weka-specific fallback); threading the same explicit seed + into two independent full parses of the same run synthesizes byte-identical + weka content (the determinism the in-process build parse and any + spawn-started pool worker rely on), and distinct seeds diverge; +* ``from_run`` honors an explicit, graph-incompatible ``--custom-dataset-type`` + over the file-sniff auto-detection (precedence), while a bare weka file with + no pinned format still auto-routes to GRAPH_IR. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from pytest import param + +from aiperf.common.enums import CreditPhase, DatasetFormat +from aiperf.config.flags.cli_config import CLIConfig +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.models import ParsedGraph +from aiperf.plugin.enums import ArrivalPattern, TimingMode +from aiperf.timing.config import TimingConfig, resolve_graph_content_seed +from tests.unit.conftest import make_run_from_cli + +WEKA_MIN = Path(__file__).parent / "fixtures" / "weka_min.json" + + +def _graph_run(**cli_overrides): + """Build a ``BenchmarkRun`` pointing at the weka fixture, with warmup set.""" + cfg = CLIConfig( + model_names=["test-model"], + input_file=str(WEKA_MIN), + warmup_request_count=2, + request_count=3, + **cli_overrides, + ) + return make_run_from_cli(cfg) + + +def _pool_contents(parsed: ParsedGraph) -> dict[str, tuple[str, str, str | None]]: + """Materialized segment content keyed by content-addressed segment id.""" + pool = parsed.segment_pool + assert pool is not None, "weka parse must surface the segment pool" + return {sid: (s.role, s.content, s.parent_id) for sid, s in pool._by_id.items()} + + +def test_graph_warmup_phase_routes_to_graph_ir() -> None: + run = _graph_run() + tc = TimingConfig.from_run(run) + warmup = [p for p in tc.phase_configs if p.phase == CreditPhase.WARMUP] + profiling = [p for p in tc.phase_configs if p.phase == CreditPhase.PROFILING] + assert warmup, "expected a warmup phase" + assert all(p.timing_mode == TimingMode.GRAPH_IR for p in warmup), ( + "graph warmup must use GRAPH_IR, not the linear REQUEST_RATE strategy" + ) + assert all(p.timing_mode == TimingMode.GRAPH_IR for p in profiling) + + +def test_agentx_window_injects_boundary_auto_warmup(monkeypatch) -> None: + """An active AgentX t* window (no explicit warmup flags) gets the auto-warmup phase. + + AgentX parity: with the t* window active (trajectory_start_max_ratio > 0) + the graph run ALWAYS primes each live chain's boundary turn before + profiling. The injected phase must carry NO stop-condition counts -- the + GraphIRReplayStrategy owns warmup completion (its ``rewrite_for_warmup`` + boundary graph fires exactly the live boundary turns, then drains). + """ + cfg = CLIConfig( + model_names=["test-model"], + input_file=str(WEKA_MIN), + request_count=3, + trajectory_start_min_ratio=0.25, + trajectory_start_max_ratio=0.75, + ) + run = make_run_from_cli(cfg) + tc = TimingConfig.from_run(run) + warmup = [p for p in tc.phase_configs if p.phase == CreditPhase.WARMUP] + assert len(warmup) == 1, "expected exactly one auto-injected warmup phase" + auto = warmup[0] + assert auto.timing_mode == TimingMode.GRAPH_IR + assert auto.total_expected_requests is None + assert auto.expected_num_sessions is None + assert auto.expected_duration_sec is None + assert auto.grace_period_sec == float("inf") + + +def test_default_graph_run_injects_no_auto_warmup(monkeypatch) -> None: + """Bare default = full replay: no t* window, no boundary auto-warmup. + + Pin the pressure duration to None explicitly: ``CLIConfig`` dual-writes it + globally, so xdist ordering can otherwise leak a duration in and inject a + (pressure-mode) warmup phase even with the t* window off. + """ + cfg = CLIConfig( + model_names=["test-model"], + input_file=str(WEKA_MIN), + request_count=3, + ) + run = make_run_from_cli(cfg) + tc = TimingConfig.from_run(run) + warmup = [p for p in tc.phase_configs if p.phase == CreditPhase.WARMUP] + assert warmup == [] + + +def test_pressure_duration_injects_warmup_even_at_tstar_zero(monkeypatch) -> None: + """A pressure duration forces the auto-warmup phase even with t* inactive. + + The extended (cache-pressure) warmup runs inside the WARMUP phase, so a + configured ``--agentic-cache-warmup-duration`` must inject the phase even + when the t* window is closed (max ratio 0.0, the bare default => empty + boundary priming). Mirrors the pressure gating in ``timing/config.py``. + """ + cfg = CLIConfig( + agentic_cache_warmup_duration=30.0, + model_names=["test-model"], + input_file=str(WEKA_MIN), + request_count=3, + ) + run = make_run_from_cli(cfg) + tc = TimingConfig.from_run(run) + warmup = [p for p in tc.phase_configs if p.phase == CreditPhase.WARMUP] + assert len(warmup) == 1, "pressure duration must inject a warmup phase at t*=0" + assert warmup[0].timing_mode == TimingMode.GRAPH_IR + + +def test_non_graph_warmup_phase_stays_request_rate() -> None: + cfg = CLIConfig( + model_names=["test-model"], + warmup_request_count=2, + request_count=3, + ) + run = make_run_from_cli(cfg) + tc = TimingConfig.from_run(run) + warmup = [p for p in tc.phase_configs if p.phase == CreditPhase.WARMUP] + assert warmup + assert all(p.timing_mode == TimingMode.REQUEST_RATE for p in warmup) + + +def test_resolve_graph_content_seed_is_the_run_seed() -> None: + run = _graph_run() + # The content seed IS the AIPerf run seed -- no weka-specific fallback. A + # default single run leaves --random-seed unset, so the content seed is None + # (ambient global RNG), stable across calls for the same run. + assert run.random_seed is None + assert resolve_graph_content_seed(run) is None + assert resolve_graph_content_seed(run) == resolve_graph_content_seed(run) + + +def test_resolve_graph_content_seed_honors_explicit_run_seed() -> None: + run = _graph_run(random_seed=7) + assert resolve_graph_content_seed(run) == 7 + + +def test_same_seed_two_parses_synthesize_byte_identical_content() -> None: + """Two independent full parses under the SAME explicit seed match byte-for-byte. + + Mirrors the content-determinism contract: the DatasetManager's in-process + build parse and any spawn-started pool worker parse the weka file + separately; threading the SAME seed makes the synthesized segment content + byte-identical. Compares the segment pool's real + materialized ``(role, content, parent_id)`` image -- NOT + ``TraceRecord.replay_outputs``, which is always empty on the weka path. + """ + parse_a = from_weka_trace(str(WEKA_MIN), content_root_seed=7) + parse_b = from_weka_trace(str(WEKA_MIN), content_root_seed=7) + assert parse_a.segment_pool is not None and parse_a.segment_pool._by_id + assert _pool_contents(parse_a) == _pool_contents(parse_b) + + +def test_different_seed_parses_synthesize_different_content() -> None: + """Distinct seeds synthesize distinct bytes (byte-identity is falsifiable).""" + parse_a = from_weka_trace(str(WEKA_MIN), content_root_seed=7) + parse_b = from_weka_trace(str(WEKA_MIN), content_root_seed=8) + assert _pool_contents(parse_a) != _pool_contents(parse_b) + + +def test_explicit_incompatible_format_overrides_weka_sniff() -> None: + """An explicit, graph-incompatible --custom-dataset-type beats the sniff. + + A user who pins ``--custom-dataset-type multi_turn`` on a file that happens + to sniff as weka must NOT be silently rerouted to the graph pipeline. + """ + run = _graph_run(custom_dataset_type=DatasetFormat.MULTI_TURN) + tc = TimingConfig.from_run(run) + profiling = [p for p in tc.phase_configs if p.phase == CreditPhase.PROFILING] + assert profiling + assert all(p.timing_mode != TimingMode.GRAPH_IR for p in profiling), ( + "an explicit incompatible --custom-dataset-type must suppress graph routing" + ) + + +def test_pressure_warmup_grace_is_min_of_duration_and_cap(monkeypatch) -> None: + """Pressure warmup drains at min(duration, cap) -- agentx MAX_WARMUP_GRACE parity. + + A pressure-mode warmup pins ``expected_duration_sec=None`` (the drain is + bounded by the finite grace, not a phase duration) and, absent an explicit + ``--warmup-grace-period``, drains for ``min(pressure duration, cap)`` where + ``cap`` is ``PRESSURE_DRAIN_GRACE_CAP`` (default 300). + """ + # duration 30 < cap 300 -> grace 30 + cfg = CLIConfig( + model_names=["test-model"], + input_file=str(WEKA_MIN), + request_count=3, + agentic_cache_warmup_duration=30.0, + ) + warmup = next( + p + for p in TimingConfig.from_run(make_run_from_cli(cfg)).phase_configs + if p.phase == CreditPhase.WARMUP + ) + assert warmup.grace_period_sec == 30.0 + assert warmup.expected_duration_sec is None + + # duration 900 > cap 300 -> grace 300 (clamped to the cap) + cfg = CLIConfig( + model_names=["test-model"], + input_file=str(WEKA_MIN), + request_count=3, + agentic_cache_warmup_duration=900.0, + ) + warmup = next( + p + for p in TimingConfig.from_run(make_run_from_cli(cfg)).phase_configs + if p.phase == CreditPhase.WARMUP + ) + assert warmup.grace_period_sec == 300.0 + + +def test_no_pressure_auto_warmup_keeps_infinite_grace() -> None: + """Without a pressure duration, the boundary-priming warmup still waits forever.""" + cfg = CLIConfig( + model_names=["test-model"], + input_file=str(WEKA_MIN), + request_count=3, + trajectory_start_min_ratio=0.25, + trajectory_start_max_ratio=0.75, + ) + warmup = next( + p + for p in TimingConfig.from_run(make_run_from_cli(cfg)).phase_configs + if p.phase == CreditPhase.WARMUP + ) + assert warmup.grace_period_sec == float("inf") + + +def test_pressure_supersedes_user_warmup_phase(monkeypatch) -> None: + """Graph + pressure: the warmup phase is MODE-OWNED (agentx parity). + + A user ``--warmup-request-count`` / warmup phase config is superseded by the + auto boundary-priming + pressure shape; agentx pins expected_duration_sec=None + unconditionally for its agentic warmup. + """ + cfg = CLIConfig( + agentic_cache_warmup_duration=30.0, + model_names=["test-model"], + input_file=str(WEKA_MIN), + warmup_request_count=2, + request_count=3, + ) + warmups = [ + p + for p in TimingConfig.from_run(make_run_from_cli(cfg)).phase_configs + if p.phase == CreditPhase.WARMUP + ] + assert len(warmups) == 1 + assert warmups[0].expected_duration_sec is None + assert warmups[0].total_expected_requests is None + assert warmups[0].arrival_pattern == ArrivalPattern.CONCURRENCY_BURST + + +def test_pressure_supersede_carries_user_grace_verbatim(monkeypatch) -> None: + """An explicit user warmup grace survives the supersede (agentx :227-229). + + The auto shape still owns duration (=None) and the burst pattern, but the + operator's explicit grace is honored verbatim -- it is NOT re-derived as + ``min(duration, cap)``. + """ + cfg = CLIConfig( + agentic_cache_warmup_duration=30.0, + model_names=["test-model"], + input_file=str(WEKA_MIN), + warmup_duration=10.0, + warmup_grace_period=45.0, + request_count=3, + ) + warmup = next( + p + for p in TimingConfig.from_run(make_run_from_cli(cfg)).phase_configs + if p.phase == CreditPhase.WARMUP + ) + assert warmup.grace_period_sec == 45.0 # NOT min(30, 300) + assert warmup.expected_duration_sec is None # still mode-owned (superseded) + + +def test_user_warmup_untouched_without_pressure() -> None: + """No pressure duration: explicit user warmup phases build exactly as today.""" + cfg = CLIConfig( + model_names=["test-model"], + input_file=str(WEKA_MIN), + warmup_duration=10.0, + warmup_grace_period=45.0, + request_count=3, + ) + warmup = next( + p + for p in TimingConfig.from_run(make_run_from_cli(cfg)).phase_configs + if p.phase == CreditPhase.WARMUP + ) + # Built verbatim from the user phase: duration and explicit grace preserved, + # NOT nulled/superseded by the auto pressure shape. + assert warmup.expected_duration_sec == 10.0 + assert warmup.grace_period_sec == 45.0 + + +def test_adaptive_scale_on_graph_workload_raises() -> None: + """An explicit adaptive_scale phase on a graph workload fails loudly. + + Adaptive scaling and the recorded graph replay both want to own pacing and + concurrency; silently routing the phase to GRAPH_IR would discard the + user's explicit adaptive_scale choice, so ``from_run`` must reject the + combination up front. + """ + run = _graph_run( + adaptive_scale=True, + concurrency=4, + benchmark_duration=10.0, + adaptive_sustain_duration=5.0, + adaptive_scale_sla=["request_latency:p95:le:30000"], + ) + with pytest.raises( + ValueError, match="adaptive_scale is not supported for graph workloads" + ): + TimingConfig.from_run(run) + + +@pytest.mark.parametrize( + "cli_overrides", + [ + param({"request_rate": 5.0}, id="request_rate_poisson"), + param( + {"request_rate": 5.0, "arrival_pattern": ArrivalPattern.CONSTANT}, + id="request_rate_constant", + ), + param({"fixed_schedule": True}, id="fixed_schedule"), + param({"warmup_request_rate": 2.0}, id="warmup_request_rate"), + ], +) # fmt: skip +def test_rate_or_schedule_phase_on_graph_workload_raises(cli_overrides) -> None: + """A rate-controlled or fixed-schedule phase on a graph workload fails loudly. + + GRAPH_IR owns pacing (decision D2): forcing the graph strategy would + silently discard the user's explicit --request-rate / --user-centric-rate / + --fixed-schedule arrival timing, so ``from_run`` must reject the + combination up front -- warmup phases included (--warmup-request-rate is + equally discarded by the graph warmup). USER_CENTRIC is covered by the + same ``type != CONCURRENCY`` check but is unreachable from the CLI with a + file dataset (the converter rejects it earlier on the turn-mean floor), + so it has no param here. + """ + run = _graph_run(**cli_overrides) + with pytest.raises(ValueError, match="is not supported for graph workloads"): + TimingConfig.from_run(run) diff --git a/tests/unit/graph/test_warmup_variants.py b/tests/unit/graph/test_warmup_variants.py new file mode 100644 index 0000000000..a698bc03a0 --- /dev/null +++ b/tests/unit/graph/test_warmup_variants.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Warmup materialization on the segment-trie IR (Task 11 / adv2 F2). + +The trie build plane persists one profiling manifest per node plus the +content-addressed segment pool in the ONE interned unified store +(``build_unified_trie_store_interned``). A WARMUP credit reuses those profiling +bytes and materializes a payload with AgentX's unconditional 1-token output +cap, avoiding a duplicate warmup copy in the store. +""" + +from pathlib import Path + +import pytest + +from aiperf.common.environment import Environment +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.segment_ir.store_builder import ( + build_unified_trie_store_interned, +) +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) +from aiperf.graph.worker_materialize import materialize_graph_request_unified + +FIX = Path(__file__).parent / "fixtures" / "weka_min.json" + + +async def _build_trie_store( + fixture: Path, tmp_path: Path, benchmark_id: str +) -> tuple[dict[str, dict[str, int]], GraphSegmentUnifiedClient]: + """Ingest ``fixture`` and drain it into an opened interned unified reader.""" + parsed = from_weka_trace(str(fixture)) + assert parsed.segment_pool is not None, "trie ingest must attach a SegmentPool" + + store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id=benchmark_id + ) + addr = await build_unified_trie_store_interned(parsed, store) + + client = GraphSegmentUnifiedClient( + base_path=tmp_path, benchmark_id=benchmark_id + ).open() + return addr, client + + +@pytest.mark.asyncio +async def test_warmup_materializes_for_every_node(tmp_path): + addr, client = await _build_trie_store(FIX, tmp_path, "bw") + t0 = next(iter(addr)) + try: + for ordinal in addr[t0].values(): + # Warmup has no dedicated manifest -- it reuses the profiling bytes. + assert client.get_node_envelope(t0, ordinal, "warmup") is None + assert client.get_node_envelope(t0, ordinal, "profiling") is not None + warmup = materialize_graph_request_unified(client, t0, ordinal, "warmup") + assert warmup is not None + assert warmup["messages"] + finally: + client.close() + + +@pytest.mark.asyncio +async def test_warmup_materialization_caps_output_to_one_token(tmp_path): + """AgentX parity: warmup UNCONDITIONALLY caps output to 1 token.""" + addr, client = await _build_trie_store(FIX, tmp_path, "bw2") + t0 = next(iter(addr)) + # The first recorded turn (arrival ordinal 0) carries the recorded out=25. + ordinal = min(addr[t0].values()) + try: + warmup = materialize_graph_request_unified(client, t0, ordinal, "warmup") + profiling = materialize_graph_request_unified(client, t0, ordinal, "profiling") + + assert warmup is not None + assert profiling is not None + cap = Environment.GRAPH.WARMUP_MAX_OUTPUT_TOKENS + assert cap == 1, "WARMUP_MAX_OUTPUT_TOKENS must default to AgentX's 1" + # The trie manifest carries the recorded ``out`` on ``max_output_tokens``, + # endpoint-mapped to ``max_completion_tokens`` (modern default); warmup + # pops the modern field and unconditionally overwrites with the 1-token + # legacy cap while profiling keeps the recorded out (25). + assert warmup.get("max_tokens") == cap + assert "max_completion_tokens" not in warmup + assert profiling.get("max_completion_tokens") == 25 + assert "max_tokens" not in profiling + # Warmup reuses the EXACT profiling input prefix -- only the cap differs. + assert warmup["messages"] == profiling["messages"] + finally: + client.close() diff --git a/tests/unit/graph/test_weka_block_size_decode.py b/tests/unit/graph/test_weka_block_size_decode.py new file mode 100644 index 0000000000..23064bc3df --- /dev/null +++ b/tests/unit/graph/test_weka_block_size_decode.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weka content decode honors the TRACE's ``block_size``. + +The trie GEOMETRY (covered-block math) uses ``trace.block_size``; if the +production content callbacks instead decoded every hash block at the +synthesizer's hardcoded 64 tokens, a ``block_size: 32`` trace with ``in: 64`` +would synthesize +a 128-token prompt (2 covered blocks x 64) instead of 64, silently corrupting +every ISL / prefix-cache measurement on non-default-block-size corpora. These +tests pin that ``_default_callbacks`` (and therefore ``build_trie_graph``) +decode at the trace's own block size, mirroring the dynamo adapter's +``dynamo_recon_callbacks``. + +Hermetic: ``fake_tokenizer`` pins ``Tokenizer.from_pretrained`` to the +deterministic ``FakeTokenizer`` (decode = ``"tok$" * n_tokens``), so prompt +token counts are recoverable from content length with no real tokenizer. +""" + +from __future__ import annotations + +import orjson +import pytest + +from aiperf.dataset.graph.adapters.shared.content import CorpusContentSynthesizer +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.adapters.weka.trie_build import _default_callbacks +from aiperf.dataset.graph.segment_ir.envelope import read_prompt_segment_ids +from tests.harness.fake_tokenizer import TOKEN_LEN + +_BLOCK32_TRACE = { + "id": "trace_bs32", + "models": ["M"], + "block_size": 32, + "hash_id_scope": "local", + "requests": [ + { + "t": 0.0, + "type": "n", + "model": "M", + "in": 64, + "out": 8, + "hash_ids": [1, 2], + "api_time": 0.5, + }, + ], +} + + +@pytest.fixture(autouse=True) +def _fresh_synth_cache(): + """Isolate the process-level synthesizer cache from other tests.""" + CorpusContentSynthesizer.reset_worker_cache() + yield + CorpusContentSynthesizer.reset_worker_cache() + + +@pytest.mark.parametrize("block_size", [16, 32, 64]) +def test_default_callbacks_decode_at_trace_block_size( + fake_tokenizer: None, # noqa: ARG001 + block_size: int, +) -> None: + """Each hash id decodes to exactly ``block_size`` tokens (not the legacy 64).""" + callbacks = _default_callbacks( + "fake-tok", "coding", 0, trace_id="trace_bs", block_size=block_size + ) + assert len(callbacks.decode_block_tokens([1])) == block_size + assert len(callbacks.decode_block_tokens([1, 2])) == 2 * block_size + + +def test_build_block32_trace_synthesizes_recorded_isl( + fake_tokenizer: None, # noqa: ARG001 + tmp_path, +) -> None: + """A ``block_size=32`` trace with ``in=64`` yields a 64-token prompt (was 128).""" + trace_file = tmp_path / "bs32.json" + trace_file.write_bytes(orjson.dumps(_BLOCK32_TRACE)) + + parsed = from_weka_trace(trace_file, content_root_seed=0) + + # The trie route carries NO inline prompt; the content is in the segment + # pool, addressed via the node's prompt_segment_ids path. + node = parsed.graph.nodes["trace_bs32:0"] + messages = parsed.segment_pool.materialize(read_prompt_segment_ids(node)) + prompt_tokens = sum(len(m["content"]) // TOKEN_LEN for m in messages) + assert prompt_tokens == 64, ( + f"covered-count ISL must equal recorded in=64 at block_size=32, " + f"got {prompt_tokens} (a 64-token legacy decode doubles it to 128)" + ) + + +def test_from_weka_trace_threads_max_osl_to_dispatch( + fake_tokenizer: None, # noqa: ARG001 + tmp_path, +) -> None: + """``max_osl`` reaches the trie build through the parse seam (finding W2).""" + capped_trace = dict(_BLOCK32_TRACE, id="trace_capped") + capped_trace["requests"] = [dict(_BLOCK32_TRACE["requests"][0], out=5000)] + trace_file = tmp_path / "capped.json" + trace_file.write_bytes(orjson.dumps(capped_trace)) + + parsed = from_weka_trace(trace_file, content_root_seed=0, max_osl=100) + assert parsed.graph.nodes["trace_capped:0"].max_tokens == 100 diff --git a/tests/unit/graph/test_weka_content_knob_wiring.py b/tests/unit/graph/test_weka_content_knob_wiring.py new file mode 100644 index 0000000000..8757333642 --- /dev/null +++ b/tests/unit/graph/test_weka_content_knob_wiring.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""HF streaming store build threads the run content knobs. + +The eager parse (``parse_graph_workload``) resolves ``content_tokenizer`` / +``prompt_corpus`` / ``max_osl`` / ``idle_gap_cap_seconds`` from the run config +and threads them into the weka parse; the HF STREAMING store build must thread +the SAME knobs into its per-row workers or the streamed store synthesizes +builtin+"coding" bytes instead of the run-resolved content. Knobs resolve from +the run config so every parse of the same run agrees. These tests pin both +wiring seams: + +* ``stream_weka_trace_segment_payloads`` forwards the knobs into + ``parse_kwargs``; +* ``GraphStoreBuilder._build_graph_store_streaming`` resolves them from the run + exactly like the eager parse (the ONE ``resolve_graph_parse_context`` + resolution, spread verbatim into the stream entry). +""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest + +from aiperf.config.flags.cli_config import CLIConfig +from tests.unit.conftest import make_run_from_cli + +WEKA_MIN = Path(__file__).parent / "fixtures" / "weka_min.json" + + +def test_stream_segment_payloads_threads_content_knobs(monkeypatch) -> None: + """The streaming entry point forwards every run content knob to workers.""" + from aiperf.dataset.graph.adapters.weka import trace as weka_trace + from aiperf.dataset.graph.adapters.weka import trace_parallel + + captured: dict[str, Any] = {} + + def fake_iter(items, *, source_label, item_count, workers, parse_kwargs): # noqa: ANN001, ARG001 + captured.update(parse_kwargs) + yield "payload" + + monkeypatch.setattr(trace_parallel, "iter_item_segment_payloads", fake_iter) + monkeypatch.setattr(weka_trace, "_load_hf_rows", lambda *a, **k: iter(())) # noqa: ARG005 + + payloads = list( + weka_trace.stream_weka_trace_segment_payloads( + "org/weka-corpus", + content_root_seed=7, + content_tokenizer="run-tok", + prompt_corpus="sonnet", + max_osl=64, + idle_gap_cap_seconds=30.0, + ) + ) + + assert payloads == ["payload"] + assert captured["content_root_seed"] == 7 + assert captured["content_tokenizer"] == "run-tok" + assert captured["prompt_corpus"] == "sonnet" + assert captured["max_osl"] == 64 + assert captured["idle_gap_cap_seconds"] == 30.0 + + +@pytest.mark.asyncio +async def test_build_graph_store_streaming_resolves_run_knobs(monkeypatch) -> None: + """The streaming store build resolves knobs from the run like the eager plane.""" + from aiperf.dataset.graph.adapters.weka import trace as weka_trace + from aiperf.dataset.graph.store_build import GraphStoreBuilder + from aiperf.timing.config import ( + resolve_graph_content_seed, + resolve_graph_content_tokenizer, + ) + + run = make_run_from_cli( + CLIConfig( + model_names=["test-model"], + input_file=str(WEKA_MIN), + tokenizer_name="builtin", + prompt_corpus="sonnet", + synthesis_max_osl=64, + synthesis_idle_gap_cap=30.0, + ) + ) + + captured: dict[str, Any] = {} + + def fake_stream(path, **kwargs): # noqa: ANN001 + captured["path"] = path + captured.update(kwargs) + return iter(()) + + monkeypatch.setattr(weka_trace, "stream_weka_trace_segment_payloads", fake_stream) + + sentinel_catalog = {"trace": {"r_0": 0}} + + async def fake_trie(payloads, base_path): # noqa: ANN001, ARG001 + return sentinel_catalog, None + + stub = SimpleNamespace(run=run, _build_graph_store_streaming_trie=fake_trie) + + catalog, merged = await GraphStoreBuilder._build_graph_store_streaming( + stub, Path("org/weka-corpus"), Path("/tmp/unused"), "weka_trace" + ) + + assert catalog is sentinel_catalog + assert merged is None + assert captured["content_root_seed"] == resolve_graph_content_seed(run) + assert captured["content_tokenizer"] == resolve_graph_content_tokenizer(run) + assert captured["prompt_corpus"] == "sonnet" + assert captured["max_osl"] == 64 + assert captured["idle_gap_cap_seconds"] == 30.0 diff --git a/tests/unit/graph/test_weka_fidelity_tool_gates.py b/tests/unit/graph/test_weka_fidelity_tool_gates.py new file mode 100644 index 0000000000..2acd95f86b --- /dev/null +++ b/tests/unit/graph/test_weka_fidelity_tool_gates.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The offline fidelity tool's content knobs are threaded, not hardcoded. + +``tools/weka_trace_fidelity.py`` used to hardcode ``gpt2`` / ``"coding"`` / +no-seed into its trie rebuild, so any run built with different knobs (or even +the bare live-run defaults, which use the builtin tokenizer) spuriously failed +every content comparison. These tests lock the fix: :func:`build_recorded_trace` +defaults to the live-run knobs (builtin / coding / no seed) and honors explicit +``tokenizer_name`` / ``prompt_corpus`` / ``root_seed`` overrides, and the CLI +``--tokenizer`` / ``--corpus`` / ``--seed`` flags reach the rebuild. + +These live in the UNIT lane on purpose: the component-integration package +patches ``Tokenizer.from_pretrained`` to a FakeTokenizer, which flattens +tokenizer/seed-dependent content and makes the knobs indistinguishable. The +unit lane's fake-tokenizer fixture is opt-in, so the real builtin synthesizer +runs here. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +from tools.weka_trace_fidelity import _main, build_recorded_trace + +_MODEL = "M" + + +def _linear_trace() -> dict: + """A two-turn linear trace: r_1 extends r_0's hash prefix 1.0s after it ends.""" + return { + "id": "knobs", + "models": [_MODEL], + "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "n", "model": _MODEL, "in": 64, "out": 8, + "hash_ids": [1], "api_time": 1.0}, + {"t": 2.0, "type": "n", "model": _MODEL, "in": 128, "out": 8, + "hash_ids": [1, 2], "api_time": 1.0}, + ], + } # fmt: skip + + +def _write_trace(tmp_path: Path) -> Path: + trace_dir = tmp_path / "traces" + trace_dir.mkdir() + path = trace_dir / "knobs.json" + path.write_text(json.dumps(_linear_trace())) + return path + + +def _faithful_raw(tmp_path: Path, trace_file: Path) -> Path: + """A faithful default-knob export: genuine materialized prompts, warped timing. + + ``r_0`` dispatches at an arbitrary origin; ``r_1`` at origin + its warped + edge delay off ``r_0`` (the zero-latency-mock collapse of end-to-start onto + start-to-start), so both criteria pass when the proof rebuilds with the SAME + content knobs. + """ + recorded = build_recorded_trace(trace_file) + t0 = 1_000_000_000_000 + delay_s = recorded.nodes["knobs:1"].pred_delay_us["knobs:0"] / 1e6 + dispatch_ns = {"knobs:0": t0, "knobs:1": t0 + int(delay_s * 1e9)} + lines = [] + for nid, ns in dispatch_ns.items(): + lines.append( + json.dumps( + { + "metadata": { + "conversation_id": "knobs#0.0", + "x_request_id": f"{nid}::deadbeefdeadbeefdeadbeefdeadbeef", + "benchmark_phase": "profiling", + "request_start_ns": ns, + "credit_issued_ns": ns - 1_000_000, + }, + "payload": {"messages": recorded.nodes[nid].messages}, + } + ) + ) + raw = tmp_path / "profile_export_raw.jsonl" + raw.write_text("\n".join(lines) + "\n", encoding="utf-8") + return raw + + +def test_build_recorded_trace_content_knobs_thread_into_rebuild( + tmp_path: Path, +) -> None: + """Explicit seed/corpus knobs change the rebuilt content; the default equals + the explicit live-run knobs (builtin / coding / no seed).""" + trace_file = _write_trace(tmp_path) + + default = build_recorded_trace(trace_file) + explicit = build_recorded_trace( + trace_file, tokenizer_name="builtin", prompt_corpus="coding", root_seed=None + ) + seeded = build_recorded_trace(trace_file, root_seed=7) + sonnet = build_recorded_trace(trace_file, prompt_corpus="sonnet") + + assert {n: default.nodes[n].messages for n in default.nodes} == { + n: explicit.nodes[n].messages for n in explicit.nodes + } + assert any( + default.nodes[n].messages != seeded.nodes[n].messages for n in default.nodes + ) + assert any( + default.nodes[n].messages != sonnet.nodes[n].messages for n in default.nodes + ) + + +def test_main_content_knob_flags_thread_into_proof(tmp_path: Path) -> None: + """The ``--tokenizer`` / ``--corpus`` / ``--seed`` CLI flags reach + ``build_recorded_trace``: a default-knob export passes under explicit default + flags but fails content under a different ``--seed``.""" + trace_file = _write_trace(tmp_path) + raw = _faithful_raw(tmp_path, trace_file) + base_argv = ["--raw", str(raw), "--trace-dir", str(trace_file.parent)] + + assert _main([*base_argv, "--tokenizer", "builtin", "--corpus", "coding"]) == 0 + assert _main([*base_argv, "--seed", "7"]) == 1 diff --git a/tests/unit/graph/test_weka_hf_real_ingest.py b/tests/unit/graph/test_weka_hf_real_ingest.py new file mode 100644 index 0000000000..81ee969dac --- /dev/null +++ b/tests/unit/graph/test_weka_hf_real_ingest.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Real-dataset ingest of the published SemiAnalysis weka corpus (062126). + +Proves the HuggingFace ``org/name`` load path end-to-end on REAL data: a small +pinned slice of ``semianalysisai/cc-traces-weka-062126`` is streamed (so the +~1.8 GB single-file ``traces.jsonl`` is never fully materialized), parsed into a +segment-trie :class:`ParsedGraph` through the SAME shared core the local-file +path uses, and drained into a :class:`GraphSegmentUnifiedBackingStore` via the +same unified-store builders :class:`DatasetManager._configure_graph_workload` +runs (interned eager build + payload-streamed build). + +Marked ``slow`` (real-content synthesis loads a tokenizer and synthesizes every +turn) and ``network`` (resolves the corpus from the HuggingFace cache / hub). +The slice + revision are pinned so a re-run is reproducible; the test skips +cleanly when the dataset cannot be resolved (no cache and no network). +""" + +from __future__ import annotations + +import orjson +import pytest + +from aiperf.common.environment import Environment +from aiperf.dataset.graph.adapters.weka.trace import ( + WekaTraceAdapterError, + _looks_like_hf_dataset_id, + from_weka_trace, +) +from aiperf.dataset.graph.models import LlmNode, ParsedGraph, resolve_trace_graph +from aiperf.dataset.graph.segment_ir.store_builder import ( + TraceSegmentPayload, + build_unified_trie_store_interned, +) +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) +from aiperf.graph.worker_materialize import materialize_graph_request_unified + +# Published SemiAnalysis weka corpus dated 062126. Public, single ``train`` +# split exposed as one ~1.8 GB ``traces.jsonl``. +WEKA_HF_REPO_ID = "semianalysisai/cc-traces-weka-062126" + +# Pinned commit SHA of the corpus snapshot this test was authored against, so a +# streamed re-run reads the identical bytes even if the dataset's default +# branch advances. Matches the locally cached revision. +WEKA_HF_REVISION = "23f152f6f0f9399a85901b89a6458def0ef16729" + +# A handful of real traces is enough to exercise the multi-graph merge, the +# trie pool union, and real-content segment synthesis without pulling the whole +# corpus. Streaming + this cap keeps the ingest fast and cheap. +SLICE_ROWS = 3 + + +def _parse_real_slice_or_skip() -> ParsedGraph: + """Parse the pinned real slice; skip ONLY on availability errors. + + hub/datasets availability failures (offline mode, cache miss, repo/auth + errors, HTTP failures) all descend from ``OSError``; the weka loader wraps + the ``load_dataset`` failure in ``WekaTraceAdapterError`` with the original + as ``__cause__``. Anything else -- a real parse/trie crash on real data -- + must fail the test loudly, never skip. + """ + unavailable = f"weka corpus {WEKA_HF_REPO_ID!r} unavailable" + try: + return from_weka_trace(WEKA_HF_REPO_ID, content_root_seed=42) + except OSError as exc: + pytest.skip(f"{unavailable}: {exc!r}") + except WekaTraceAdapterError as exc: + if isinstance(exc.__cause__, OSError): + pytest.skip(f"{unavailable}: {exc!r}") + raise + + +def test_repo_id_routes_to_weka_hf_loader() -> None: + """The published 062126 repo id is recognized as a weka HF dataset id.""" + assert _looks_like_hf_dataset_id(WEKA_HF_REPO_ID) + + +@pytest.mark.slow +@pytest.mark.network +@pytest.mark.asyncio +async def test_real_062126_slice_ingests_to_unified_store( + tmp_path, monkeypatch +) -> None: + """Stream a real 062126 slice -> trie ParsedGraph -> unified segment store. + + Asserts the real traces parse through the segment-trie IR (non-None + ``segment_pool``), produce a populated unified store with valid dense + 0..n-1 node ordinals per trace, and worker-materialize to non-empty + prompts — the build-time proof that real weka data flows through the + graph-IR-v1 trie ingest path. + """ + # Pin the streamed slice + revision so the ingest is reproducible and cheap. + # Bound the streamed rows via the split slice (WEKA_HF_MAX_ROWS was removed). + monkeypatch.setattr(Environment.DATASET, "WEKA_HF_SPLIT", f"train[:{SLICE_ROWS}]") + monkeypatch.setattr(Environment.DATASET, "WEKA_HF_REVISION", WEKA_HF_REVISION) + + parsed = _parse_real_slice_or_skip() + + # Streaming slice yields exactly SLICE_ROWS traces, each its own graph, + # every one parsed through the trie builder (pool always present). + assert len(parsed.traces) == SLICE_ROWS + assert len(parsed.graphs) == SLICE_ROWS + assert parsed.segment_pool is not None, "weka parse must surface the trie pool" + assert parsed.segment_pool._by_id, "trie pool must carry real content segments" + + # Every trace's graph lowers to LLM nodes with the recorded model resolved + # into dispatch_overrides -- the real "endpoint resolved" proof. + for trace in parsed.traces: + graph = resolve_trace_graph(parsed, trace) + llm_nodes = [n for n in graph.nodes.values() if isinstance(n, LlmNode)] + assert llm_nodes, f"trace {trace.id} must lower to at least one LLM node" + for node in llm_nodes: + assert (node.dispatch_overrides or {}).get("model"), ( + f"trace {trace.id}: LLM node missing dispatch_overrides['model']" + ) + + store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id="real062126" + ) + catalog = await build_unified_trie_store_interned(parsed, store) + + # Every parsed trace is addressable in the catalog with valid ordinals. + assert set(catalog) == {t.id for t in parsed.traces} + total_nodes = 0 + for nodes in catalog.values(): + assert nodes, "expected at least one node per trace" + ordinals = sorted(nodes.values()) + assert ordinals == list(range(len(nodes))), "ordinals must be dense 0..n-1" + total_nodes += len(nodes) + assert total_nodes > 0 + + # The unified store materializes non-empty real-content prompts for the + # overwhelming majority of nodes (nodes without a prompt path write no + # manifest, so a strong floor is asserted rather than 100%). + non_empty = 0 + with GraphSegmentUnifiedClient(tmp_path, "real062126").open() as client: + for trace_id, nodes in catalog.items(): + for ordinal in nodes.values(): + payload = materialize_graph_request_unified( + client, trace_id, ordinal, "profiling" + ) + if payload is None: + continue + messages = payload["messages"] + if messages and messages[0].get("content"): + non_empty += 1 + assert non_empty >= total_nodes - SLICE_ROWS + assert non_empty > 0 + + +@pytest.mark.slow +@pytest.mark.network +@pytest.mark.asyncio +async def test_streaming_ingest_matches_eager_byte_for_byte( + tmp_path, monkeypatch +) -> None: + """Streaming ingest == eager ingest, byte-for-byte, on the same real slice. + + The streaming build plane (``build_unified_trie_store_from_payloads`` over + ``stream_weka_trace_segment_payloads``) keeps resident memory at ~one trace so + the full real corpus ingests without OOM. This asserts the streamed payloads + carry the trie shape (:class:`TraceSegmentPayload` with segments + + ``prompt_segment_ids`` envelopes) and that draining them produces the + IDENTICAL unified store as the eager whole-corpus + ``build_unified_trie_store_interned``: same catalog (per-trace ordinal maps) + and byte-identical store files. + """ + # Bound the streamed rows via the split slice (WEKA_HF_MAX_ROWS was removed). + monkeypatch.setattr(Environment.DATASET, "WEKA_HF_SPLIT", f"train[:{SLICE_ROWS}]") + monkeypatch.setattr(Environment.DATASET, "WEKA_HF_REVISION", WEKA_HF_REVISION) + + from aiperf.dataset.graph.adapters.weka.trace import ( + stream_weka_trace_segment_payloads, + ) + from aiperf.dataset.graph.segment_ir.store_builder import ( + build_unified_trie_store_from_payloads, + ) + + eager_parsed = _parse_real_slice_or_skip() + + assert eager_parsed.segment_pool is not None + eager_store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id="eager" + ) + eager_catalog = await build_unified_trie_store_interned(eager_parsed, eager_store) + + payloads = list( + stream_weka_trace_segment_payloads(WEKA_HF_REPO_ID, content_root_seed=42) + ) + assert len(payloads) == SLICE_ROWS + assert all(isinstance(p, TraceSegmentPayload) for p in payloads) + assert any(p.segments for p in payloads), "trie payloads must carry segments" + for payload in payloads: + assert payload.envelopes, f"trace {payload.trace_id} must carry envelopes" + for node_envelope in payload.envelopes: + env = orjson.loads(node_envelope.envelope_bytes) + assert "prompt_segment_ids" in env, "trie envelope, not legacy" + assert "messages_delta" not in env, "must NOT be a legacy delta envelope" + + stream_store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id="stream" + ) + stream_catalog = await build_unified_trie_store_from_payloads( + payloads, stream_store + ) + + assert stream_catalog == eager_catalog, "streaming catalog must equal eager" + + eager_dir = tmp_path / "aiperf_graph_segments_eager" + stream_dir = tmp_path / "aiperf_graph_segments_stream" + eager_files = sorted(p.name for p in eager_dir.iterdir()) + stream_files = sorted(p.name for p in stream_dir.iterdir()) + assert eager_files == stream_files and eager_files, ( + f"unified store file sets differ: {eager_files} vs {stream_files}" + ) + for name in eager_files: + assert (eager_dir / name).read_bytes() == (stream_dir / name).read_bytes(), ( + f"unified store file {name!r} differs between streaming and eager" + ) diff --git a/tests/unit/graph/test_weka_ingest.py b/tests/unit/graph/test_weka_ingest.py new file mode 100644 index 0000000000..fe6f42ba59 --- /dev/null +++ b/tests/unit/graph/test_weka_ingest.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from pathlib import Path + +import msgspec +import pytest + +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.models import GraphRecord, LlmNode, ParsedGraph +from aiperf.dataset.graph.segment_ir.store_builder import ( + build_unified_trie_store_from_payloads, + build_unified_trie_store_interned, + iter_trace_segment_payloads, +) +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, +) + +FIX = Path(__file__).parent / "fixtures" / "weka_min.json" + + +def test_from_weka_trace_yields_parsed_graph_real_content(): + # The segment-trie IR is the only weka path: ingest attaches a SegmentPool and + # every LlmNode carries a ``prompt_segment_ids`` path that materializes to real + # (non-placeholder) conversation text -- no legacy replay_outputs/subgraphs. + parsed = from_weka_trace(str(FIX)) + assert parsed.traces, "expected at least one trace" + assert parsed.segment_pool is not None, "trie ingest must attach a SegmentPool" + + pool = parsed.segment_pool + graph = parsed.graph + llm_nodes = [n for n in graph.nodes.values() if isinstance(n, LlmNode)] + assert llm_nodes, "expected at least one LlmNode" + + for node in llm_nodes: + trie = node.metadata.get("trie") + assert trie is not None, "trie node must carry trie metadata" + path = trie["prompt_segment_ids"] + assert path, "trie node must carry a non-empty prompt segment path" + + # The trie route carries NO inline prompt: content lives ONLY in the + # segment pool, reached via the prompt_segment_ids path. + assert node.prompt == [] + materialized = pool.materialize(path) + # Real content, not the empty/placeholder prompt: every segment has text. + for msg in materialized: + assert msg["role"] in {"system", "user", "assistant", "tool"} + assert isinstance(msg["content"], str) and msg["content"] + + # The node's recorded response is a real assistant pool entry chained + # onto the prompt tip (content-addressed; no metadata handle needed). + tip = path[-1] + assert any( + s.role == "assistant" and s.parent_id == tip for s in pool.by_id.values() + ) + + +def _with_sentinel_prompts(parsed: ParsedGraph) -> ParsedGraph: + """Copy ``parsed`` with a non-empty sentinel on every LlmNode.prompt, across + ``parsed.graph`` AND every ``parsed.graphs`` value.""" + sentinel: list = [{"role": "user", "content": "SENTINEL"}] + + def _stamp(graph: GraphRecord) -> GraphRecord: + nodes = { + nid: ( + msgspec.structs.replace(node, prompt=sentinel) + if isinstance(node, LlmNode) + else node + ) + for nid, node in graph.nodes.items() + } + return msgspec.structs.replace(graph, nodes=nodes) + + return msgspec.structs.replace( + parsed, + graph=_stamp(parsed.graph), + graphs={ref: _stamp(g) for ref, g in parsed.graphs.items()}, + ) + + +@pytest.mark.asyncio +async def test_weka_store_bytes_independent_of_inline_prompt(tmp_path: Path) -> None: + """Weka trie store bytes are a function of (segment pool, trie envelope) ONLY + -- never the inline ``LlmNode.prompt`` -- through BOTH the eager and streaming + drains. Mirrors the dynamo pin (``test_dynamo_streaming_store_parity``) on a + weka-shaped parse so the ``prompt=[]`` convention is pinned per format.""" + parsed = from_weka_trace(str(FIX)) + sentinel = _with_sentinel_prompts(parsed) + # Guard the guard: EVERY LlmNode in the sentinel copy must carry the non-empty + # sentinel prompt -- an ``all(... == sentinel)`` (not ``any``) so a partial + # ``_stamp`` failure cannot pass this vacuously (mirrors the dynamo twin in + # test_dynamo_streaming_store_parity.py). + sentinel_nodes = [ + node + for graph in (sentinel.graph, *sentinel.graphs.values()) + for node in graph.nodes.values() + if isinstance(node, LlmNode) + ] + assert sentinel_nodes and all( + node.prompt == [{"role": "user", "content": "SENTINEL"}] + for node in sentinel_nodes + ) + + async def _dirs(p: ParsedGraph, tag: str) -> tuple[Path, Path]: + interned = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id=f"{tag}-i" + ) + await build_unified_trie_store_interned(p, interned) + stream = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id=f"{tag}-s" + ) + await build_unified_trie_store_from_payloads( + iter_trace_segment_payloads(p), stream + ) + return ( + tmp_path / f"aiperf_graph_segments_{tag}-i", + tmp_path / f"aiperf_graph_segments_{tag}-s", + ) + + base_i, base_s = await _dirs(parsed, "base") + sent_i, sent_s = await _dirs(sentinel, "sent") + + for real_dir, sent_dir in ((base_i, sent_i), (base_s, sent_s)): + names = sorted(p.name for p in real_dir.iterdir()) + assert names == sorted(p.name for p in sent_dir.iterdir()) and names + for name in names: + assert (real_dir / name).read_bytes() == (sent_dir / name).read_bytes(), ( + f"weka store file {name!r} differs -- a drain read inline " + f"node.prompt instead of the segment pool + trie envelope" + ) diff --git a/tests/unit/graph/test_weka_isl_endpoint_osl_fidelity.py b/tests/unit/graph/test_weka_isl_endpoint_osl_fidelity.py new file mode 100644 index 0000000000..70e963ed88 --- /dev/null +++ b/tests/unit/graph/test_weka_isl_endpoint_osl_fidelity.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ISL recovery from the wire payload when ``Turn.texts`` is empty (Fix 1). + +The graph / weka dispatch path wraps the prompt in a ``raw_payload`` (chat +``messages``) with EMPTY ``Turn.texts``, so the parser's turns-walk yields no +tokenizable text and ``InputSequenceLengthMetric`` would be dropped -- ISL +silently absent. Agentx computes ISL by tokenizing the WIRE payload +(``extract_payload_inputs`` over ``payload_bytes``); the fix mirrors that as a +fallback when ``texts`` is empty. These locks are parser-level and independent +of the weka IR shape. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import orjson +import pytest + +from aiperf.common.models import RequestRecord, Text, Turn +from aiperf.common.models.record_models import ( + ParsedResponse, + RecordContext, + TextResponseData, +) +from aiperf.common.tokenizer import Tokenizer +from aiperf.records.inference_result_parser import InferenceResultParser + +# --------------------------------------------------------------------------- +# Fix 1 -- ISL recovered from the wire payload when Turn.texts is empty +# --------------------------------------------------------------------------- + + +def _raw_payload_record(messages: list[dict], model: str) -> RequestRecord: + """A record whose only turn is a raw_payload chat body (no texts). + + Mirrors the graph / weka dispatch shape: ``Turn(raw_payload=...)`` with empty + ``texts``, and ``request_info.payload_bytes`` = the exact wire JSON. + """ + payload = {"messages": messages, "max_completion_tokens": 25, "model": model} + return RequestRecord( + turns=[Turn(role="user", raw_payload=payload)], + request_info=RecordContext( + payload_bytes=orjson.dumps(payload), + credit_num=0, + credit_phase="profiling", + conversation_id="conv-0", + turn_index=0, + x_request_id="req-0", + x_correlation_id="t-1#0|deadbeef|t-1|parent_0|profiling", + ), + model_name=model, + ) + + +@pytest.fixture +def parser(benchmark_run) -> InferenceResultParser: + """A parser with the REAL chat endpoint (so ``extract_payload_inputs`` runs + for real over the wire payload) but a word-count fake tokenizer (so no HF + model download is needed). The endpoint's payload extraction is what the + Fix 1 fallback exercises; the tokenizer is incidental to the count. + """ + parser = InferenceResultParser(run=benchmark_run) + fake = MagicMock(spec=Tokenizer) + fake.encode.side_effect = lambda text: list(range(len(text.split()))) + parser.get_tokenizer = AsyncMock(return_value=fake) + return parser + + +@pytest.mark.asyncio +async def test_isl_from_wire_payload_when_texts_empty( + parser: InferenceResultParser, +) -> None: + """A raw_payload chat body (empty texts) still yields a tokenized ISL. + + The turns-walk finds no ``texts``; the fallback runs the endpoint's + ``extract_payload_inputs`` over the wire payload and tokenizes the recovered + prompt text, so ISL is the wire-prompt token count rather than ``None``. + """ + record = _raw_payload_record( + [{"role": "user", "content": "hello world this is the prompt"}], + parser.model_endpoint.primary_model_name, + ) + isl = await parser.compute_input_token_count(record) + assert isl is not None and isl > 0 + + +@pytest.mark.asyncio +async def test_isl_text_turn_path_unchanged(parser: InferenceResultParser) -> None: + """A turn carrying ``texts`` uses the turns-walk, not the wire fallback. + + Locks that Fix 1 is a strict fallback: when ``texts`` is present the existing + space-joined tokenization is used (the linear path is not regressed). + """ + record = RequestRecord( + turns=[Turn(role="user", texts=[Text(name="p", contents=["hello world"])])], + model_name=parser.model_endpoint.primary_model_name, + ) + isl = await parser.compute_input_token_count(record) + assert isl is not None and isl > 0 + + +@pytest.mark.asyncio +async def test_isl_none_when_no_texts_and_no_payload( + parser: InferenceResultParser, +) -> None: + """No texts and no wire payload -> ISL is None (unavailable), not a crash.""" + record = RequestRecord( + turns=[Turn(role="user")], + model_name=parser.model_endpoint.primary_model_name, + ) + isl = await parser.compute_input_token_count(record) + assert isl is None + + +# --------------------------------------------------------------------------- +# OSL -- output tokens still counted on the raw_payload (graph/weka) shape +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_osl_counted_from_response_text_on_raw_payload_record( + parser: InferenceResultParser, +) -> None: + """OSL fidelity: the raw_payload record shape yields a real OUTPUT count. + + The graph/weka dispatch shape (empty ``Turn.texts``) must not degrade the + client-side output tokenization -- OSL is counted from the response text + (5 words -> 5 tokens under the word-count fake) alongside the + wire-recovered ISL in the same ``TokenCounts``. + """ + record = _raw_payload_record( + [{"role": "user", "content": "hello world this is the prompt"}], + parser.model_endpoint.primary_model_name, + ) + responses = [ + ParsedResponse( + perf_ns=1, data=TextResponseData(text="five words of model output") + ), + ] + + counts = await parser._compute_client_side_token_counts(record, responses) + + assert counts.output == 5 + assert counts.input is not None and counts.input > 0 + + +@pytest.mark.asyncio +async def test_osl_concatenates_streamed_response_chunks( + parser: InferenceResultParser, +) -> None: + """Streamed chunks are DELTAS: joined with no separator before tokenizing.""" + record = _raw_payload_record( + [{"role": "user", "content": "hi"}], + parser.model_endpoint.primary_model_name, + ) + responses = [ + ParsedResponse(perf_ns=1, data=TextResponseData(text="hello wor")), + ParsedResponse(perf_ns=2, data=TextResponseData(text="ld once again")), + ] + + counts = await parser._compute_client_side_token_counts(record, responses) + + # "hello wor" + "ld once again" -> "hello world once again" -> 4 words. + assert counts.output == 4 diff --git a/tests/unit/graph/test_weka_memory_regressions.py b/tests/unit/graph/test_weka_memory_regressions.py new file mode 100644 index 0000000000..fdefc4a657 --- /dev/null +++ b/tests/unit/graph/test_weka_memory_regressions.py @@ -0,0 +1,240 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import sys +import types +from collections.abc import Iterator +from typing import Any + +import numpy as np + +from aiperf.common.environment import Environment + + +def test_hf_rows_are_streamed_lazily(monkeypatch) -> None: + from aiperf.dataset.graph.adapters.weka import trace as weka_trace + + calls: list[dict[str, Any]] = [] + + def load_dataset(*_args, **kwargs): # noqa: ANN001 + calls.append(dict(kwargs)) + return iter([{"id": "row-0"}, {"id": "row-1"}]) + + monkeypatch.setitem( + sys.modules, + "datasets", + types.SimpleNamespace(load_dataset=load_dataset), + ) + + rows = weka_trace._load_hf_rows("org/weka-corpus", split="train", revision="rev") + assert calls == [] + + assert next(rows) == {"id": "row-0"} + assert calls == [{"split": "train", "streaming": True, "revision": "rev"}] + + +def test_parse_items_streams_decoded_results_into_merge(monkeypatch) -> None: + from aiperf.dataset.graph.adapters.weka import trace_parallel as parallel + + events: list[str] = [] + + def fake_run_pool_streaming( + _worker_fn: Any, + _work_items: Any, + **_kwargs: Any, + ) -> Iterator[bytes]: + events.append("pool_start") + yield b"one" + events.append("after_one") + yield b"two" + + def fake_decode(blob: bytes) -> str: + decoded = blob.decode() + events.append(f"decode:{decoded}") + return decoded + + def fake_merge(per_row: Any) -> str: + assert not isinstance(per_row, list) + events.append(f"merge_arg:{type(per_row).__name__}") + iterator = iter(per_row) + assert next(iterator) == "one" + assert events == ["merge_arg:generator", "pool_start", "decode:one"] + assert next(iterator) == "two" + assert events == [ + "merge_arg:generator", + "pool_start", + "decode:one", + "after_one", + "decode:two", + ] + return "merged" + + monkeypatch.setattr(parallel, "_run_pool_streaming", fake_run_pool_streaming) + monkeypatch.setattr(parallel, "decode_parsed_graph_msgpack", fake_decode) + monkeypatch.setattr(parallel, "merge_parsed_graphs", fake_merge) + + assert ( + parallel.parse_items( + [parallel._WorkItem(source="src#0", path=None, row={"id": "row-0"})], + source_label="org/weka-corpus", + threshold=0, + workers=1, + parse_kwargs={}, + ) + == "merged" + ) + + +def test_prefetch_multiplier_field_default_ratchet() -> None: + """The Weka prefetch-window multiplier default must not regress below 16 + (window 256 at the auto 16 workers -- covers the rows remaining behind the + heaviest full-corpus trace so fast workers do not stall head-of-line). + + Assert the declared FIELD default on the settings model, NOT the runtime + ``Environment.DATASET`` value: the latter is env-overridable + (``AIPERF_DATASET_WEKA_GRAPH_PARALLEL_PREFETCH_MULTIPLIER``), so the + per-leg ABAB overrides used to measure this change would otherwise make the + ratchet lie.""" + field = type(Environment.DATASET).model_fields[ + "WEKA_GRAPH_PARALLEL_PREFETCH_MULTIPLIER" + ] + assert field.default >= 16 + + +def test_weka_worker_auto_matches_agentx(monkeypatch) -> None: + from aiperf.dataset.graph.adapters.weka import trace_parallel as parallel + + monkeypatch.setattr(parallel.os, "cpu_count", lambda: 32) + monkeypatch.setattr(Environment.DATASET, "WEKA_GRAPH_PARALLEL_WORKERS", 0) + monkeypatch.setattr(Environment.DATASET, "WEKA_GRAPH_PARALLEL_AUTO_MAX_WORKERS", 16) + + assert parallel._get_workers(item_count=393, override=None) == 16 + assert parallel._get_workers(item_count=4, override=None) == 4 + assert parallel._get_workers(item_count=393, override=2) == 2 + assert parallel._get_workers(item_count=393, override=0) == 16 + + monkeypatch.setattr(Environment.DATASET, "WEKA_GRAPH_PARALLEL_WORKERS", 3) + assert parallel._get_workers(item_count=393, override=None) == 3 + + +def test_numpy_shared_corpus_wraparound_matches_list_corpus() -> None: + from aiperf.dataset.graph.adapters.shared.content import ( + CorpusContentSynthesizer, + ) + + class _Generator: + def __init__(self, corpus) -> None: # noqa: ANN001 + self._tokenized_corpus = corpus + self._corpus_size = len(corpus) + self._cache: dict[int, list[int]] = {} + + class _HashRng: + def reseed_for_hash_id(self, _hash_id: int) -> None: + return + + def randrange(self, _upper: int) -> int: + return 8 + + def decode(corpus) -> list[int]: # noqa: ANN001 + synth = object.__new__(CorpusContentSynthesizer) + synth._pg = _Generator(corpus) + synth._hash_id_corpus_rng = _HashRng() + synth._block_size = 5 + return CorpusContentSynthesizer._decode_block_tokens(synth, [123]) + + expected = [8, 9, 0, 1, 2] + assert decode(list(range(10))) == expected + assert decode(np.asarray(list(range(10)), dtype=np.int32)) == expected + + +def _offset_synth(corpus, *, block_size: int, seed: int = 98765): + """A CorpusContentSynthesizer bound to a tiny stub corpus + a REAL + HashIdRandomGenerator (so offsets vary per hash id), mirroring + ``test_numpy_shared_corpus_wraparound_matches_list_corpus`` but exercising + the real reseed/randrange path both decoders share.""" + import types + + from aiperf.common.hash_id_random_generator import HashIdRandomGenerator + from aiperf.dataset.graph.adapters.shared.content import ( + CorpusContentSynthesizer, + ) + + synth = object.__new__(CorpusContentSynthesizer) + synth._pg = types.SimpleNamespace( + _tokenized_corpus=corpus, _corpus_size=len(corpus), _cache={} + ) + synth._hash_id_corpus_rng = HashIdRandomGenerator(seed, _internal=True) + synth._block_size = block_size + synth._offset_cache_corpus_size = None + return synth + + +def test_offset_cache_decode_matches_list_cache_across_backings() -> None: + """The offset-cached decode is a byte-identical, memory-lean twin of the + list-cache ``_decode_block_tokens`` on the MISS and the REPEAT (hit) path, for + both a list and an ``np.int32`` corpus backing, wraparound blocks included. + + Both decoders issue the same ``reseed_for_hash_id(h)`` + ``randrange`` pair on + a miss (a full reseed -> identical ``start``) and touch no RNG on a repeat, so + the token streams must match exactly; only the cache SHAPE differs + (``int`` offset vs decoded ``list``). The tiny corpus forces frequent + wraparound, so BOTH branches of the offset decode (the extend-from-slice + fast path and the verbatim two-step wraparound path) are pinned.""" + bs = 16 + base = list(range(40)) # tiny corpus forces frequent wraparound at bs=16 + # ~500 hash ids: positive u64-ish + negative virtual ids, deterministic. + hash_ids = [((i * 2654435761) % (2**63)) + 1 for i in range(300)] + hash_ids += [-(i + 1) for i in range(200)] + + for corpus in (base, np.asarray(base, dtype=np.int32)): + s_list = _offset_synth(corpus, block_size=bs) + s_off = _offset_synth(corpus, block_size=bs) + dict_cache: dict[int, list[int]] = {} + off_cache: dict[int, int] = {} + + first_list = s_list._decode_block_tokens( + hash_ids, block_size=bs, cache=dict_cache + ) + first_off = s_off._decode_block_tokens_offset_cached( + hash_ids, block_size=bs, offset_cache=off_cache + ) + assert first_off == first_list + + # Repeat decode: both caches hit; each must reproduce its first decode. + repeat_list = s_list._decode_block_tokens( + hash_ids, block_size=bs, cache=dict_cache + ) + repeat_off = s_off._decode_block_tokens_offset_cached( + hash_ids, block_size=bs, offset_cache=off_cache + ) + assert repeat_list == first_list + assert repeat_off == first_off == repeat_list + + # Every covered block is exactly block_size (wraparound blocks included). + assert len(first_off) == len(hash_ids) * bs + # Guard the guard: BOTH offset-decode branches were actually exercised + # (some offsets wrap past the corpus end, some do not). + assert any(v > len(base) - bs for v in off_cache.values()) + assert any(v <= len(base) - bs for v in off_cache.values()) + # Cache shape: plain int offsets keyed by the (int) hash id. + assert all(type(v) is int for v in off_cache.values()) + + +def test_offset_cache_fails_loud_on_corpus_size_change() -> None: + """A rebind that changes ``_corpus_size`` under a populated offset cache fails + loud: the cached offsets would otherwise silently reproduce wrong bytes (the + corpus-immutability contract on ``_decode_block_tokens_offset_cached``).""" + import pytest + + synth = _offset_synth(list(range(40)), block_size=16, seed=123) + off: dict[int, int] = {} + synth._decode_block_tokens_offset_cached([1, 2, 3], block_size=16, offset_cache=off) + + synth._pg._tokenized_corpus = list(range(80)) + synth._pg._corpus_size = 80 + with pytest.raises(RuntimeError, match="corpus size changed"): + synth._decode_block_tokens_offset_cached( + [1, 2, 3], block_size=16, offset_cache=off + ) diff --git a/tests/unit/graph/test_weka_parallel_hardening.py b/tests/unit/graph/test_weka_parallel_hardening.py new file mode 100644 index 0000000000..19d8b9ae2c --- /dev/null +++ b/tests/unit/graph/test_weka_parallel_hardening.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weka parallel-parse hardening. + +* Bounded waits: ``_bounded_ordered_pool_map`` waits on each pool result with a bounded + ``get(timeout=...)`` -- a raw ``multiprocessing.Pool`` never completes the + in-flight task of a killed worker (OOM / SIGKILL), so an unbounded ``get()`` + presents as a silent indefinite CLI hang. Expiry must raise a clear + ``RuntimeError`` naming that cause; ordering is preserved. + +* Shared-memory attach: a pool worker attaches the parent-built shared-memory corpus WITHOUT + first paying a private corpus build -- ``get_or_build_synthesizer(..., + shared_corpus=...)`` constructs the generator directly on the supplied + array, and decoded content stays byte-identical to a self-built synthesizer. +""" + +from __future__ import annotations + +import multiprocessing +from multiprocessing import shared_memory +from typing import Any + +import numpy as np +import pytest + +from aiperf.common.environment import Environment +from aiperf.dataset.graph.adapters.shared.content import ( + _WORKER_SYNTH_CACHE, + CorpusContentSynthesizer, + get_or_build_synthesizer, +) +from aiperf.dataset.graph.adapters.weka import trace_parallel as parallel + + +@pytest.fixture(autouse=True) +def _fresh_synth_cache(): + """Isolate the process-level synthesizer cache from other tests.""" + CorpusContentSynthesizer.reset_worker_cache() + yield + CorpusContentSynthesizer.reset_worker_cache() + + +# --- W5: bounded pool result wait ------------------------------------------- + + +class _ReadyResult: + """AsyncResult stub that completes immediately.""" + + def __init__(self, value: bytes) -> None: + self._value = value + + def get(self, timeout: float | None = None) -> bytes: + assert timeout is not None, "pool result wait must be bounded" + return self._value + + +class _DeadWorkerResult: + """AsyncResult stub for a task whose worker was killed: never completes.""" + + def get(self, timeout: float | None = None) -> bytes: + assert timeout is not None, "pool result wait must be bounded" + raise multiprocessing.TimeoutError + + +class _FakePool: + """apply_async returns the pre-scripted result for each submitted item.""" + + def __init__(self, results: list[Any]) -> None: + self._results = list(results) + + def apply_async(self, fn, args): # noqa: ANN001, ARG002 + return self._results.pop(0) + + +def test_bounded_pool_map_preserves_input_order() -> None: + results = [_ReadyResult(b"a"), _ReadyResult(b"b"), _ReadyResult(b"c")] + out = list( + parallel._bounded_ordered_pool_map( + _FakePool(results), lambda item: item, ["a", "b", "c"], window_size=2 + ) + ) + assert out == [b"a", b"b", b"c"] + + +def test_bounded_pool_map_raises_runtime_error_on_dead_worker(monkeypatch) -> None: + """A killed worker's never-completing result raises instead of hanging.""" + monkeypatch.setattr( + Environment.DATASET, "WEKA_GRAPH_PARALLEL_ITEM_TIMEOUT_SECONDS", 0.01 + ) + results = [_ReadyResult(b"a"), _DeadWorkerResult()] + stream = parallel._bounded_ordered_pool_map( + _FakePool(results), lambda item: item, ["a", "b"], window_size=2 + ) + assert next(stream) == b"a" + with pytest.raises(RuntimeError, match="killed"): + next(stream) + + +# --- W6: shared corpus without private rebuild ------------------------------ + + +def test_shared_corpus_synthesizer_skips_private_build_and_matches_bytes( + monkeypatch, + fake_tokenizer: None, # noqa: ARG001 +) -> None: + """Constructing on a shared corpus pays NO private build; bytes identical.""" + from aiperf.dataset.generator.coding_content import CodingContentGenerator + + builds: list[int] = [] + original_build = CodingContentGenerator._build_tool_pool + + def counting_build(self) -> None: # noqa: ANN001 + builds.append(1) + return original_build(self) + + monkeypatch.setattr(CodingContentGenerator, "_build_tool_pool", counting_build) + + # Parent role: pays the corpus build exactly once. + parent = CorpusContentSynthesizer("tok-w6", prompt_corpus="coding", root_seed=0) + assert builds == [1] + + corpus = np.asarray(parent.corpus_tokens(), dtype=np.int32) + + # Worker role: constructed ON the shared array -- no second build. + worker = CorpusContentSynthesizer( + "tok-w6", prompt_corpus="coding", root_seed=0, shared_corpus=corpus + ) + assert builds == [1], "shared-corpus construction must not rebuild the corpus" + assert worker._pg._corpus_size == len(corpus) + + # Determinism is sacred: decode + partial tails byte-identical. + parent_cache: dict[int, list[int]] = {} + worker_cache: dict[int, list[int]] = {} + parent_tokens = parent._decode_block_tokens( + [1, 2, 3], cache=parent_cache, trace_id="t" + ) + worker_tokens = worker._decode_block_tokens( + [1, 2, 3], cache=worker_cache, trace_id="t" + ) + assert [int(t) for t in worker_tokens] == [int(t) for t in parent_tokens] + assert [int(t) for t in worker._sample_partial_tail_tokens(7, "n0:response")] == [ + int(t) for t in parent._sample_partial_tail_tokens(7, "n0:response") + ] + + +def test_init_worker_attaches_shared_corpus_without_rebuild( + monkeypatch, + fake_tokenizer: None, # noqa: ARG001 +) -> None: + """``_init_worker`` warms the cache on the shm corpus with zero build calls.""" + from aiperf.dataset.generator.coding_content import CodingContentGenerator + + # Keep the test process's signal handlers and HF env untouched. + monkeypatch.setattr(parallel, "_install_hard_exit_on_sigterm", lambda: None) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + + def forbidden_build(self) -> None: # noqa: ANN001 + raise AssertionError("worker init must not build a private corpus") + + monkeypatch.setattr(CodingContentGenerator, "_build_tool_pool", forbidden_build) + + corpus = np.arange(512, dtype=np.int32) + shm = shared_memory.SharedMemory(create=True, size=corpus.nbytes) + try: + np.ndarray(corpus.shape, dtype=np.int32, buffer=shm.buf)[:] = corpus + args = parallel._WorkerInitArgs( + tokenizer_name="tok-init", + prompt_corpus="coding", + root_seed=3, + shm_name=shm.name, + corpus_len=512, + ) + + parallel._init_worker(args) + + # _init_worker is fail-soft, so an empty cache means the private-build + # path fired (and was swallowed); the cache MUST be warm on the shm array. + synth = _WORKER_SYNTH_CACHE[("tok-init", "coding", 3)] + assert synth._pg._corpus_size == 512 + assert [int(t) for t in synth._pg._tokenized_corpus[:4]] == [0, 1, 2, 3] + finally: + parallel._worker_corpus_shm = None + shm.close() + shm.unlink() + + +def test_get_or_build_rebinds_cached_synthesizer_to_shared_corpus( + fake_tokenizer: None, # noqa: ARG001 +) -> None: + """A cache HIT with a shared corpus rebinds instead of keeping the private one.""" + first = get_or_build_synthesizer("tok-hit", prompt_corpus="coding", root_seed=0) + replacement = np.arange(64, dtype=np.int32) + second = get_or_build_synthesizer( + "tok-hit", prompt_corpus="coding", root_seed=0, shared_corpus=replacement + ) + assert second is first + assert second._pg._corpus_size == 64 diff --git a/tests/unit/graph/test_weka_parse_kwargs_unification.py b/tests/unit/graph/test_weka_parse_kwargs_unification.py new file mode 100644 index 0000000000..22c31408b1 --- /dev/null +++ b/tests/unit/graph/test_weka_parse_kwargs_unification.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import inspect +from pathlib import Path +from typing import Any + +from aiperf.common import random_generator as rng +from aiperf.dataset.graph.adapters.weka import trace as weka_trace + +FIX_MIN = Path(__file__).parent / "fixtures" / "weka_min.json" + + +def test_resolve_parse_kwargs_keys_match_parse_trace_dict_signature() -> None: + # Drift guard: the ONE kwargs dict must cover exactly the keyword set of + # the shared per-trace core (minus the per-item ``source`` label). + kwargs = weka_trace._resolve_parse_kwargs( + tag="t", + idle_gap_cap_seconds=weka_trace._USE_DEFAULT, + content_root_seed=None, + content_tokenizer=None, + prompt_corpus=None, + max_osl=None, + ) + sig = inspect.signature(weka_trace._parse_trace_dict) + expected = { + name + for name, p in sig.parameters.items() + if p.kind is inspect.Parameter.KEYWORD_ONLY and name != "source" + } + assert set(kwargs) == expected + assert "delay_cap_seconds" not in kwargs + + +def test_resolve_parse_kwargs_resolves_sentinel_and_seed() -> None: + rng.reset() + rng.init(777) + kwargs = weka_trace._resolve_parse_kwargs( + tag="t", + idle_gap_cap_seconds=weka_trace._USE_DEFAULT, + content_root_seed=None, + content_tokenizer=None, + prompt_corpus=None, + max_osl=None, + ) + assert kwargs["idle_gap_cap_seconds"] == weka_trace._DEFAULT_IDLE_GAP_CAP_SECONDS + assert kwargs["content_root_seed"] == 777 + + +def test_from_weka_trace_threads_resolved_kwargs_to_core(monkeypatch) -> None: + rng.reset() + rng.init(777) + captured: list[dict[str, Any]] = [] + real = weka_trace._parse_trace_dict + + def spy(raw: dict[str, Any], *, source: str, **kwargs: Any): + captured.append(dict(kwargs)) + return real(raw, source=source, **kwargs) + + monkeypatch.setattr(weka_trace, "_parse_trace_dict", spy) + weka_trace.from_weka_trace(str(FIX_MIN)) + assert len(captured) == 1 + assert captured[0]["content_root_seed"] == 777 + assert ( + captured[0]["idle_gap_cap_seconds"] == weka_trace._DEFAULT_IDLE_GAP_CAP_SECONDS + ) diff --git a/tests/unit/graph/test_weka_segment_pool_merge_codec.py b/tests/unit/graph/test_weka_segment_pool_merge_codec.py new file mode 100644 index 0000000000..c6c882cfc6 --- /dev/null +++ b/tests/unit/graph/test_weka_segment_pool_merge_codec.py @@ -0,0 +1,139 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""segment_pool survives the directory merge and the cross-process msgpack codec. + +Regression guards for the directory/parallel weka trie-IR data-loss bug: the +merge dropped ``ParsedGraph.segment_pool`` (no ``segment_pool=`` arg) and the +codec decoded the (then ``Any``-typed) field to a bare ``dict`` instead of a +``SegmentPool``. Either drop knocks the trie graph off the trie ordinal scheme +and yields zero traces. +""" + +from __future__ import annotations + +import pytest + +from aiperf.dataset.graph.codecs import ( + decode_parsed_graph_msgpack, + encode_parsed_graph_msgpack, +) +from aiperf.dataset.graph.merge import ( + GraphMergeError, + merge_parsed_graphs, +) +from aiperf.dataset.graph.models import ParsedGraph, TraceRecord +from aiperf.dataset.graph.segment_ir.pool import Segment, SegmentPool + + +def _pool_with(entries: list[tuple[str, str, list[int], str | None]]) -> SegmentPool: + pool = SegmentPool() + for role, content, tokens, parent_id in entries: + pool.add(role=role, content=content, tokens=tokens, parent_id=parent_id) + return pool + + +def test_merge_unions_per_file_segment_pools() -> None: + # A shared root segment (same role/content/tokens/parent -> same content id) + # plus a file-distinct leaf in each file. + shared = ("system", "sys", [1, 2], None) + pool_a = _pool_with([shared, ("user", "alpha", [3], "ignored")]) + pool_b = _pool_with([shared, ("user", "beta", [4], "ignored")]) + + pg_a = ParsedGraph( + traces=[TraceRecord(id="t-a")], + segment_pool=pool_a, + ) + pg_b = ParsedGraph( + traces=[TraceRecord(id="t-b")], + segment_pool=pool_b, + ) + + merged = merge_parsed_graphs([pg_a, pg_b]) + + assert isinstance(merged.segment_pool, SegmentPool) + expected_ids = set(pool_a._by_id) | set(pool_b._by_id) + assert set(merged.segment_pool._by_id) == expected_ids + # Shared id present exactly once; union is strictly larger than either input. + assert len(merged.segment_pool._by_id) == len(expected_ids) + assert len(merged.segment_pool._by_id) > len(pool_a._by_id) + + # Every file's path still materializes against the merged pool. + for pool in (pool_a, pool_b): + for sid, seg in pool._by_id.items(): + assert merged.segment_pool.materialize([sid]) == [ + {"role": seg.role, "content": seg.content} + ] + + +def test_merge_divergent_content_same_id_raises() -> None: + # Ids are content-addressed (blake2b over parent_id/role/tokens), so the SAME + # id across pools MUST carry identical content. Two pools mapping one id to + # divergent segments can only mean a content-addressing / hash break, so the + # merge must fail loud rather than silently keep whichever entry wins. + seg_a = Segment(id="collide", role="user", content="alpha", parent_id=None) + seg_b = Segment(id="collide", role="user", content="beta", parent_id=None) + pool_a = SegmentPool(_by_id={"collide": seg_a}) + pool_b = SegmentPool(_by_id={"collide": seg_b}) + + pg_a = ParsedGraph(traces=[TraceRecord(id="t-a")], segment_pool=pool_a) + pg_b = ParsedGraph(traces=[TraceRecord(id="t-b")], segment_pool=pool_b) + + with pytest.raises(GraphMergeError): + merge_parsed_graphs([pg_a, pg_b]) + + +def test_merge_identical_content_same_id_dedups() -> None: + # The normal case: two pools share an id whose content is identical (Segment + # is a frozen dataclass, so == is a value comparison). The union dedups to a + # single entry and never raises. + expected = Segment(id="shared", role="user", content="alpha", parent_id=None) + pool_a = SegmentPool( + _by_id={ + "shared": Segment(id="shared", role="user", content="alpha", parent_id=None) + } + ) + pool_b = SegmentPool( + _by_id={ + "shared": Segment(id="shared", role="user", content="alpha", parent_id=None) + } + ) + + pg_a = ParsedGraph(traces=[TraceRecord(id="t-a")], segment_pool=pool_a) + pg_b = ParsedGraph(traces=[TraceRecord(id="t-b")], segment_pool=pool_b) + + merged = merge_parsed_graphs([pg_a, pg_b]) + + assert isinstance(merged.segment_pool, SegmentPool) + assert set(merged.segment_pool._by_id) == {"shared"} + assert merged.segment_pool._by_id["shared"] == expected + + +def test_merge_without_pool_leaves_segment_pool_none() -> None: + pg = ParsedGraph(traces=[TraceRecord(id="t-a")]) + merged = merge_parsed_graphs([pg]) + assert merged.segment_pool is None + + +def test_codec_round_trips_real_segment_pool() -> None: + pool = _pool_with( + [ + ("system", "sys", [1, 2], None), + ("user", "hello", [3, 4, 5], "x"), + ] + ) + pg = ParsedGraph(segment_pool=pool) + + decoded = decode_parsed_graph_msgpack(encode_parsed_graph_msgpack(pg)) + + assert type(decoded.segment_pool).__name__ == "SegmentPool" + assert set(decoded.segment_pool._by_id) == set(pool._by_id) + for sid, seg in pool._by_id.items(): + assert decoded.segment_pool.materialize([sid]) == [ + {"role": seg.role, "content": seg.content} + ] + + +def test_codec_round_trips_none_pool_unchanged() -> None: + pg = ParsedGraph() + decoded = decode_parsed_graph_msgpack(encode_parsed_graph_msgpack(pg)) + assert decoded.segment_pool is None diff --git a/tests/unit/graph/test_weka_segments.py b/tests/unit/graph/test_weka_segments.py new file mode 100644 index 0000000000..6ce74348a9 --- /dev/null +++ b/tests/unit/graph/test_weka_segments.py @@ -0,0 +1,36 @@ +from aiperf.dataset.graph.segment_ir.pool import SegmentPool + + +def test_rolling_id_is_prefix_dependent_and_deterministic(): + pool = SegmentPool() + a = pool.add(role="system", content="sys", tokens=[1, 2], parent_id=None) + b1 = pool.add(role="user", content="u", tokens=[3], parent_id=a) + # same content+prefix -> same id (dedup) + b2 = pool.add(role="user", content="u", tokens=[3], parent_id=a) + assert b1 == b2 + # same content, DIFFERENT prefix -> different id (rolling) + c = pool.add(role="system", content="other", tokens=[9], parent_id=None) + b3 = pool.add(role="user", content="u", tokens=[3], parent_id=c) + assert b3 != b1 + + +def test_materialize_path_is_exact_messages_in_order(): + pool = SegmentPool() + a = pool.add(role="system", content="sys", tokens=[1], parent_id=None) + b = pool.add(role="user", content="hi", tokens=[2], parent_id=a) + d = pool.add(role="assistant", content="yo", tokens=[3], parent_id=b) + assert pool.materialize([a, b, d]) == [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "yo"}, + ] + + +def test_truncation_branches_at_shared_prefix(): + pool = SegmentPool() + s1 = pool.add(role="system", content="1", tokens=[1], parent_id=None) + s2 = pool.add(role="user", content="2", tokens=[2], parent_id=s1) + s3 = pool.add(role="assistant", content="3", tokens=[3], parent_id=s2) + # truncation reuses [s1,s2], drops s3, adds s5 off s2 + s5 = pool.add(role="user", content="5", tokens=[5], parent_id=s2) + assert s5 != s3 and pool.get(s5).parent_id == s2 diff --git a/tests/unit/graph/test_weka_trace_models_finite.py b/tests/unit/graph/test_weka_trace_models_finite.py new file mode 100644 index 0000000000..43e63bc464 --- /dev/null +++ b/tests/unit/graph/test_weka_trace_models_finite.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weka trace-model NaN/Inf discipline. + +A bare ``ge`` bound on ``api_time`` / ``think_time`` / +``ttft`` would accept ``+inf`` (``inf >= 0`` satisfies it). ``inf`` +ingress is real via HuggingFace ``datasets`` rows (Arrow floats bypass orjson's +``Infinity`` rejection) and flows into ``raw_end = t + api_time`` edge-delay +math, stamping non-finite microsecond delays into ``StaticEdge`` floats that +cross the msgpack boundary. Every duration field must reject non-finite values +at parse time. +""" + +from __future__ import annotations + +import pytest +from pydantic import ValidationError +from pytest import param + +from aiperf.dataset.graph.adapters.weka.trace_models import ( + WekaNormalRequest, + WekaStreamingRequest, +) + + +def _request(type_: str, **overrides: object) -> dict: + base: dict = { + "t": 0.0, + "type": type_, + "model": "m", + "in": 64, + "out": 8, + "hash_ids": [1], + } + base.update(overrides) + return base + + +@pytest.mark.parametrize( + "model_cls,type_,field", + [ + param(WekaNormalRequest, "n", "api_time", id="n-api_time"), + param(WekaNormalRequest, "n", "think_time", id="n-think_time"), + param(WekaStreamingRequest, "s", "api_time", id="s-api_time"), + param(WekaStreamingRequest, "s", "think_time", id="s-think_time"), + param(WekaStreamingRequest, "s", "ttft", id="s-ttft"), + ], +) # fmt: skip +@pytest.mark.parametrize( + "bad", + [ + param(float("inf"), id="inf"), + param(float("nan"), id="nan"), + ], +) # fmt: skip +def test_duration_fields_reject_non_finite( + model_cls: type, type_: str, field: str, bad: float +) -> None: + with pytest.raises(ValidationError): + model_cls.model_validate(_request(type_, **{field: bad})) + + +@pytest.mark.parametrize( + "model_cls,type_", + [ + param(WekaNormalRequest, "n", id="normal"), + param(WekaStreamingRequest, "s", id="streaming"), + ], +) # fmt: skip +def test_finite_duration_fields_still_accepted(model_cls: type, type_: str) -> None: + req = model_cls.model_validate(_request(type_, api_time=0.5, think_time=1.5)) + assert req.api_time == 0.5 + assert req.think_time == 1.5 diff --git a/tests/unit/graph/test_weka_trie_build.py b/tests/unit/graph/test_weka_trie_build.py new file mode 100644 index 0000000000..cdc4d7f083 --- /dev/null +++ b/tests/unit/graph/test_weka_trie_build.py @@ -0,0 +1,482 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dependency-only ``ParsedGraph`` construction from a ``WekaTrace`` (Task 3). + +:func:`build_trie_graph` walks a Weka trace (top-level + nested subagent +requests) in recorded time order and emits a trivial ``ParsedGraph``: one +``LlmNode`` per recorded ``n``/``s`` request and plain ``StaticEdge`` +"waits-for" dependency edges. NO reducers / channels / subgraphs / spawn / +await nodes; NO chain-detection / ``::fa``-``::aux`` classification. + +The edge rule (the heart) is interval order plus a start-anchor carve-out: + * ``A -> R`` iff A FINISHED before R STARTED (raw ``A.t + A.api_time <= R.t``) + AND ``rank(A) < rank(R)``, after async exclusion and a frontier + (transitive-reduction) filter. The latest-ending frontier predecessor + carries the warped end-to-start ``delay_after_predecessor_us``; every other + frontier predecessor is an AND-fan-in wait (delay 0). + * A request with NO finished-before cause roots at ``StaticEdge(START, R)`` + with ``min_start_delay_us = R.t * 1e6``. + * Start-anchor carve-out: when R's recorded causal parent (spawner / chain-prev) + was still IN FLIGHT at R's start, R's incoming edges collapse to ONE + start-anchored ``StaticEdge(parent -> R, delay_after_predecessor_start_us = + warped start-to-start gap)`` -- 0 for coincident-start siblings, which + therefore dispatch together (see + ``test_coincident_fanout_second_branch_start_anchors_to_chain_prev``). + * Concurrency is otherwise emergent: two requests sharing a cause with no edge + between them stay edge-free; the only inter-sibling edge the builder adds is + the start-anchor above. + +These tests drive the builder with the same deterministic stub callbacks the +Task-2 segment-emission test uses, so no tokenizer / corpus build is needed; +they assert dependency STRUCTURE (edges, fan-in, fan-out), not content bytes. +""" + +from __future__ import annotations + +import pytest + +from aiperf.dataset.graph.adapters.weka.trace import EmptyWekaTraceError +from aiperf.dataset.graph.adapters.weka.trace_models import WekaTrace +from aiperf.dataset.graph.adapters.weka.trie_build import ( + ReconCallbacks, + build_trie_graph, +) +from aiperf.dataset.graph.models import ( + ChannelRequirement, + LlmNode, + ParsedGraph, + StaticEdge, +) +from aiperf.dataset.graph.segment_ir.pool import SegmentPool + +BLOCK_SIZE = 64 + + +def _stub_decode_block_tokens(hash_ids: list[int]) -> list[int]: + """Return ``BLOCK_SIZE`` distinct token IDs per hash id (agentx stub parity).""" + out: list[int] = [] + for h in hash_ids: + out.extend(range(h * 100, h * 100 + BLOCK_SIZE)) + return out + + +def _stub_partial_tail_tokens(n_tokens: int, seed: str) -> list[int]: + """Return ``n_tokens`` deterministic token IDs keyed on ``seed`` (stub).""" + base = sum(ord(c) for c in seed) * 1000 + return list(range(base, base + n_tokens)) + + +def _stub_decode_tokens_to_text(tokens: list[int]) -> str: + """Decode a token list to text by joining IDs (collision-free stub).""" + return "|".join(str(t) for t in tokens) + + +_STUB_CALLBACKS = ReconCallbacks( + decode_block_tokens=_stub_decode_block_tokens, + sample_partial_tail_tokens=_stub_partial_tail_tokens, + decode_tokens_to_text=_stub_decode_tokens_to_text, +) + + +def _build(trace: WekaTrace) -> tuple[ParsedGraph, SegmentPool]: + """Run the builder with the deterministic stub reconstructor callbacks.""" + return build_trie_graph(trace, callbacks=_STUB_CALLBACKS) + + +def _llm_nodes(graph: ParsedGraph) -> dict[str, LlmNode]: + """All ``LlmNode``s in the single top-level graph keyed by node id.""" + return {nid: n for nid, n in graph.graph.nodes.items() if isinstance(n, LlmNode)} + + +def _edges(graph: ParsedGraph) -> list[StaticEdge]: + """All ``StaticEdge``s in the single top-level graph.""" + return [e for e in graph.graph.edges if isinstance(e, StaticEdge)] + + +def _edge(edges: list[StaticEdge], source: str, target: str) -> StaticEdge | None: + """Find the (first) edge ``source -> target`` if present.""" + for e in edges: + if e.source == source and e.target == target: + return e + return None + + +def _n(t: float, hashes: list[int], in_len: int, out: int, api_time: float, + model: str = "M", type_: str = "n") -> dict: # fmt: skip + """One normal/streaming request dict in the Weka wire shape.""" + return { + "t": t, + "type": type_, + "model": model, + "in": in_len, + "out": out, + "hash_ids": hashes, + "api_time": api_time, + } + + +def _subagent(t: float, agent_id: str, requests: list[dict], status: str, + duration_ms: int | None = None) -> dict: # fmt: skip + """One subagent marker dict with nested inner requests.""" + return { + "t": t, + "type": "subagent", + "agent_id": agent_id, + "subagent_type": "Explore", + "status": status, + "duration_ms": duration_ms, + "requests": requests, + "models": ["M"], + } + + +def _response_segment(pool, node): + """The node's assistant response segment: the assistant pool entry chained + onto the node's prompt tip (content-addressed, so a successor's identical + history message dedups onto the same entry).""" + tip = node.metadata["trie"]["prompt_segment_ids"][-1] + return next( + s for s in pool.by_id.values() if s.role == "assistant" and s.parent_id == tip + ) + + +def _trace(requests: list[dict]) -> WekaTrace: + """Schema-validate a minimal local-scope Weka trace.""" + return WekaTrace.model_validate( + { + "id": "trace_0", + "models": ["M"], + "block_size": BLOCK_SIZE, + "hash_id_scope": "local", + "requests": requests, + } + ) + + +def test_sequential_continuation_single_predecessor() -> None: + """Two turns where turn1 starts after turn0 finishes -> edge turn0->turn1. + + turn0: t=0, api_time=2 -> completes at 2.0. turn1: t=3 (>= 2.0) shares + turn0's hash prefix (content-parent). The completed-before test passes, so + a single ``StaticEdge`` turn0->turn1 is emitted with + ``delay_after_predecessor_us = (3 - 2) * 1e6``. turn0 itself roots at START. + """ + trace = _trace( + [ + _n(t=0.0, hashes=[1, 2], in_len=128, out=64, api_time=2.0), + _n(t=3.0, hashes=[1, 2, 3, 4], in_len=256, out=64, api_time=1.0), + ] + ) + graph, pool = _build(trace) + + nodes = _llm_nodes(graph) + assert len(nodes) == 2 + ids = list(nodes.keys()) + n0, n1 = ids[0], ids[1] + + edges = _edges(graph) + # turn0 roots at START; turn1 waits for turn0. + assert _edge(edges, "START", n0) is not None + e = _edge(edges, n0, n1) + assert e is not None + assert e.delay_after_predecessor_us == (3.0 - 2.0) * 1e6 + # No spurious edge back from turn1 or START->turn1. + assert _edge(edges, "START", n1) is None + assert _edge(edges, n1, n0) is None + + # Dispatch overrides + arrival offset on the successor. + assert nodes[n1].arrival_offset_us == int(3.0 * 1e6) + assert nodes[n1].max_tokens == 64 + assert nodes[n1].model == "M" + # trie metadata carries the prompt path; the response segment is the + # assistant pool entry chained onto the prompt tip. + trie = nodes[n1].metadata["trie"] + assert set(trie) == {"prompt_segment_ids"} + assert trie["prompt_segment_ids"] + assert _response_segment(pool, nodes[n1]).role == "assistant" + + +def test_coincident_fanout_second_branch_start_anchors_to_chain_prev() -> None: + """B and C both branch off A at the SAME instant -> C start-anchors to B. + + A: t=0, api_time=1 -> done at 1.0. B (t=2) and C (t=2) both extend A's hash + prefix and run concurrently. B completed-before-waits on A (end-anchored + ``A -> B``). C's positional chain-prev is B, and B is still in flight at C's + coincident start (2 <= 2 < 7), so the start-anchor post-pass collapses C's + incoming edges to ONE start-anchored ``B -> C`` edge with a ZERO start-to-start + delay. C therefore dispatches WITH B (same warped instant t=2) -- the recorded + concurrency is preserved, just expressed as a start-anchored edge instead of a + second end-anchored fan-out from A. No backward ``C -> B`` edge; neither + sibling re-roots at START. + """ + trace = _trace( + [ + _n(t=0.0, hashes=[1, 2], in_len=128, out=64, api_time=1.0), + _n(t=2.0, hashes=[1, 2, 3], in_len=192, out=32, api_time=5.0), + _n(t=2.0, hashes=[1, 2, 4], in_len=192, out=32, api_time=5.0), + ] + ) + graph, _ = _build(trace) + + nodes = list(_llm_nodes(graph).keys()) + a, b, c = nodes[0], nodes[1], nodes[2] + + edges = _edges(graph) + assert _edge(edges, "START", a) is not None + # B end-anchored on its completed-before cause A (A finished at 1.0 < 2.0). + ab = _edge(edges, a, b) + assert ab is not None and ab.delay_after_predecessor_start_us is None + # C's incoming collapses to a single start-anchored edge from its chain-prev B. + assert _edge(edges, a, c) is None + c_incoming = [e for e in edges if e.target == c] + assert len(c_incoming) == 1 + (bc,) = c_incoming + assert bc.source == b + assert bc.delay_after_predecessor_start_us == 0.0 + # No backward inter-sibling edge; neither sibling re-roots at START. + assert _edge(edges, c, b) is None + assert _edge(edges, "START", b) is None + assert _edge(edges, "START", c) is None + + +def test_blocking_join_and_fan_in() -> None: + """Parent spawns subagents s1,s2; resume turn waits for BOTH last-turns. + + The parent's first turn (p0) precedes two completed subagent markers. The + subagents' fresh-context first turns get a structural-spawner edge from p0. + The parent's resume turn (p1) starts after both subagents finish, so it + AND-fans-in: an edge from s1's last turn AND from s2's last turn. + """ + s1 = _subagent( + t=1.0, + agent_id="agent_1", + status="completed", + duration_ms=2000, + requests=[_n(t=1.0, hashes=[50, 51], in_len=128, out=16, api_time=1.5)], + ) + s2 = _subagent( + t=1.0, + agent_id="agent_2", + status="completed", + duration_ms=2000, + requests=[_n(t=1.2, hashes=[60, 61], in_len=128, out=16, api_time=1.5)], + ) + trace = _trace( + [ + _n(t=0.0, hashes=[1, 2], in_len=128, out=64, api_time=0.5), + s1, + s2, + # Resume turn extends the parent prefix and starts after both + # subagents complete (s1 done ~2.5, s2 done ~2.7). + _n(t=4.0, hashes=[1, 2, 3], in_len=192, out=32, api_time=1.0), + ] + ) + graph, _ = _build(trace) + + nodes = _llm_nodes(graph) + # 4 LlmNodes: p0, s1-inner, s2-inner, p1 (resume). + assert len(nodes) == 4 + + edges = _edges(graph) + # Identify nodes by their hash prefix via the trie prompt path is overkill; + # instead find by dispatch arrival offset / structural position. + by_offset = {n.arrival_offset_us: nid for nid, n in nodes.items()} + p0 = by_offset[int(0.0 * 1e6)] + s1_inner = by_offset[int(1.0 * 1e6)] + s2_inner = by_offset[int(1.2 * 1e6)] + p1 = by_offset[int(4.0 * 1e6)] + + # Subagent first turns derive from the recorded spawner (p0). + assert _edge(edges, p0, s1_inner) is not None + assert _edge(edges, p0, s2_inner) is not None + # The resume turn AND-fans-in on BOTH subagent last turns. + assert _edge(edges, s1_inner, p1) is not None + assert _edge(edges, s2_inner, p1) is not None + + +def test_node_inputs_match_predecessor_edges() -> None: + """Each node's AND-fan-in ``inputs`` mirror its non-START interval-order edges. + + A multi-predecessor join node must declare one + ``ChannelRequirement(channel="{src}_out", count=1)`` per predecessor source + so the executor's ``await_inputs`` gate enforces the AND-join (trie LlmNodes + otherwise declare no inputs, so the gate is a no-op and the node early-fires + on its FIRST completing predecessor). A single-predecessor node has exactly + one requirement; a START-rooted node has none. + + Reuses the blocking-join topology: p0 roots at START (no inputs); the two + subagent first turns each wait on p0 (one input); the resume turn AND-fans-in + on the interval-order FRONTIER of its finished-before causes. + + Under interval-order timing (:func:`_build_interval_edges`) p1's predecessor + set is the MAXIMAL finished-before frontier, NOT every completed-before cause. + p0 (end 0.5) is transitively covered -- it finished before s1_inner started + (0.5 <= 1.0), so ``p0 -> s1_inner -> p1`` drops p0 from p1's frontier. The + content-parent p0 is therefore NOT a direct timing predecessor of p1; p1 + fans in on the two subagent last turns ONLY (two inputs), which the executor + still gates as an AND-join. + """ + s1 = _subagent( + t=1.0, + agent_id="agent_1", + status="completed", + duration_ms=2000, + requests=[_n(t=1.0, hashes=[50, 51], in_len=128, out=16, api_time=1.5)], + ) + s2 = _subagent( + t=1.0, + agent_id="agent_2", + status="completed", + duration_ms=2000, + requests=[_n(t=1.2, hashes=[60, 61], in_len=128, out=16, api_time=1.5)], + ) + trace = _trace( + [ + _n(t=0.0, hashes=[1, 2], in_len=128, out=64, api_time=0.5), + s1, + s2, + _n(t=4.0, hashes=[1, 2, 3], in_len=192, out=32, api_time=1.0), + ] + ) + graph, _ = _build(trace) + + nodes = _llm_nodes(graph) + by_offset = {n.arrival_offset_us: nid for nid, n in nodes.items()} + p0 = by_offset[int(0.0 * 1e6)] + s1_inner = by_offset[int(1.0 * 1e6)] + s2_inner = by_offset[int(1.2 * 1e6)] + p1 = by_offset[int(4.0 * 1e6)] + + # START-rooted node has no input requirements (fires at min_start_delay_us). + assert nodes[p0].inputs == [] + + # Single-predecessor nodes wait on exactly their one spawner (p0). + assert nodes[s1_inner].inputs == [ChannelRequirement(channel=f"{p0}_out", count=1)] + assert nodes[s2_inner].inputs == [ChannelRequirement(channel=f"{p0}_out", count=1)] + + # The resume turn AND-fans-in on its interval-order frontier: both subagent + # last turns. The content-parent p0 is NOT a direct predecessor -- it + # finished before s1_inner started, so it is transitively covered + # (p0 -> s1_inner -> p1) and dropped from p1's frontier. + assert set(nodes[p1].inputs) == { + ChannelRequirement(channel=f"{s1_inner}_out", count=1), + ChannelRequirement(channel=f"{s2_inner}_out", count=1), + } + # p0 is not a direct timing predecessor of p1 (frontier-dropped). + assert ChannelRequirement(channel=f"{p0}_out", count=1) not in nodes[p1].inputs + # Inputs match the non-START predecessor edges exactly (no extras / dups). + edges = _edges(graph) + pred_sources = {e.source for e in edges if e.target == p1 and e.source != "START"} + assert pred_sources == {s1_inner, s2_inner} + assert {req.channel for req in nodes[p1].inputs} == { + f"{src}_out" for src in pred_sources + } + assert len(nodes[p1].inputs) == len(pred_sources) + + +def test_fresh_context_subagent_uses_recorded_spawner() -> None: + """A subagent first turn sharing NO content prefix still derives from the marker. + + The subagent's inner turn has hash_ids disjoint from the parent's, so it has + NO content-parent. The recorded structural spawner (the parent turn before + the marker) supplies the dependency edge instead. p0 completes at 0.5; the + subagent starts at 1.0, so the completed-before test passes. + """ + sa = _subagent( + t=1.0, + agent_id="agent_x", + status="completed", + duration_ms=1000, + requests=[_n(t=1.0, hashes=[900, 901], in_len=128, out=16, api_time=0.3)], + ) + trace = _trace( + [ + _n(t=0.0, hashes=[1, 2], in_len=128, out=64, api_time=0.5), + sa, + ] + ) + graph, _ = _build(trace) + + nodes = _llm_nodes(graph) + assert len(nodes) == 2 + by_offset = {n.arrival_offset_us: nid for nid, n in nodes.items()} + p0 = by_offset[int(0.0 * 1e6)] + inner = by_offset[int(1.0 * 1e6)] + + edges = _edges(graph) + # No content prefix -> the only cause is the recorded spawner p0. + e = _edge(edges, p0, inner) + assert e is not None + assert e.delay_after_predecessor_us == (1.0 - 0.5) * 1e6 + # The fresh-context inner turn does NOT re-root at START (spawner completed). + assert _edge(edges, "START", inner) is None + + +# --- max_osl dispatch cap (W2) ---------------------------------------------- + + +def test_max_osl_caps_top_level_dispatch_not_subagent_body() -> None: + """``max_osl`` caps top-level chain ``max_tokens``; subagent bodies stay uncapped. + + agentx ``_cap_output`` parity: the wire cap applies to the TOP-LEVEL chain + requests only. The synthesized response SEGMENT stays sized to the recorded + ``out`` so successor prompt content (and ISL) is unchanged. + """ + sa = _subagent( + t=1.0, + agent_id="agent_x", + status="completed", + duration_ms=1000, + requests=[_n(t=1.2, hashes=[900, 901], in_len=128, out=4000, api_time=0.3)], + ) + trace = _trace( + [ + _n(t=0.0, hashes=[1, 2], in_len=128, out=5000, api_time=0.5), + sa, + ] + ) + graph, pool = build_trie_graph(trace, callbacks=_STUB_CALLBACKS, max_osl=100) + nodes = _llm_nodes(graph) + + top = nodes["trace_0:0"] + inner = nodes["agent_x:0"] + assert top.max_tokens == 100, "top-level chain capped" + assert inner.max_tokens == 4000, "subagent body uncapped" + # The response segment is NOT capped: successor prompt bytes stay recorded. + top_response = _response_segment(pool, top) + assert len(top_response.content.split("|")) == 5000 + + +def test_max_osl_above_recorded_out_is_a_noop() -> None: + trace = _trace([_n(t=0.0, hashes=[1, 2], in_len=128, out=64, api_time=0.5)]) + graph, _ = build_trie_graph(trace, callbacks=_STUB_CALLBACKS, max_osl=10_000) + assert _llm_nodes(graph)["trace_0:0"].max_tokens == 64 + + +def test_max_osl_none_leaves_recorded_out_uncapped() -> None: + trace = _trace([_n(t=0.0, hashes=[1, 2], in_len=128, out=5000, api_time=0.5)]) + graph, _ = build_trie_graph(trace, callbacks=_STUB_CALLBACKS, max_osl=None) + assert _llm_nodes(graph)["trace_0:0"].max_tokens == 5000 + + +# --- empty flattened leaf set (W7) ------------------------------------------ + + +def test_all_empty_subagent_trace_raises_empty_weka_trace_error() -> None: + """Non-empty ``requests`` that flatten to ZERO leaves must fail the parse. + + A trace whose only entries are subagent markers with empty bodies would + otherwise produce a schedulable 0-node graph that can never fire. + """ + trace = _trace( + [ + _subagent( + t=0.5, + agent_id="agent_async", + status="async_launched", + requests=[], + ) + ] + ) + with pytest.raises(EmptyWekaTraceError, match="zero normal/streaming leaf"): + _build(trace) diff --git a/tests/unit/graph/test_weka_trie_build_resolution.py b/tests/unit/graph/test_weka_trie_build_resolution.py new file mode 100644 index 0000000000..308469e3c4 --- /dev/null +++ b/tests/unit/graph/test_weka_trie_build_resolution.py @@ -0,0 +1,374 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Content-parent resolution: O(n) trie equivalence + perf gate. + +:func:`~aiperf.dataset.graph.segment_ir.trie_content.resolve_content_parents` +was an O(n^2 * m) double loop (every node scanned against all earlier nodes, +each comparison O(m) over ``hash_ids`` lists 2000+ long on real corpus traces). +It is now an incremental prefix-trie pass, O(sum of hash_ids lengths). + +These tests pin the trie pass byte-for-byte to the OLD double-loop semantics: +``_old_resolve_reference`` is the verbatim pre-fix algorithm kept here as an +ORACLE, and the equivalence test asserts the new pass returns the IDENTICAL +``content_parent`` mapping on hand-built multi-branch node sets (sequential +continuation, parallel fan-out, truncation branch, tie-break toward most-recent, +disjoint roots, exact duplicates). The perf test asserts a corpus-scale node set +resolves in well under the wall the old loop hit. +""" + +from __future__ import annotations + +import glob +import os +import time +from pathlib import Path + +import orjson +import pytest +from hypothesis import example, given, settings +from hypothesis import strategies as st + +from aiperf.dataset.graph.adapters.weka.trace_models import ( + WekaNormalRequest, + WekaTrace, +) +from aiperf.dataset.graph.adapters.weka.trie_build import _flatten_requests +from aiperf.dataset.graph.segment_ir.trie_content import ( + TrieNode, + TrieRequest, + resolve_content_parents, +) + + +def _is_full_prefix(prefix: list[int], seq: list[int]) -> bool: + """``True`` when ``prefix`` is a full leading slice of ``seq``.""" + return len(prefix) <= len(seq) and seq[: len(prefix)] == prefix + + +def _lcp_len(a: list[int], b: list[int]) -> int: + """Length of the longest common prefix of two hash-id lists.""" + n = min(len(a), len(b)) + i = 0 + while i < n and a[i] == b[i]: + i += 1 + return i + + +def _old_resolve_reference(nodes: list[TrieNode]) -> None: + """Verbatim pre-fix O(n^2 * m) resolution -- the equivalence ORACLE. + + For node R, scan all earlier nodes and pick the one whose ``hash_ids`` is the + longest FULL prefix of R's, tie-broken toward the most recent (``>=``). When + no earlier node is a full prefix, fall back to the earlier node with the + longest partial LCP, tie-broken toward the earliest (strict ``>``). With no + overlap at all, R stays a fresh root. + """ + for idx, r in enumerate(nodes): + r_hashes = r.request.hash_ids + best_full: TrieNode | None = None + best_full_len = -1 + best_partial: TrieNode | None = None + best_partial_lcp = 0 + for p in nodes[:idx]: + p_hashes = p.request.hash_ids + if p_hashes and _is_full_prefix(p_hashes, r_hashes): + if len(p_hashes) >= best_full_len: + best_full_len = len(p_hashes) + best_full = p + continue + lcp = _lcp_len(p_hashes, r_hashes) + if lcp > best_partial_lcp: + best_partial_lcp = lcp + best_partial = p + if best_full is not None: + r.content_parent = best_full + elif best_partial is not None: + r.content_parent = best_partial + + +def _mk_nodes(hash_id_lists: list[list[int]]) -> list[TrieNode]: + """Build a recorded-order ``TrieNode`` list from hash-id sequences only. + + The resolution pass touches only ``request.hash_ids`` + ``order``; the other + request fields are filled with valid placeholders. + """ + nodes: list[TrieNode] = [] + for i, hashes in enumerate(hash_id_lists): + req = WekaNormalRequest.model_validate( + { + "t": float(i), + "type": "n", + "model": "M", + "in": 1, + "out": 1, + "hash_ids": hashes, + } + ) + nodes.append(TrieNode(node_id=f"r_{i}", request=req, order=i)) + return nodes + + +def _parent_orders(nodes: list[TrieNode]) -> list[int | None]: + """Map each node's resolved content-parent to the parent's ``order`` (or None).""" + order_of = {id(n): n.order for n in nodes} + return [ + None if n.content_parent is None else order_of[id(n.content_parent)] + for n in nodes + ] + + +# Hand-built multi-branch corpus stressing every resolution branch. The trailing +# comment on each row names the case the row exercises. +_CASES = [ + pytest.param( + [[1, 2], [1, 2, 3, 4], [1, 2, 3, 4, 5]], + id="sequential_continuation", + ), + pytest.param( + [[1, 2], [1, 2, 3], [1, 2, 4]], + id="parallel_fanout_branch_point", + ), + pytest.param( + # node3 [1,2,3] is a SHORTER (truncated) reuse of node1 [1,2,3,4,5]: + # node1 is not a full prefix of node3 (longer); node0 [1,2] IS -> parent 0. + [[1, 2], [1, 2, 3, 4, 5], [1, 2, 3]], + id="truncation_branch_prefers_shorter_full_prefix", + ), + pytest.param( + # Two identical [1,2] precede R [1,2,9]; both are full prefixes of R of + # equal length -> tie-break toward the MOST RECENT (order 1, not 0). + [[1, 2], [1, 2], [1, 2, 9]], + id="full_prefix_tie_break_most_recent", + ), + pytest.param( + # Two nodes share R's partial LCP of 2 with no full prefix; the partial + # tie-break favors the EARLIEST (order 0). + [[1, 2, 7], [1, 2, 8], [1, 2, 9, 9]], + id="partial_lcp_tie_break_earliest", + ), + pytest.param( + [[1, 2], [3, 4], [5, 6]], + id="disjoint_roots_no_overlap", + ), + pytest.param( + [[], [1, 2], [], [1, 2, 3]], + id="empty_hash_ids_never_parent_or_child", + ), + pytest.param( + # A deep shared trunk with a late short reuse + a most-recent equal-length + # duplicate competing on the full-prefix tie-break. + [ + [1, 2, 3, 4, 5], + [1, 2, 3], + [1, 2, 3, 4, 5], + [1, 2, 3, 4, 5, 6], + [1, 2, 3, 7], + ], + id="mixed_trunk_reuse_and_duplicate", + ), +] + + +@pytest.mark.parametrize("hash_id_lists", _CASES) +def test_trie_resolution_matches_old_oracle( + hash_id_lists: list[list[int]], +) -> None: + """New trie pass == old O(n^2) oracle on every hand-built multi-branch set.""" + new_nodes = _mk_nodes(hash_id_lists) + resolve_content_parents(new_nodes) + + old_nodes = _mk_nodes(hash_id_lists) + _old_resolve_reference(old_nodes) + + assert _parent_orders(new_nodes) == _parent_orders(old_nodes) + + +def test_trie_resolution_explicit_expected_parents() -> None: + """Lock the resolved parents to explicit ordinals (not just oracle parity).""" + nodes = _mk_nodes( + [ + [1, 2], # 0: root (no earlier) + [1, 2, 3, 4, 5], # 1: extends 0 (full prefix 0) + [1, 2, 3], # 2: 1 is longer (partial); 0 is full prefix -> 0 + [1, 2, 3, 4, 5], # 3: equals 1 (full prefix, most recent of {1}) -> 1 + [9, 9], # 4: disjoint root + ] + ) + resolve_content_parents(nodes) + assert _parent_orders(nodes) == [None, 0, 0, 1, None] + + +def _mk_nodes_direct(hash_id_lists: list[list[int]]) -> list[TrieNode]: + """Build nodes straight from :class:`TrieRequest` (no ``ge=0`` model gate). + + :func:`resolve_content_parents` operates on the FULL opaque-key domain of + ``TrieRequest.hash_ids``: dynamo emits negative VIRTUAL ids + (``itertools.count(-1, -1)`` in ``adapters/dynamo/trie_lowering.py``) and + weka JSON ids need not fit 64 bits. ``WekaNormalRequest``'s ``ge=0`` + constraint (used by :func:`_mk_nodes`) would reject the negative-id + adversarial cases, so the differential property builds nodes directly. + """ + return [ + TrieNode( + node_id=f"r_{i}", + request=TrieRequest( + hash_ids=list(hashes), + input_length=1, + output_length=1, + t=float(i), + api_time=0.0, + ), + order=i, + ) + for i, hashes in enumerate(hash_id_lists) + ] + + +# Small-alphabet strategy: a narrow id range makes shared prefixes, exact +# duplicates, and branch points DENSE (the resolution branches that matter), +# and the bounds keep the O(n^2 * m) oracle sub-second (<=30 nodes x <=40 ids). +_HASH_ID_LISTS = st.lists( + st.lists(st.integers(min_value=-3, max_value=6), max_size=40), + max_size=30, +) + + +@settings(deadline=None, max_examples=250) +@given(node_lists=_HASH_ID_LISTS) +# hash(-1) == hash(-2) == -2 in CPython: the collision pair that distinguishes +# nested-int-dict hashing from tuple-key ((state, h)) hashing. -1 and -2 must +# stay DISTINCT sibling transitions despite the equal hash; a trailing duplicate +# [-1, -2] also pins the most-recent full-prefix tie-break under the collision. +@example(node_lists=[[-1, -2], [-1, -2, 5], [-2, -1], [-1, -2]]) +# A >2**64 id: arbitrary-width Python ints as keys (tuple-key hashing must not +# alias them -- the combined-single-int-key variant was REJECTED for this). +@example(node_lists=[[2**64 + 1, 7], [2**64 + 1, 7, 9], [2**64 + 1]]) +# Duplicate hash_ids node pair: nodes 0 and 1 are identical full prefixes of +# node 2 -> the full-prefix tie-break must select the MOST RECENT (order 1). +@example(node_lists=[[1, 2, 3], [1, 2, 3], [1, 2, 3, 4]]) +def test_trie_resolution_matches_oracle_property( + node_lists: list[list[int]], +) -> None: + """The resolution pass == the verbatim O(n^2 * m) oracle on arbitrary inputs. + + The differential harness for the flat-int-automaton rewrite: for every + generated (and every ``@example``-pinned adversarial) node set, the + content-parent mapping the trie pass produces must be IDENTICAL to the one + :func:`_old_resolve_reference` produces by brute force. This is the property + that makes the frozen-behavior claim a theorem rather than a hope. + """ + new_nodes = _mk_nodes_direct(node_lists) + resolve_content_parents(new_nodes) + + old_nodes = _mk_nodes_direct(node_lists) + _old_resolve_reference(old_nodes) + + assert _parent_orders(new_nodes) == _parent_orders(old_nodes) + + +def _synth_corpus_node_lists( + n_nodes: int, base_len: int, branch_stride: int +) -> list[list[int]]: + """A corpus-scale stressor: many nodes with long shared-prefix ``hash_ids``. + + Reproduces the O(n^2 * m) wall shape -- a deep shared trunk (so every pairwise + comparison scans ~``base_len`` ids) with periodic branches. ``branch_stride`` + forks a fresh sub-trunk so both the full-prefix and partial-LCP branches fire. + """ + lists: list[list[int]] = [] + trunk = list(range(base_len)) + for i in range(n_nodes): + if i % branch_stride == 0 and i > 0: + # Fork: replace the tail so this node branches off the shared trunk. + seq = trunk[: base_len // 2] + list(range(10_000 + i, 10_000 + i + 8)) + else: + seq = trunk + [base_len + i] + lists.append(seq) + return lists + + +def test_trie_resolution_corpus_scale_is_fast() -> None: + """O(n) trie pass resolves a corpus-scale node set well under the old wall. + + A 466-node trace with ~2000-id ``hash_ids`` is the real corpus shape that the + old O(n^2 * m) loop ground through in seconds-to-minutes; this synthetic set + is comparably sized. The trie pass must finish in a small fraction of that. + Equivalence to the old oracle is asserted alongside the timing so a faster but + wrong implementation cannot pass. + """ + lists = _synth_corpus_node_lists(n_nodes=500, base_len=2000, branch_stride=37) + + new_nodes = _mk_nodes(lists) + t0 = time.perf_counter() + resolve_content_parents(new_nodes) + elapsed = time.perf_counter() - t0 + + old_nodes = _mk_nodes(lists) + _old_resolve_reference(old_nodes) + assert _parent_orders(new_nodes) == _parent_orders(old_nodes) + + assert elapsed < 2.0, f"trie resolution took {elapsed:.3f}s (expected < 2s)" + + +# Real corpus traces are not committed (multi-MB each). The real-corpus perf +# test reads the largest suitable trace from the directory named by the +# AIPERF_TEST_WEKA_CORPUS_DIR env var and skips when it is unset, so CI stays +# green without the corpus while a developer with a local corpus can reproduce +# the unblock claim. +_CORPUS_DIR_ENV = "AIPERF_TEST_WEKA_CORPUS_DIR" + + +# Cap the oracle-equivalence trace so the deliberately-slow O(n^2 * m) reference +# (the thing we replaced) stays bounded; the new pass is timed on it too. +_MAX_ORACLE_TRACE_BYTES = 10 * 1024 * 1024 + + +def _largest_local_trace(max_bytes: int | None = None) -> Path | None: + """Largest ``*.json`` (optionally <= ``max_bytes``) in the env-named corpus dir.""" + corpus_dir = os.environ.get(_CORPUS_DIR_ENV) + if not corpus_dir: + return None + candidates = glob.glob(os.path.join(corpus_dir, "*.json")) + if max_bytes is not None: + candidates = [c for c in candidates if os.path.getsize(c) <= max_bytes] + if not candidates: + return None + return Path(max(candidates, key=os.path.getsize)) + + +def test_real_corpus_parent_resolution_under_10s() -> None: + """Content-parent resolution of a REAL corpus trace completes fast. + + The pre-fix O(n^2 * m) double loop ground through real corpus traces + (hundreds of nodes, ``hash_ids`` thousands long) for seconds-to-minutes and + blocked corpus-scale builds. The trie pass must resolve the largest local + trace in well under 10s. Equivalence to the old oracle is re-asserted on the + real node set so a wrong-but-fast pass cannot slip through. + + NOTE (second bottleneck, out of scope here): the FULL ``build_trie_graph`` of + this same trace still exceeds 120s, dominated by per-node content-synthesis + re-replay (``_replay_lineage`` -> ``SegmentPool.add`` -> ``segment_id`` token + hashing), NOT by parent resolution. + """ + trace_path = _largest_local_trace(max_bytes=_MAX_ORACLE_TRACE_BYTES) + if trace_path is None: + pytest.skip( + f"no local weka corpus trace available: set {_CORPUS_DIR_ENV} " + "to a directory of raw weka trace .json files to run this test" + ) + + trace = WekaTrace.model_validate(orjson.loads(trace_path.read_bytes())) + nodes = _flatten_requests(trace.requests, root_scope=trace.id) + + t0 = time.perf_counter() + resolve_content_parents(nodes) + elapsed = time.perf_counter() - t0 + + oracle = _flatten_requests(trace.requests, root_scope=trace.id) + _old_resolve_reference(oracle) + assert _parent_orders(nodes) == _parent_orders(oracle) + + assert elapsed < 10.0, ( + f"real-corpus parent resolution took {elapsed:.3f}s " + f"(trace {trace_path.name}, {len(nodes)} nodes, expected < 10s)" + ) diff --git a/tests/unit/graph/test_weka_trie_hash_scope.py b/tests/unit/graph/test_weka_trie_hash_scope.py new file mode 100644 index 0000000000..194391328e --- /dev/null +++ b/tests/unit/graph/test_weka_trie_hash_scope.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hash-id scope: per-trace vs cross-trace content identity in the weka trie path. + +Weka traces declare ``hash_id_scope`` as either ``"local"`` or ``"global"``: + +* ``"local"``: hash ids are a PER-TRACE namespace, so hash_id 0 in trace A and + hash_id 0 in trace B are different logical blocks and must synthesize + different bytes. Deciding content from the hash id alone would manufacture + cross-trace KV-cache prefix sharing at the server that the recorded workload + never had. +* ``"global"``: ONE hash namespace is shared across every trace in the corpus, + so equal hash ids must synthesize byte-identical blocks -- and, because + segment ids are content-addressed, identical pool segment ids -- reproducing + the recorded cross-trace KV-cache sharing. Response/tail content stays + trace-scoped in both modes: node ids recur across traces and are not part of + the recorded hash namespace. + +Any other scope value is rejected with :class:`WekaHashScopeError`. + +These tests drive the REAL synthesizer callbacks (builtin tokenizer), not the +stub callbacks the topology tests inject, because the scoping lives in the +production callback layer. +""" + +from __future__ import annotations + +from pathlib import Path + +import orjson +import pytest + +from aiperf.dataset.graph.adapters.weka.trace import ( + WekaHashScopeError, + from_weka_trace, +) + +_SEED = 1234 + + +def _trace_dict(trace_id: str, scope: str = "local") -> dict: + return { + "id": trace_id, + "models": ["m"], + "block_size": 64, + "hash_id_scope": scope, + "requests": [ + { + "type": "n", + "t": 0.0, + "api_time": 1.0, + "model": "m", + "hash_ids": [0, 1], + "in": 128, + "out": 8, + }, + { + "type": "n", + "t": 5.0, + "api_time": 1.0, + "model": "m", + "hash_ids": [0, 1, 2], + "in": 192, + "out": 8, + }, + ], + } + + +def _parse( + tmp_path: Path, trace_id: str, scope: str = "local" +) -> tuple[list[str], list[str], str]: + """Parse one trace; return (r_0 prompt path, r_0 prompt contents, r_0 response content).""" + p = tmp_path / f"{trace_id}-{scope}.json" + p.write_bytes(orjson.dumps(_trace_dict(trace_id, scope))) + parsed = from_weka_trace( + str(p), content_root_seed=_SEED, content_tokenizer="builtin" + ) + pool = parsed.segment_pool + assert pool is not None + node = parsed.graph.nodes[f"{trace_id}:0"] + path = node.metadata["trie"]["prompt_segment_ids"] + contents = [m["content"] for m in pool.materialize(path)] + response = next( + s + for s in pool.by_id.values() + if s.role == "assistant" and s.parent_id == path[-1] + ) + return path, contents, response.content + + +def test_local_scope_traces_synthesize_distinct_content(tmp_path: Path) -> None: + """Same local hash ids in two different traces => different bytes + segment ids.""" + path_a, contents_a, response_a = _parse(tmp_path, "trace-A") + path_b, contents_b, response_b = _parse(tmp_path, "trace-B") + + assert contents_a != contents_b, ( + "hash_id_scope='local' requires per-trace content identity; identical " + "bytes across traces manufacture cross-trace KV-cache sharing" + ) + assert path_a != path_b, ( + "content-addressed segment ids must not collide across local-scope traces" + ) + assert response_a != response_b, ( + "response segments seeded only by node_id collide across traces" + ) + + +def test_global_scope_traces_synthesize_identical_content(tmp_path: Path) -> None: + """Same global hash ids in two different traces => identical bytes + segment ids. + + This is the cross-trace KV-cache sharing contract: under + ``hash_id_scope='global'`` equal hash ids are the SAME logical block, so + two traces sharing a hash prefix must materialize byte-identical prompt + content -- and, segment ids being content-addressed, the SAME pool entries. + Responses stay trace-scoped: they are per-request output, not part of the + recorded hash namespace. + """ + path_a, contents_a, response_a = _parse(tmp_path, "trace-A", scope="global") + path_b, contents_b, response_b = _parse(tmp_path, "trace-B", scope="global") + + assert contents_a == contents_b, ( + "hash_id_scope='global' requires one cross-trace hash namespace; " + "differing bytes for equal hash ids drop the recorded sharing" + ) + assert path_a == path_b, ( + "content-addressed segment ids must dedup across global-scope traces" + ) + assert response_a != response_b, ( + "response tails must stay trace-scoped under global scope; node ids " + "recur in every trace and are not part of the hash namespace" + ) + + +def test_global_scope_differs_from_local_scope_bytes(tmp_path: Path) -> None: + """Scope selects the decode namespace: same trace, different scope => different bytes.""" + _path_l, contents_local, _resp_l = _parse(tmp_path, "trace-A", scope="local") + _path_g, contents_global, _resp_g = _parse(tmp_path, "trace-A", scope="global") + + assert contents_local != contents_global, ( + "global scope must decode from the bare hash-id namespace, not the " + "trace-scoped one" + ) + + +@pytest.mark.parametrize("scope", ["local", "global"]) +def test_same_trace_reparse_is_byte_identical(tmp_path: Path, scope: str) -> None: + """Scoped synthesis stays deterministic: same trace => same bytes.""" + path_1, contents_1, response_1 = _parse(tmp_path, "trace-A", scope=scope) + path_2, contents_2, response_2 = _parse(tmp_path, "trace-A", scope=scope) + + assert path_1 == path_2 + assert contents_1 == contents_2 + assert response_1 == response_2 + + +@pytest.mark.parametrize("scope", ["local", "global"]) +def test_token_counts_unaffected_by_hash_scope(tmp_path: Path, scope: str) -> None: + """Scoping changes bytes, never geometry: covered-count ISL holds per trace.""" + from aiperf.dataset.graph.adapters.shared.content import ( + get_or_build_synthesizer, + ) + + synth = get_or_build_synthesizer("builtin", prompt_corpus="coding", root_seed=_SEED) + encode = synth._pg.tokenizer.encode + for trace_id in ("trace-A", "trace-B"): + _path, contents, _resp = _parse(tmp_path, trace_id, scope=scope) + assert sum(len(encode(c)) for c in contents) == 128 + + +def test_unrecognized_hash_scope_rejected(tmp_path: Path) -> None: + """A scope outside {'local', 'global'} raises WekaHashScopeError with the cause.""" + p = tmp_path / "bad-scope.json" + p.write_bytes(orjson.dumps(_trace_dict("trace-A", scope="cluster"))) + + with pytest.raises(WekaHashScopeError, match="hash_id_scope='cluster'"): + from_weka_trace(str(p), content_root_seed=_SEED, content_tokenizer="builtin") + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit/graph/test_weka_trie_idle_gap_cap.py b/tests/unit/graph/test_weka_trie_idle_gap_cap.py new file mode 100644 index 0000000000..bcc56b463e --- /dev/null +++ b/tests/unit/graph/test_weka_trie_idle_gap_cap.py @@ -0,0 +1,330 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``build_trie_graph`` must honor the per-trace idle-gap cap (warp). + +The trie IR derives node arrival offsets and ``StaticEdge`` delays from the +recorded ``request.t`` / ``api_time``. A recorded multi-hour IDLE gap (dead air +between one request finishing and the next starting) would otherwise survive +verbatim into the warmup phase and park it forever. ``idle_gap_cap_seconds`` +builds an ``_ActiveIdleWarp`` over the UNION of every request's active interval +in the tree and places node start times, edge ``delay_after_predecessor_us``, +and root ``min_start_delay_us`` on the same idle-compressed clock. Only true +idle stretches are cut; active processing and overlapping subagents keep their +exact temporal shape. + +These tests drive the builder with the deterministic stub callbacks the other +trie-build tests use, so no tokenizer / corpus build is needed; they assert the +warped TIMING geometry, not content bytes. +""" + +from __future__ import annotations + +from aiperf.dataset.graph.adapters.weka.trace_models import WekaTrace +from aiperf.dataset.graph.adapters.weka.trie_build import ( + ReconCallbacks, + build_trie_graph, +) +from aiperf.dataset.graph.models import LlmNode, ParsedGraph, StaticEdge +from aiperf.dataset.graph.segment_ir.pool import SegmentPool + +BLOCK_SIZE = 64 + + +def _stub_decode_block_tokens(hash_ids: list[int]) -> list[int]: + out: list[int] = [] + for h in hash_ids: + out.extend(range(h * 100, h * 100 + BLOCK_SIZE)) + return out + + +def _stub_partial_tail_tokens(n_tokens: int, seed: str) -> list[int]: + base = sum(ord(c) for c in seed) * 1000 + return list(range(base, base + n_tokens)) + + +def _stub_decode_tokens_to_text(tokens: list[int]) -> str: + return "|".join(str(t) for t in tokens) + + +_STUB_CALLBACKS = ReconCallbacks( + decode_block_tokens=_stub_decode_block_tokens, + sample_partial_tail_tokens=_stub_partial_tail_tokens, + decode_tokens_to_text=_stub_decode_tokens_to_text, +) + + +def _n(t: float, hashes: list[int], in_len: int, out: int, api_time: float, + model: str = "M", type_: str = "n") -> dict: # fmt: skip + return { + "t": t, + "type": type_, + "model": model, + "in": in_len, + "out": out, + "hash_ids": hashes, + "api_time": api_time, + } + + +def _trace(requests: list[dict]) -> WekaTrace: + return WekaTrace.model_validate( + { + "id": "trace_0", + "models": ["M"], + "block_size": BLOCK_SIZE, + "hash_id_scope": "local", + "requests": requests, + } + ) + + +def _build(trace: WekaTrace, **kwargs: object) -> tuple[ParsedGraph, SegmentPool]: + return build_trie_graph(trace, callbacks=_STUB_CALLBACKS, **kwargs) + + +def _llm_nodes(graph: ParsedGraph) -> dict[str, LlmNode]: + return {nid: n for nid, n in graph.graph.nodes.items() if isinstance(n, LlmNode)} + + +def _edges(graph: ParsedGraph) -> list[StaticEdge]: + return [e for e in graph.graph.edges if isinstance(e, StaticEdge)] + + +def _edge(edges: list[StaticEdge], source: str, target: str) -> StaticEdge | None: + for e in edges: + if e.source == source and e.target == target: + return e + return None + + +def test_gap_over_cap_is_compressed() -> None: + """A raw IDLE gap >> cap is compressed to the cap (cut between FINISH/START). + + turn0: t=0, api_time=2 -> completes at 2.0 (raw). turn1: raw t=137124, sharing + turn0's hash prefix. The idle warp cuts the dead air between turn0 FINISHING + (raw end 2.0) and turn1 STARTING (raw 137124) -- 137122s of idle -- down to + the 60s cap. So turn1's warped start is ``end(2) + 60 = 62.0`` and the + end-to-start edge delay is the capped idle ``(62 - 2) * 1e6 = 60s`` -- NOT the + raw 137122s, and NOT ``0 + 60`` (capping start-to-start would wrongly eat into + turn0's own 2s of processing). + """ + trace = _trace( + [ + _n(t=0.0, hashes=[1, 2], in_len=128, out=64, api_time=2.0), + _n(t=137124.0, hashes=[1, 2, 3, 4], in_len=256, out=64, api_time=1.0), + ] + ) + graph, _ = _build(trace, idle_gap_cap_seconds=60.0) + + nodes = _llm_nodes(graph) + ids = list(nodes.keys()) + n0, n1 = ids[0], ids[1] + + # turn1 starts 60s (capped idle) after turn0's recorded END (2.0). + assert nodes[n1].arrival_offset_us == int(round(62.0 * 1e6)) + + edges = _edges(graph) + e = _edge(edges, n0, n1) + assert e is not None + # End-to-start gap on the warped clock: warped_start(62) - warped_end(2) = 60. + assert e.delay_after_predecessor_us == 60.0 * 1e6 + + +def test_sub_cap_gap_passes_through() -> None: + """A raw gap below the cap is unchanged -- the warp leaves it intact.""" + trace = _trace( + [ + _n(t=0.0, hashes=[1, 2], in_len=128, out=64, api_time=2.0), + _n(t=3.0, hashes=[1, 2, 3, 4], in_len=256, out=64, api_time=1.0), + ] + ) + graph, _ = _build(trace, idle_gap_cap_seconds=60.0) + + nodes = _llm_nodes(graph) + ids = list(nodes.keys()) + n0, n1 = ids[0], ids[1] + + assert nodes[n1].arrival_offset_us == int(round(3.0 * 1e6)) + e = _edge(_edges(graph), n0, n1) + assert e is not None + assert e.delay_after_predecessor_us == (3.0 - 2.0) * 1e6 + + +def test_cap_none_is_raw_passthrough() -> None: + """``idle_gap_cap_seconds=None`` reproduces the pre-change RAW geometry exactly. + + Builds the same trace twice -- once with the param omitted (the historic call + shape) and once explicitly ``None`` -- and asserts both yield the RAW edge + delays and arrival offsets (no warp applied). Regression guard for the no-cap + path. + """ + trace = _trace( + [ + _n(t=0.0, hashes=[1, 2], in_len=128, out=64, api_time=2.0), + _n(t=137124.0, hashes=[1, 2, 3, 4], in_len=256, out=64, api_time=1.0), + ] + ) + graph_default, _ = _build(trace) + graph_none, _ = _build(trace, idle_gap_cap_seconds=None) + + for graph in (graph_default, graph_none): + nodes = _llm_nodes(graph) + ids = list(nodes.keys()) + n0, n1 = ids[0], ids[1] + # Raw arrival offset survives -- no compression. + assert nodes[n1].arrival_offset_us == int(round(137124.0 * 1e6)) + e = _edge(_edges(graph), n0, n1) + assert e is not None + assert e.delay_after_predecessor_us == (137124.0 - 2.0) * 1e6 + + +def test_multi_gap_cumulative_shift_matches_warp() -> None: + """3 requests with TWO >cap IDLE gaps: cumulative shift accrues across both. + + Proves the build accumulates the idle-warp shift (each prior over-cap idle + gap shifts later events left), not a per-edge ``min(gap, cap)``: a per-edge + clamp would lose the first gap's shift when placing the third node. Each + capped idle sits AFTER the preceding request's END (interval warp), so: + turn0 [0,1]; idle 1->1000 (999>cap) -> turn1 start = end(1)+60 = 61; + turn1 [61,62] warped; idle 1001->3000 raw (1999>cap) -> turn2 start = + warped_end(62)+60 = 122. + """ + trace = _trace( + [ + _n(t=0.0, hashes=[1, 2], in_len=128, out=64, api_time=1.0), + _n(t=1000.0, hashes=[1, 2, 3], in_len=192, out=32, api_time=1.0), + _n(t=3000.0, hashes=[1, 2, 3, 4], in_len=256, out=32, api_time=1.0), + ] + ) + graph, _ = _build(trace, idle_gap_cap_seconds=60.0) + + nodes = _llm_nodes(graph) + vals = list(nodes.values()) + # turn1: end(1) + 60s capped idle = 61. + assert vals[1].arrival_offset_us == int(round(61.0 * 1e6)) + # turn2: warped_end(62) + 60s capped idle = 122 -- cumulative across BOTH gaps. + assert vals[2].arrival_offset_us == int(round(122.0 * 1e6)) + + +def _subagent(t: float, agent_id: str, requests: list[dict], status: str) -> dict: + """A blocking/async subagent marker with nested inner leaf requests.""" + return { + "t": t, + "type": "subagent", + "agent_id": agent_id, + "subagent_type": "Explore", + "status": status, + "duration_ms": None, + "requests": requests, + "models": ["M"], + } + + +def _replay_starts(graph: ParsedGraph, api: dict[str, float]) -> dict[str, float]: + """Mini discrete-event replay: fire each node at max over incoming edges of + (sim_end(pred) + delay) [START-roots at min_start_delay; a start-anchored + edge fires at sim_start(pred) + start_delay], occupying its recorded + ``api_time``. Returns each node's reconstructed start. Mirrors + tools/weka_trie_timing_sim.py so the byte-exact invariant is unit-locked.""" + incoming: dict[str, list[StaticEdge]] = {nid: [] for nid in graph.graph.nodes} + for e in _edges(graph): + if e.target in incoming: + incoming[e.target].append(e) + nodes = _llm_nodes(graph) + order = sorted(nodes, key=lambda nid: (nodes[nid].arrival_offset_us or 0, nid)) + sim_start: dict[str, float] = {} + sim_end: dict[str, float] = {} + for nid in order: + gate = 0.0 + for e in incoming[nid]: + if e.source == "START": + gate = max(gate, (e.min_start_delay_us or 0.0) / 1e6) + elif e.delay_after_predecessor_start_us is not None: + # Start-anchored: gate off the predecessor's DISPATCH, not its end. + gate = max( + gate, + sim_start.get(e.source, 0.0) + + e.delay_after_predecessor_start_us / 1e6, + ) + elif e.source in sim_end: + gate = max( + gate, + sim_end[e.source] + (e.delay_after_predecessor_us or 0.0) / 1e6, + ) + sim_start[nid] = gate + sim_end[nid] = gate + api.get(nid, 0.0) + return sim_start + + +def test_concurrent_turn_overlapping_its_cause_roots_at_warped_offset() -> None: + """A subagent launched WHILE its spawner is still running start-anchors to it. + + p0 runs [0, 5]; its subagent's first turn starts at t=2 -- before p0 finished. + It did not WAIT for p0's completion, so it must NOT get an end-anchored + ``p0 -> s0`` completion edge; instead the causal-parent stamping + start-anchor + post-pass collapse its incoming edges to ONE start-anchored ``p0 -> s0`` edge + (``delay_after_predecessor_start_us``) whose delay is the warped start-to-start + gap (t=2). The runtime schedules s0 at p0's DISPATCH and gates it 2 s later, + so s0 still fires at its warped arrival t=2 -- concurrency preserved, no + START re-root. + """ + trace = _trace( + [ + _n(t=0.0, hashes=[1, 2], in_len=128, out=64, api_time=5.0), + _subagent( + t=2.0, + agent_id="a1", + status="completed", + requests=[ + _n(t=2.0, hashes=[1, 2, 3], in_len=192, out=32, api_time=3.0), + _n(t=6.0, hashes=[1, 2, 3, 4], in_len=256, out=32, api_time=2.0), + ], + ), + _n(t=9.0, hashes=[1, 2, 9], in_len=320, out=16, api_time=1.0), + ] + ) + graph, _ = _build(trace, idle_gap_cap_seconds=60.0) + edges = _edges(graph) + # s0 = r_1_0 (first inner of the subagent at index 1). It overlapped p0 (r_0), + # so its incoming edges collapse to a single start-anchored p0 -> r_1_0 edge; + # it does NOT re-root at START. + assert _edge(edges, "START", "a1:0") is None + incoming = [e for e in edges if e.target == "a1:0"] + assert len(incoming) == 1 + (anchor_edge,) = incoming + assert anchor_edge.source == "trace_0:0" + assert anchor_edge.delay_after_predecessor_start_us == 2.0 * 1e6 + assert anchor_edge.delay_after_predecessor_us is None + assert anchor_edge.min_start_delay_us is None + + +def test_replay_reconstructs_recorded_timeline_byte_exact() -> None: + """Replaying the trie edges with recorded api_time lands every node on its + recorded (warp-free here -- no >cap idle) start: p0=0, s0=2, s1=6, resume=9. + + Locks the end-to-end timing invariant the standalone simulator proves on the + corpus (1039/1039): concurrent overlap (s0), sequential continuation (s1), + and a parent-resume binding the final child (resume after s1 completes). + """ + trace = _trace( + [ + _n(t=0.0, hashes=[1, 2], in_len=128, out=64, api_time=5.0), + _subagent( + t=2.0, + agent_id="a1", + status="completed", + requests=[ + _n(t=2.0, hashes=[1, 2, 3], in_len=192, out=32, api_time=3.0), + _n(t=6.0, hashes=[1, 2, 3, 4], in_len=256, out=32, api_time=2.0), + ], + ), + _n(t=9.0, hashes=[1, 2, 9], in_len=320, out=16, api_time=1.0), + ] + ) + graph, _ = _build(trace, idle_gap_cap_seconds=60.0) + api = {"trace_0:0": 5.0, "a1:0": 3.0, "a1:1": 2.0, "trace_0:1": 1.0} + sim = _replay_starts(graph, api) + assert sim["trace_0:0"] == 0.0 + assert sim["a1:0"] == 2.0 # concurrent: launched during p0 + assert sim["a1:1"] == 6.0 # sequential after s0 (end 5) + 1s think + assert sim["trace_0:1"] == 9.0 # parent-resume: 1s after final child s1 (end 8) diff --git a/tests/unit/graph/test_weka_trie_interval_order.py b/tests/unit/graph/test_weka_trie_interval_order.py new file mode 100644 index 0000000000..6a36f02b85 --- /dev/null +++ b/tests/unit/graph/test_weka_trie_interval_order.py @@ -0,0 +1,540 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Structural invariants of the interval-order + message-unit Weka trie builder. + +Covers the two reconstruction contracts (with deterministic stub callbacks): + +* **Timing** -- global interval-order edges: ``async_ancestors`` stamping, a + time-consistent ``rank``, the finished-before frontier (transitive reduction), + async-subtree exclusion, coincident/concurrent handling, and START-rooting. +* **Content** -- per-block frozen ``(role, starts_new_message)`` tags + (creator-frozen, inherited verbatim, trailing-user at creation), message-unit + content-addressed emission, shared-prefix identical message-id chains, boundary + preservation, and the block-aligned covered-count ISL gate. + +Scope: these assert the *structural* invariant only -- properties of the +reconstructed IR. The end-to-end cache-hit claim (a shared block prefix yielding +a real KV prefix-cache hit) is provable only against a real prefix-caching engine +(vLLM/SGLang/real Dynamo); AIPerf's mock server is throughput-only and cannot +validate it. See ``docs/reference/graph-ingest-build-pipeline.md`` -> +"Validation boundary". +""" + +import pytest + +from aiperf.dataset.graph.adapters.weka.trace_models import ( + WekaNormalRequest, + WekaSubagentEntry, +) +from aiperf.dataset.graph.adapters.weka.trie_build import ( + _flatten_requests, +) +from aiperf.dataset.graph.segment_ir.interval_order import ( + build_interval_edges as _build_interval_edges, +) +from aiperf.dataset.graph.segment_ir.interval_order import ( + compute_ranks as _compute_ranks, +) + + +def _n(t, hid, out=0, api=0.0, in_=0): + return WekaNormalRequest( + t=t, type="n", model="m", **{"in": in_}, out=out, hash_ids=hid, api_time=api + ) + + +def test_flatten_stamps_async_ancestors(): + reqs = [ + _n(0.0, [1]), + WekaSubagentEntry( + t=1.0, + type="subagent", + agent_id="a", + subagent_type="X", + status="completed", + requests=[_n(1.1, [1, 2])], + models=["m"], + ), + WekaSubagentEntry( + t=2.0, + type="subagent", + agent_id="b", + subagent_type="X", + status="async_launched", + requests=[_n(2.1, [1, 3])], + models=["m"], + ), + _n(3.0, [1, 4]), + ] + by = { + tuple(n.request.hash_ids): n for n in _flatten_requests(reqs, root_scope="tr") + } + assert by[(1, 3)].async_ancestors and not by[(1, 2)].async_ancestors + assert not by[(1,)].async_ancestors and not by[(1, 4)].async_ancestors + + +def test_rank_total_order_and_coincident_tiebreak(): + nodes = _flatten_requests( + [_n(0.0, [1], api=0.0), _n(0.0, [2], api=0.0)], root_scope="tr" + ) # coincident zero-duration + for x in nodes: + x.warped_start = x.request.t + _compute_ranks(nodes) + assert sorted(n.rank for n in nodes) == [0, 1] # unique, total order + a, b = sorted(nodes, key=lambda n: n.node_id) + assert a.rank < b.rank # lower node_id -> lower rank on a full tie + + +def _prep(reqs): + nodes = _flatten_requests(reqs, root_scope="tr") + for x in nodes: + x.warped_start = x.request.t + _compute_ranks(nodes) + return nodes + + +def test_overlap_concurrent_finishedbefore_andjoin(): + nodes = _prep( + [ + _n(0.0, [1], api=1.0), # P0 [0,1] + _n(1.2, [1, 2], api=2.8), # A [1.2,4] + _n(1.3, [1, 3], api=3.7), # B [1.3,5] + _n(5.2, [1, 4], api=1.8), # C [5.2,7] + ] + ) + edges = _build_interval_edges(nodes) + + def srcs(n): + return {e.source for e in edges[n.node_id]} + + a, b, c = nodes[1], nodes[2], nodes[3] + assert b.node_id not in srcs(a) and a.node_id not in srcs(b) # racers concurrent + assert srcs(c) == {a.node_id, b.node_id} # AND-join both + ce = {e.source: e for e in edges[c.node_id]} + assert ce[b.node_id].delay_after_predecessor_us == pytest.approx( + (5.2 - 5.0) * 1e6 + ) # binding=B (latest end) + assert ce[a.node_id].delay_after_predecessor_us == 0.0 + + +def test_empty_frontier_roots_at_start(): + nodes = _prep([_n(0.0, [1], api=1.0)]) + e = _build_interval_edges(nodes)[nodes[0].node_id] + assert ( + len(e) == 1 + and e[0].source == "START" + and e[0].min_start_delay_us == pytest.approx(0.0) + ) + + +def test_async_excluded_cross_subtree(): + nodes = _prep( + [ + _n(0.0, [1], api=0.5), + WekaSubagentEntry( + t=0.6, + type="subagent", + agent_id="b", + subagent_type="X", + status="async_launched", + requests=[_n(0.6, [1, 9], api=0.2)], + models=["m"], + ), + _n(3.0, [1, 4], api=0.5), # P1 starts after the async child ended + ] + ) + edges = _build_interval_edges(nodes) + p1 = next(n for n in nodes if tuple(n.request.hash_ids) == (1, 4)) + b = next(n for n in nodes if tuple(n.request.hash_ids) == (1, 9)) + assert b.node_id not in { + e.source for e in edges[p1.node_id] + } # fire-and-forget child not a pred + + +# --- Task 4B: block geometry + role split + trailing-user caps (foundation) --- + +from aiperf.dataset.graph.segment_ir.trie_content import ( # noqa: E402 + block_role_split, + compute_asst_caps, + compute_turn_block_geometry, + resolve_content_parents, +) + + +def test_geometry_block_aligned_no_partial_tail(): + # in=5, bs=2 -> in//bs=2 covered blocks; 3 hash ids but only 2 covered; NO in%bs tail + g = compute_turn_block_geometry([1], [1, 2, 3], 5, 2) + assert ( + g.lcp == 1 + and g.m_curr_covered == 2 + and g.new_blocks_count == 1 + and g.synth_tail_n == 0 + ) + # missing whole blocks: in//bs=3 but only 2 hash ids -> 1 missing block = 2 tokens + g2 = compute_turn_block_geometry([], [1, 2], 6, 2) + assert g2.m_curr_covered == 2 and g2.synth_tail_n == 2 + + +def test_block_role_split_trailing_user_at_creation(): + # prev_out large: ceil(6/2)=3 asst, clamped to 2 new blocks -> asst==new_n, so the + # last new block flips to user (trailing-user frozen at creation). + inherited, roles = block_role_split( + prev_hash_ids=[1], + curr_hash_ids=[1, 2, 3], + curr_in_tokens=6, + prev_out_tokens=6, + block_size=2, + max_asst_blocks=None, + parent_has_user=True, + ) + # last new block flipped to user + assert inherited == 1 + assert roles == ["assistant", "user"] + # context-loss: parent_has_user False -> all user + _, roles2 = block_role_split( + prev_hash_ids=[1], + curr_hash_ids=[1, 2, 3], + curr_in_tokens=6, + prev_out_tokens=6, + block_size=2, + max_asst_blocks=None, + parent_has_user=False, + ) + assert roles2 == ["user", "user"] + # single all-assistant new block flips wholesale to user (no assistant survives) + _, roles3 = block_role_split( + prev_hash_ids=[1], + curr_hash_ids=[1, 2], + curr_in_tokens=4, + prev_out_tokens=10, + block_size=2, + max_asst_blocks=None, + parent_has_user=True, + ) + assert roles3 == ["user"] + + +def test_compute_asst_caps_skips_root_owner(): + # single root node, degenerate: no cap should be recorded on a root owner + nodes = _flatten_requests( + [_n(0.0, [1, 2], in_=4), _n(0.2, [1, 2], in_=4)], root_scope="tr" + ) + resolve_content_parents(nodes) + caps = compute_asst_caps(nodes, 2) + roots = [n for n in nodes if n.content_parent is None] + for r in roots: + assert caps.get(r.node_id) is None # root owner never capped + + +# --- Task 4: per-block (role, starts_new_message) tag pass, frozen at creation --- + +from aiperf.dataset.graph.segment_ir.trie_content import ( # noqa: E402 + assign_block_tags, +) + + +def _tags_for(reqs, bs): + nodes = _flatten_requests(reqs, root_scope="tr") + resolve_content_parents(nodes) + caps = compute_asst_caps(nodes, bs) + return nodes, assign_block_tags(nodes, bs, caps) + + +def test_shared_prefix_tags_identical_and_no_relabel(): + # A long parent chain that per-turn would relabel a shared block, + a fork inheriting it. + reqs = [ + _n(0.0, [10, 11, 12], in_=6, out=0, api=0.1), # P0 user blocks + _n(0.2, [10, 11], in_=4, out=400, api=0.1), # P1 pull-back w/ big prev_out + _n(0.4, [10, 11, 20], in_=6, out=0, api=0.1), # P2 continues, shares [10,11] + WekaSubagentEntry( + t=1.0, + type="subagent", + agent_id="a", + subagent_type="X", + status="completed", + requests=[_n(1.1, [10, 11, 30], in_=6, out=0, api=0.1)], + models=["m"], + ), # fork inherits [10,11] + ] + nodes, tags = _tags_for(reqs, 2) + by = lambda h: next(n for n in nodes if tuple(n.request.hash_ids) == h) # noqa: E731 + fork, p2 = by((10, 11, 30)), by((10, 11, 20)) + # blocks [10,11] (indices 0,1) have IDENTICAL (role, starts_new_message) in fork and P2 + assert tags[fork.node_id][:2] == tags[p2.node_id][:2] + + +def test_contiguous_same_role_turns_keep_boundary(): + # two consecutive user turns (prev_out=0) -> the 2nd turn's first new block starts a new message + reqs = [_n(0.0, [1], in_=2, out=0, api=0.1), _n(0.2, [1, 2], in_=4, out=0, api=0.1)] + nodes, tags = _tags_for(reqs, 2) + p = next(n for n in nodes if tuple(n.request.hash_ids) == (1, 2)) + # block 0 (user, starts), block 1 is the new turn's block -> user, starts_new_message True + assert tags[p.node_id][1] == ("user", True) + + +# --- Task 5: message-unit content emission from frozen per-block tags --- + +from aiperf.dataset.graph.segment_ir.pool import ( # noqa: E402 + SegmentPool, +) +from aiperf.dataset.graph.segment_ir.trie_content import ( # noqa: E402 + assemble_messages, +) + + +def _decode(hids): + return [h * 1000 + k for h in hids for k in range(2)] # bs=2, collision-free + + +def _text(toks): + return " ".join(map(str, toks)) + + +def test_shared_prefix_identical_message_ids(): + reqs = [ + _n(0.0, [10, 11], in_=4, out=0, api=0.1), # P0 + WekaSubagentEntry( + t=1.0, + type="subagent", + agent_id="a", + subagent_type="X", + status="completed", + requests=[_n(1.1, [10, 11, 30], in_=6, out=0, api=0.1)], + models=["m"], + ), # fork inherits [10,11] + _n(2.0, [10, 11, 20], in_=6, out=0, api=0.1), # P1 shares [10,11] + ] + nodes, tags = _tags_for(reqs, 2) + by = lambda h: next( # noqa: E731 + n for n in nodes if tuple(n.request.hash_ids) == h + ) + + def ids(node): + return assemble_messages( + node.request.hash_ids, tags[node.node_id], pool, _decode, _text + )[0] + + pool = SegmentPool() + fork_ids = ids(by((10, 11, 30))) + p1_ids = ids(by((10, 11, 20))) + p0_ids = ids(by((10, 11))) + # the message(s) covering the shared [10,11] prefix are byte-identical ids + # across all three + assert fork_ids[: len(p0_ids)] == p0_ids == p1_ids[: len(p0_ids)] + + +def test_contiguous_same_role_not_coalesced(): + # two user turns -> block1 has starts_new_message True -> two user messages + reqs = [ + _n(0.0, [1], in_=2, out=0, api=0.1), + _n(0.2, [1, 2], in_=4, out=0, api=0.1), + ] + nodes, tags = _tags_for(reqs, 2) + pool = SegmentPool() + p = next(n for n in nodes if tuple(n.request.hash_ids) == (1, 2)) + msg_ids, _ = assemble_messages( + p.request.hash_ids, tags[p.node_id], pool, _decode, _text + ) + roles = [pool.materialize([mid])[0]["role"] for mid in msg_ids] + assert roles == ["user", "user"] and len(msg_ids) == 2 # boundary preserved + + +def test_last_message_is_user(): + nodes, tags = _tags_for([_n(0.0, [1, 2], in_=4, out=0, api=0.1)], 2) + pool = SegmentPool() + n0 = nodes[0] + ids, _ = assemble_messages( + n0.request.hash_ids, tags[n0.node_id], pool, _decode, _text + ) + # trailing-user (frozen in tags) + assert pool.materialize([ids[-1]])[0]["role"] == "user" + + +# --- Task 6: block-aligned covered-count ISL build-abort gate --- + +from aiperf.dataset.graph.segment_ir.trie_content import ( # noqa: E402 + TrieISLMismatchError, + assert_covered_isl, +) + + +def test_isl_covered_count_pass_and_fail(): + node = _flatten_requests([_n(0.0, [1, 2], in_=40)], root_scope="tr")[ + 0 + ] # 2 hash blocks, in//16=2 -> covered 2 -> 32 + assert_covered_isl(node, 32, 16) # exact covered-count -> no raise + with pytest.raises(TrieISLMismatchError): + assert_covered_isl(node, 40, 16) # 40 != 32 + + +def test_isl_truncated_hash_does_not_abort(): + node = _flatten_requests([_n(0.0, [1], in_=40)], root_scope="tr")[ + 0 + ] # only 1 hash block, in//16=2 -> covered min(1,2)=1 -> 16 + assert_covered_isl(node, 16, 16) # covered-count 16, NOT 32 -> must NOT raise + with pytest.raises(TrieISLMismatchError): + assert_covered_isl( + node, 32, 16 + ) # demanding (in//bs)*bs would be wrong; 32 != 16 + + +# --- Task 8: end-to-end build_trie_graph switch (interval-order + message-unit) --- + +from aiperf.dataset.graph.adapters.weka.trace_models import ( # noqa: E402 + WekaTrace, +) +from aiperf.dataset.graph.adapters.weka.trie_build import ( # noqa: E402 + ReconCallbacks, + build_trie_graph, +) + +# Collision-free stub callbacks at bs=2 (no tokenizer / corpus build). Each hash +# id decodes to two distinct tokens (h*1000, h*1000+1) so shared-prefix blocks +# yield byte-identical tokens -> identical content-addressed pool ids. +_STUB_CB = ReconCallbacks( + decode_block_tokens=lambda hids: [h * 1000 + k for h in hids for k in range(2)], + sample_partial_tail_tokens=lambda n, s: [0] * n, + decode_tokens_to_text=lambda t: " ".join(map(str, t)), +) + + +def _trace(reqs, block_size=2): + return WekaTrace( + id="t", + models=["m"], + block_size=block_size, + hash_id_scope="local", + requests=reqs, + ) + + +def test_e2e_racers_fork_and_dynamic_join(): + reqs = [ + _n(0.0, [1], in_=2, out=0, api=1.0), + WekaSubagentEntry( + t=1.1, + type="subagent", + agent_id="a", + subagent_type="X", + status="completed", + requests=[_n(1.2, [1, 2], in_=4, out=0, api=2.8)], + models=["m"], + ), + WekaSubagentEntry( + t=1.1, + type="subagent", + agent_id="b", + subagent_type="X", + status="completed", + requests=[_n(1.3, [1, 3], in_=4, out=0, api=3.7)], + models=["m"], + ), + WekaSubagentEntry( + t=5.1, + type="subagent", + agent_id="c", + subagent_type="X", + status="completed", + requests=[_n(5.2, [1, 5], in_=4, out=0, api=1.8)], + models=["m"], + ), + ] + trace = _trace(reqs, block_size=2) + parsed, _ = build_trie_graph(trace, callbacks=_STUB_CB, idle_gap_cap_seconds=None) + # Key nodes by their recorded hashes via the flattener (same walk the + # build ran), not node metadata -- the trie stamp carries only the path. + ids = { + tuple(t.request.hash_ids): t.node_id + for t in _flatten_requests(trace.requests, root_scope=trace.id) + } + preds = lambda h: { # noqa: E731 + e.source for e in parsed.graph.edges if e.target == ids[h] + } + assert ids[(1, 3)] not in preds((1, 2)) and ids[(1, 2)] not in preds( + (1, 3) + ) # racers concurrent + assert preds((1, 5)) == {ids[(1, 2)], ids[(1, 3)]} # dynamic join, no parent turn + + +def test_e2e_shared_prefix_identical_prompt_ids(): + reqs = [ + _n(0.0, [10, 11], in_=4, out=0, api=0.5), + WekaSubagentEntry( + t=1.0, + type="subagent", + agent_id="a", + subagent_type="X", + status="completed", + requests=[_n(1.1, [10, 11, 30], in_=6, out=0, api=0.5)], + models=["m"], + ), + _n(2.0, [10, 11, 20], in_=6, out=0, api=0.5), + ] + trace = _trace(reqs, block_size=2) + parsed, _ = build_trie_graph(trace, callbacks=_STUB_CB, idle_gap_cap_seconds=None) + paths = { + tuple(t.request.hash_ids): parsed.graph.nodes[t.node_id].metadata["trie"][ + "prompt_segment_ids" + ] + for t in _flatten_requests(trace.requests, root_scope=trace.id) + } + lead = paths[(10, 11)] + for h, p in paths.items(): + if h[:2] == (10, 11): + assert ( + p[: len(lead)] == lead + ) # shared [10,11] prefix -> identical leading message ids + + +def test_e2e_coincident_zero_duration_single_edge(): + parsed, _ = build_trie_graph( + _trace([_n(0.0, [1], in_=2), _n(0.0, [2], in_=2)], block_size=2), + callbacks=_STUB_CB, + idle_gap_cap_seconds=None, + ) + assert len([e for e in parsed.graph.edges if e.source != "START"]) <= 1 + + +# --- Task 8 fix: content-parent under-covers, child must tag ALL covered blocks --- + + +def test_under_covering_parent_child_fully_tagged(): + # parent: 3 hash blocks but in=4 (in//2=2) -> covers only 2 blocks (under-covers block index 2). + # child: shares all 3 (lcp=3) and in=6 (covers 3). Child must tag all 3 covered blocks. + reqs = [ + _n( + 0.0, [1, 2, 3], in_=4, out=0, api=0.1 + ), # parent under-covers (2 of 3 blocks) + _n(0.2, [1, 2, 3], in_=6, out=0, api=0.1), # child covers all 3 + ] + nodes, tags = _tags_for(reqs, 2) + child = nodes[1] + covered = min(len(child.request.hash_ids), child.request.input_length // 2) # = 3 + assert len(tags[child.node_id]) == covered + # end-to-end: no ISL abort + parsed, _ = build_trie_graph( + _trace(reqs, block_size=2), callbacks=_STUB_CB, idle_gap_cap_seconds=None + ) + assert parsed is not None + + +def test_over_inheriting_child_clamped_to_covered_count(): + # Symmetric to the under-covering case: the child shares a FULL block prefix + # (lcp=3) with its parent but declares a smaller in (in//2=1 < lcp). inherited + # must clamp to the child's OWN covered count (1), not the lcp (3) -- else it + # would carry 3 tags for a 1-block prompt and abort the build on assert_covered_isl. + reqs = [ + _n(0.0, [1, 2, 3], in_=6, out=0, api=0.1), # parent covers all 3 blocks + _n(0.2, [1, 2, 3], in_=2, out=0, api=0.1), # child covers only 1 (in//2=1) + ] + nodes, tags = _tags_for(reqs, 2) + child = nodes[1] + covered = min(len(child.request.hash_ids), child.request.input_length // 2) # = 1 + assert len(tags[child.node_id]) == covered + # end-to-end: no ISL abort, and the child's single tag is the parent's first + # tag (shared-prefix identity preserved for the blocks the child does cover). + assert tags[child.node_id][0] == tags[nodes[0].node_id][0] + parsed, _ = build_trie_graph( + _trace(reqs, block_size=2), callbacks=_STUB_CB, idle_gap_cap_seconds=None + ) + assert parsed is not None diff --git a/tests/unit/graph/test_weka_trie_route.py b/tests/unit/graph/test_weka_trie_route.py new file mode 100644 index 0000000000..d6e6ec9564 --- /dev/null +++ b/tests/unit/graph/test_weka_trie_route.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Routing gate: :func:`from_weka_trace` builds the segment-trie IR. + +The dependency-only segment-trie builder (:func:`build_trie_graph`) is the +SOLE weka path: the returned :class:`ParsedGraph` carries ONLY :class:`LlmNode` +nodes and :class:`StaticEdge` edges. +""" + +from __future__ import annotations + +from pathlib import Path + +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.models import LlmNode, ParsedGraph, StaticEdge + +FIX_SUBAGENT = Path(__file__).parent / "fixtures" / "weka_subagent.json" + + +def _all_nodes(pg: ParsedGraph) -> list: + """Every node across the top-level graph and any subgraphs.""" + nodes: list = list(pg.graph.nodes.values()) + return nodes + + +def _all_edges(pg: ParsedGraph) -> list: + """Every edge across the top-level graph and any subgraphs.""" + edges: list = list(pg.graph.edges) + return edges + + +def test_weka_parse_routes_to_trie_builder(monkeypatch) -> None: + """Every weka parse -> graph is purely LlmNode + StaticEdge (trie IR). + + Hermetic: ``build_trie_graph``'s default callbacks build a real gpt2-backed + :class:`CorpusContentSynthesizer`. ``HF_HUB_OFFLINE`` / ``TRANSFORMERS_OFFLINE`` + pin the tokenizer load to the local HuggingFace cache so the unit run never + issues a live ``huggingface.co`` HEAD -- the same offline mechanism the other + ``tests/unit/graph`` weka tokenizer paths use. + """ + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + + pg = from_weka_trace(FIX_SUBAGENT) + + nodes = _all_nodes(pg) + edges = _all_edges(pg) + assert nodes, "trie graph must have at least one node" + assert all(isinstance(n, LlmNode) for n in nodes), [type(n).__name__ for n in nodes] + assert all(isinstance(e, StaticEdge) for e in edges), [ + type(e).__name__ for e in edges + ] + # The trie builder produces no subgraphs (no subagent primitives). diff --git a/tests/unit/graph/test_weka_trie_snapshot.py b/tests/unit/graph/test_weka_trie_snapshot.py new file mode 100644 index 0000000000..7bbaa162c9 --- /dev/null +++ b/tests/unit/graph/test_weka_trie_snapshot.py @@ -0,0 +1,315 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``chop_trie_at_tstar`` t* snapshot chop on the trie ``ParsedGraph`` (Task 5). + +The trie graph (Task 3) is the trivial ``LlmNode`` + ``StaticEdge`` realization +of a Weka trace. The t* chop drops every pre-``t*`` turn (warmed, not profiled) +and re-roots each surviving frontier turn from ``START`` at its t*-relative +offset, leaving each surviving node's ``prompt_segment_ids`` UNCHANGED (the full +pre-``t*`` prefix stays in the path so the worker still materializes the exact +resume prompt). ``t* <= 0`` returns the graph unchanged (full replay). + +These tests build a small multi-turn trie graph with the same deterministic stub +callbacks the Task-3 build test uses (no tokenizer / corpus build) and assert the +chop STRUCTURE. +""" + +from __future__ import annotations + +from aiperf.dataset.graph.adapters.weka.trace_models import WekaTrace +from aiperf.dataset.graph.adapters.weka.trie_build import ( + ReconCallbacks, + build_trie_graph, +) +from aiperf.dataset.graph.models import ( + ChannelRequirement, + LlmNode, + ParsedGraph, + StaticEdge, +) +from aiperf.dataset.graph.segment_ir.pool import SegmentPool +from aiperf.timing.snapshot_chop import chop_trie_at_tstar + +BLOCK_SIZE = 64 + + +def _stub_decode_block_tokens(hash_ids: list[int]) -> list[int]: + out: list[int] = [] + for h in hash_ids: + out.extend(range(h * 100, h * 100 + BLOCK_SIZE)) + return out + + +def _stub_partial_tail_tokens(n_tokens: int, seed: str) -> list[int]: + base = sum(ord(c) for c in seed) * 1000 + return list(range(base, base + n_tokens)) + + +def _stub_decode_tokens_to_text(tokens: list[int]) -> str: + return "|".join(str(t) for t in tokens) + + +_STUB_CALLBACKS = ReconCallbacks( + decode_block_tokens=_stub_decode_block_tokens, + sample_partial_tail_tokens=_stub_partial_tail_tokens, + decode_tokens_to_text=_stub_decode_tokens_to_text, +) + + +def _build(trace: WekaTrace) -> tuple[ParsedGraph, SegmentPool]: + return build_trie_graph(trace, callbacks=_STUB_CALLBACKS) + + +def _llm_nodes(graph: ParsedGraph) -> dict[str, LlmNode]: + return {nid: n for nid, n in graph.graph.nodes.items() if isinstance(n, LlmNode)} + + +def _edges(graph: ParsedGraph) -> list[StaticEdge]: + return [e for e in graph.graph.edges if isinstance(e, StaticEdge)] + + +def _edge(edges: list[StaticEdge], source: str, target: str) -> StaticEdge | None: + for e in edges: + if e.source == source and e.target == target: + return e + return None + + +def _n(t: float, hashes: list[int], in_len: int, out: int, api_time: float, + model: str = "M", type_: str = "n") -> dict: # fmt: skip + return { + "t": t, + "type": type_, + "model": model, + "in": in_len, + "out": out, + "hash_ids": hashes, + "api_time": api_time, + } + + +def _subagent(t: float, agent_id: str, requests: list[dict], status: str, + duration_ms: int | None = None) -> dict: # fmt: skip + return { + "t": t, + "type": "subagent", + "agent_id": agent_id, + "subagent_type": "Explore", + "status": status, + "duration_ms": duration_ms, + "requests": requests, + "models": ["M"], + } + + +def _trace(requests: list[dict]) -> WekaTrace: + return WekaTrace.model_validate( + { + "id": "trace_0", + "models": ["M"], + "block_size": BLOCK_SIZE, + "hash_id_scope": "local", + "requests": requests, + } + ) + + +def _join_trace() -> WekaTrace: + """Parent p0 spawns subagents s1,s2; resume turn p1 AND-fans-in on both. + + Arrival offsets: p0=0 s, s1_inner=1 s, s2_inner=1.2 s, p1=4 s. Under + interval-order timing p1 declares TWO ``inputs`` (``s1_inner_out``, + ``s2_inner_out``); the content-parent p0 (end 0.5) is transitively covered + (p0 -> s1_inner -> p1) and dropped from p1's finished-before frontier. Used + to exercise the chop's input-rescoping when some of p1's predecessors are + dropped. + """ + s1 = _subagent( + t=1.0, + agent_id="agent_1", + status="completed", + duration_ms=2000, + requests=[_n(t=1.0, hashes=[50, 51], in_len=128, out=16, api_time=1.5)], + ) + s2 = _subagent( + t=1.0, + agent_id="agent_2", + status="completed", + duration_ms=2000, + requests=[_n(t=1.2, hashes=[60, 61], in_len=128, out=16, api_time=1.5)], + ) + return _trace( + [ + _n(t=0.0, hashes=[1, 2], in_len=128, out=64, api_time=0.5), + s1, + s2, + _n(t=4.0, hashes=[1, 2, 3], in_len=192, out=32, api_time=1.0), + ] + ) + + +def _four_turn_trace() -> WekaTrace: + """Linear 4-turn conversation; each turn waits for its predecessor to finish. + + Turn k starts at t=k*10, runs api_time=1, so the inter-turn gap is + ``k*10 - ((k-1)*10 + 1) = 9`` seconds. Arrival offsets are 0, 10, 20, 30 s. + """ + return _trace( + [ + _n(t=0.0, hashes=[1, 2], in_len=128, out=64, api_time=1.0), + _n(t=10.0, hashes=[1, 2, 3], in_len=192, out=32, api_time=1.0), + _n(t=20.0, hashes=[1, 2, 3, 4], in_len=256, out=32, api_time=1.0), + _n(t=30.0, hashes=[1, 2, 3, 4, 5], in_len=320, out=32, api_time=1.0), + ] + ) + + +def _ordered_node_ids(graph: ParsedGraph) -> list[str]: + """Node ids in ascending recorded arrival offset (turn order).""" + nodes = _llm_nodes(graph) + return sorted(nodes, key=lambda nid: nodes[nid].arrival_offset_us) + + +def test_chop_drops_pre_tstar_and_reroots_frontier() -> None: + """t* between turn1 and turn2 drops turns 0,1; re-roots turn2 from START. + + With offsets [0, 10, 20, 30] s and t* = 25 s, turns 0,1,2 are pre-t* (dropped) + and turn3 (offset 30 s) survives. turn3 re-roots from START with + ``min_start_delay_us = (30 - 25) * 1e6``. + """ + graph, _ = _build(_four_turn_trace()) + ids = _ordered_node_ids(graph) + t_star_us = int(25.0 * 1e6) + + chopped = chop_trie_at_tstar(graph, t_star_us) + + surviving = _llm_nodes(chopped) + # Only the offset-30 turn survives. + assert set(surviving) == {ids[3]} + # Re-rooted from START with the t*-relative arrival offset. + edges = _edges(chopped) + start_edge = _edge(edges, "START", ids[3]) + assert start_edge is not None + assert start_edge.min_start_delay_us == (30.0 - 25.0) * 1e6 + # No dangling edges to the dropped predecessors. + assert all(e.source in surviving or e.source == "START" for e in edges) + assert all(e.target in surviving for e in edges) + + +def test_chop_keeps_surviving_inter_turn_edge_delay() -> None: + """Two surviving turns keep their recorded inter-turn ``delay_after_predecessor_us``. + + t* = 5 s drops only turn 0 (offset 0). Turns 1,2,3 survive. turn1 (offset 10) + re-roots from START at offset 5 s. The surviving turn1->turn2 and turn2->turn3 + edges keep their recorded 9 s ``delay_after_predecessor_us``. + """ + graph, _ = _build(_four_turn_trace()) + ids = _ordered_node_ids(graph) + t_star_us = int(5.0 * 1e6) + + chopped = chop_trie_at_tstar(graph, t_star_us) + + surviving = _llm_nodes(chopped) + assert set(surviving) == {ids[1], ids[2], ids[3]} + + edges = _edges(chopped) + # turn1 re-rooted from START at t*-relative offset (10 - 5 = 5 s). + start_edge = _edge(edges, "START", ids[1]) + assert start_edge is not None + assert start_edge.min_start_delay_us == (10.0 - 5.0) * 1e6 + # turn1 is no longer rooted from its dropped predecessor (turn0). + assert _edge(edges, ids[0], ids[1]) is None + + # Surviving inter-turn edges keep their recorded 9 s gap. + e12 = _edge(edges, ids[1], ids[2]) + e23 = _edge(edges, ids[2], ids[3]) + assert e12 is not None and e12.delay_after_predecessor_us == 9.0 * 1e6 + assert e23 is not None and e23.delay_after_predecessor_us == 9.0 * 1e6 + # Surviving inter-turn edges are NOT re-rooted from START. + assert _edge(edges, "START", ids[2]) is None + assert _edge(edges, "START", ids[3]) is None + + +def test_chop_leaves_prompt_segment_ids_unchanged() -> None: + """A surviving re-rooted node keeps its FULL pre-t* prompt path verbatim.""" + graph, _ = _build(_four_turn_trace()) + ids = _ordered_node_ids(graph) + before = _llm_nodes(graph)[ids[3]].metadata["trie"]["prompt_segment_ids"] + + chopped = chop_trie_at_tstar(graph, int(25.0 * 1e6)) + + after = _llm_nodes(chopped)[ids[3]].metadata["trie"]["prompt_segment_ids"] + assert after == before + assert len(after) > 1 # full lineage prefix, not truncated to the resume turn + + +def test_chop_tstar_zero_returns_graph_unchanged() -> None: + """``t* <= 0`` is full replay: the graph is returned identically.""" + graph, _ = _build(_four_turn_trace()) + + assert chop_trie_at_tstar(graph, 0) is graph + assert chop_trie_at_tstar(graph, -1) is graph + + +def test_chop_drops_dropped_predecessor_input() -> None: + """A 2-pred node losing ONE pred to the chop keeps only the surviving input. + + t* = 1.1 s drops p0 (offset 0) and s1_inner (offset 1.0) but keeps s2_inner + (1.2 s) and the resume turn p1 (4 s). p1's ``inputs`` must drop the dropped + ``s1_inner_out`` requirement and keep only ``s2_inner_out`` -- the dropped + channel is never written post-chop, so a stale requirement would DEADLOCK + ``await_inputs``. p1 still has a surviving predecessor (s2_inner), so it is + NOT re-rooted from START. + """ + graph, _ = _build(_join_trace()) + nodes = _llm_nodes(graph) + by_offset = {n.arrival_offset_us: nid for nid, n in nodes.items()} + s1_inner = by_offset[int(1.0 * 1e6)] + s2_inner = by_offset[int(1.2 * 1e6)] + p1 = by_offset[int(4.0 * 1e6)] + # Precondition: p1 fans in on its interval-order frontier -- BOTH subagent + # last turns (two predecessors) -- before the chop. The content-parent p0 is + # transitively covered and is NOT a direct predecessor under interval-order. + assert set(nodes[p1].inputs) == { + ChannelRequirement(channel=f"{s1_inner}_out", count=1), + ChannelRequirement(channel=f"{s2_inner}_out", count=1), + } + + chopped = chop_trie_at_tstar(graph, int(1.1 * 1e6)) + + survivors = _llm_nodes(chopped) + assert p1 in survivors and s2_inner in survivors + # p1 keeps ONLY the surviving predecessor's input requirement. + assert survivors[p1].inputs == [ + ChannelRequirement(channel=f"{s2_inner}_out", count=1) + ] + # The surviving inter-turn edge s2_inner -> p1 remains; p1 is not re-rooted. + edges = _edges(chopped) + assert _edge(edges, s2_inner, p1) is not None + assert _edge(edges, "START", p1) is None + + +def test_chop_dropping_all_predecessors_reroots_from_start_with_no_inputs() -> None: + """A node losing ALL preds to the chop re-roots from START with empty inputs. + + t* = 2.0 s drops p0, s1_inner, s2_inner (offsets 0, 1.0, 1.2) but keeps p1 + (offset 4 s). p1 lost BOTH predecessors, so it re-roots from START and its + ``inputs`` collapse to empty -- otherwise ``await_inputs`` would block on + channels no surviving node ever writes. + """ + graph, _ = _build(_join_trace()) + nodes = _llm_nodes(graph) + by_offset = {n.arrival_offset_us: nid for nid, n in nodes.items()} + p1 = by_offset[int(4.0 * 1e6)] + + chopped = chop_trie_at_tstar(graph, int(2.0 * 1e6)) + + survivors = _llm_nodes(chopped) + assert set(survivors) == {p1} + # No predecessor survived: inputs collapse to empty. + assert survivors[p1].inputs == [] + # Re-rooted from START at its t*-relative offset (4 - 2 = 2 s). + edges = _edges(chopped) + start_edge = _edge(edges, "START", p1) + assert start_edge is not None + assert start_edge.min_start_delay_us == (4.0 - 2.0) * 1e6 diff --git a/tests/unit/graph/test_weka_unified_dispatch.py b/tests/unit/graph/test_weka_unified_dispatch.py new file mode 100644 index 0000000000..d9c815bddf --- /dev/null +++ b/tests/unit/graph/test_weka_unified_dispatch.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import orjson +import pytest + +from aiperf.dataset.graph.adapters.weka import trace_parallel as parallel +from aiperf.dataset.graph.adapters.weka.trace import EmptyWekaTraceError + +FIX_MIN = Path(__file__).parent / "fixtures" / "weka_min.json" + +PARSE_KWARGS: dict[str, Any] = { + "tag": "from-weka-trace", + "idle_gap_cap_seconds": None, + "content_root_seed": 42, + "content_tokenizer": None, + "prompt_corpus": None, + "max_osl": None, +} + + +def _row_items(count: int) -> list[parallel._WorkItem]: + raw = orjson.loads(FIX_MIN.read_bytes()) + items = [] + for i in range(count): + row = dict(raw) + row["id"] = f"trace-{i}" + items.append(parallel._WorkItem(source=f"src#{i}", path=None, row=row)) + return items + + +def test_parse_items_serial_below_threshold_merges_all_traces() -> None: + parsed = parallel.parse_items( + _row_items(2), + source_label="src", + threshold=8, + parse_kwargs=dict(PARSE_KWARGS), + ) + assert [t.id for t in parsed.traces] == ["trace-0", "trace-1"] + assert set(parsed.graphs) == {"trace-0", "trace-1"} + + +def test_parse_items_file_and_row_items_produce_identical_traces() -> None: + raw = orjson.loads(FIX_MIN.read_bytes()) + by_row = parallel.parse_items( + [parallel._WorkItem(source="r", path=None, row=dict(raw))], + source_label="r", + threshold=8, + parse_kwargs=dict(PARSE_KWARGS), + ) + by_file = parallel.parse_items( + [parallel._WorkItem(source=str(FIX_MIN), path=str(FIX_MIN), row=None)], + source_label=str(FIX_MIN), + threshold=8, + parse_kwargs=dict(PARSE_KWARGS), + ) + assert [t.id for t in by_row.traces] == [t.id for t in by_file.traces] + assert by_row.graph.nodes.keys() == by_file.graph.nodes.keys() + + +def test_parse_items_above_threshold_routes_through_streaming_pool( + monkeypatch, +) -> None: + seen_workers: list[int] = [] + + def fake_pool(worker_fn, work_items, *, workers, **_kwargs): # noqa: ANN001 + seen_workers.append(workers) + for task in work_items: + yield worker_fn(task) + + monkeypatch.setattr(parallel, "_run_pool_streaming", fake_pool) + items = _row_items(3) + parsed = parallel.parse_items( + items, + source_label="src", + item_count=len(items), + threshold=2, + workers=8, + parse_kwargs=dict(PARSE_KWARGS), + ) + assert len(parsed.traces) == 3 + # Known item_count caps the worker fan-out (previously only the directory + # path applied this cap; row sources spawned the full configured count). + assert seen_workers == [3] + + +def test_parse_items_zero_items_raises_empty_error() -> None: + with pytest.raises(EmptyWekaTraceError, match="src-label"): + parallel.parse_items( + [], + source_label="src-label", + threshold=8, + parse_kwargs=dict(PARSE_KWARGS), + ) + + +def test_iter_item_segment_payloads_first_yield_consumes_only_prefetch_window( + monkeypatch, +) -> None: + raw = orjson.loads(FIX_MIN.read_bytes()) + consumed: list[int] = [] + + def items(): + for index in range(10): + consumed.append(index) + row = dict(raw) + row["id"] = f"trace-{index}" + yield parallel._WorkItem(source=f"src#{index}", path=None, row=row) + + def fake_pool(worker_fn, work_items, **_kwargs): # noqa: ANN001 + assert consumed == [0, 1, 2] + iterator = iter(work_items) + yield worker_fn(next(iterator)) + + monkeypatch.setattr(parallel, "_run_pool_streaming", fake_pool) + payloads_iter = parallel.iter_item_segment_payloads( + items(), + source_label="src", + threshold=2, + workers=1, + parse_kwargs=dict(PARSE_KWARGS), + ) + first = next(payloads_iter) + assert first.trace_id == "trace-0" + assert consumed == [0, 1, 2] + + +def test_workers_never_reenter_public_route(monkeypatch) -> None: + # The pool worker core must call the per-item parsers directly, not the + # public dispatcher (which re-runs HF-id and is_dir detection per file). + from aiperf.dataset.graph.adapters.weka import trace as weka_trace + + def boom(*_a: Any, **_k: Any): + raise AssertionError("worker re-entered from_weka_trace") + + monkeypatch.setattr(weka_trace, "from_weka_trace", boom) + item = parallel._WorkItem(source=str(FIX_MIN), path=str(FIX_MIN), row=None) + blob = parallel._parse_item_to_msgpack((item, dict(PARSE_KWARGS))) + assert isinstance(blob, bytes) and blob diff --git a/tests/unit/graph/test_worker_bytes_path.py b/tests/unit/graph/test_worker_bytes_path.py new file mode 100644 index 0000000000..f4e95cec22 --- /dev/null +++ b/tests/unit/graph/test_worker_bytes_path.py @@ -0,0 +1,485 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Worker graph-IR bytes path: pre-serialized body parity + record correctness. + +The interned unified store is the sole trie store shape: on a graph credit the +worker's segment reader is a :class:`GraphSegmentUnifiedClient` and, when no +content mutation is needed (``cache_bust == NONE``), the worker builds the +request body ONCE from content-pool slices via +:func:`materialize_graph_request_unified_bytes` and sends it verbatim as a +``Turn.raw_payload_bytes`` -- no per-segment ``orjson.loads`` and no +``orjson.dumps`` of the messages array. + +Three gates this test protects: + +1. **Body parity** -- the bytes the bytes path builds, when ``orjson.loads``-ed, + equal the dict the dict path (``materialize_graph_request_unified`` + + ``apply_run_level_payload_options``) produces for the same node + same + endpoint settings (stream / extra / server-token-count / warmup all folded + identically). This is the core correctness proof: a pre-serialized body can't + be mutated after build, so every outer field must be folded in at build time. +2. **payload_bytes record correctness** -- ``inference_client`` records a bytes + body VERBATIM (not ``orjson.dumps`` of it, which would corrupt the recorded + JSON into a string and break ISL / raw-export); a dict ``raw_payload`` still + records ``orjson.dumps(dict)``. +3. **cache-bust fallback** -- a cache-bust run mutates message CONTENT (prepends a + marker to the first user message), which a pre-serialized body cannot do; so + the dict path is the only one that carries the marker, and the worker gates the + bytes path on ``cache_bust == NONE``. + +Builds a 2-trace weka trie corpus (merged segment pool drained straight into a +``GraphSegmentUnifiedBackingStore`` with the node's interned manifest written +per-test). +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import orjson +import pytest + +from aiperf.common.enums import ( + CacheBustTarget, + ModelSelectionStrategy, +) +from aiperf.common.models.dataset_models import Turn +from aiperf.common.models.model_endpoint_info import ( + EndpointInfo, + ModelEndpointInfo, + ModelInfo, + ModelListInfo, +) +from aiperf.common.models.record_models import RequestRecord +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.segment_ir.pool import SegmentPool +from aiperf.dataset.graph.segment_ir.store_builder import ( + _prompt_segment_ids, + _trie_llm_nodes, +) +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) +from aiperf.graph.worker_materialize import ( + apply_run_level_payload_options, + encode_overrides_inner, + materialize_graph_request_unified, + materialize_graph_request_unified_bytes, + stamp_cache_bust_marker, +) +from aiperf.plugin.enums import EndpointType, TransportType +from aiperf.workers.inference_client import InferenceClient + +FIXTURES = Path(__file__).parent / "fixtures" + +_TRACE_ID = "t-1#0" +_NODE_ORDINAL = 7 + + +def _merge_pool(*pools: SegmentPool) -> SegmentPool: + """Union content-addressed pools (ids collide iff content matches).""" + merged = SegmentPool() + for pool in pools: + merged._by_id.update(pool._by_id) + return merged + + +@pytest.fixture +def two_trace_corpus(monkeypatch): + """A 2-trace weka trie corpus: merged pool + one node's known prompt path.""" + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + + pg_a = from_weka_trace(FIXTURES / "weka_subagent.json") + pg_b = from_weka_trace(FIXTURES / "weka_min.json") + assert pg_a.segment_pool is not None + assert pg_b.segment_pool is not None + pool = _merge_pool(pg_a.segment_pool, pg_b.segment_pool) + + known_path: list[str] | None = None + for trace in pg_a.traces: + for node in _trie_llm_nodes(pg_a, trace).values(): + path = _prompt_segment_ids(node) + if path and (known_path is None or len(path) > len(known_path)): + known_path = path + assert known_path, "fixture must yield a node with a prompt_segment_ids path" + return pool, known_path + + +async def _build_unified_client( + tmp_path: Path, + pool: SegmentPool, + known_path: list[str], + benchmark_id: str, + *, + dispatch_overrides: dict[str, Any], + stream: bool = True, +) -> GraphSegmentUnifiedClient: + """Drain the pool + one interned node manifest into a unified store; open it. + + Mirrors ``build_unified_trie_store_interned``: every pool segment is + ``put_segment``'d (assigning dense int handles), and the node's hex + ``known_path`` is resolved to those handles for the interned manifest. + """ + store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id=benchmark_id + ) + for segment in pool._by_id.values(): + store.put_segment(segment.id, segment.role, segment.content) + handles = [store.segment_handle(sid) for sid in known_path] + assert all(h is not None for h in handles) + store.add_node_manifest_interned( + _TRACE_ID, _NODE_ORDINAL, "profiling", handles, dispatch_overrides, stream + ) + await store.finalize() + return GraphSegmentUnifiedClient(tmp_path, benchmark_id).open() + + +def _endpoint( + *, + streaming: bool = True, + extra: list[tuple[str, Any]] | None = None, + use_server_token_count: bool = False, + use_legacy_max_tokens: bool = False, + cache_bust: CacheBustTarget = CacheBustTarget.NONE, +) -> EndpointInfo: + return EndpointInfo( + type=EndpointType.CHAT, + base_url="http://localhost:8000/v1/chat", + streaming=streaming, + extra=extra or [], + use_server_token_count=use_server_token_count, + use_legacy_max_tokens=use_legacy_max_tokens, + cache_bust=cache_bust, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "streaming,extra,use_server_token_count,use_legacy_max_tokens,phase_variant", + [ + pytest.param(True, [], False, False, "profiling", id="stream_modern"), + pytest.param(False, [], False, True, "profiling", id="nostream_legacy"), + pytest.param(True, [], True, False, "profiling", id="stream_usage"), + pytest.param( + True, [("temperature", 0.0), ("top_p", 0.9)], False, False, "profiling", + id="stream_extra", + ), + pytest.param(True, [], True, False, "warmup", id="warmup_usage"), + pytest.param(False, [("seed", 5)], False, True, "warmup", id="warmup_extra_legacy"), + ], +) # fmt: skip +async def test_bytes_body_matches_dict_path( + tmp_path, + two_trace_corpus, + streaming, + extra, + use_server_token_count, + use_legacy_max_tokens, + phase_variant, +): + """``materialize_graph_request_unified_bytes`` bytes round-trip to the dict-path dict. + + The dict path = ``materialize_graph_request_unified`` + + ``apply_run_level_payload_options`` (cache-bust off, so no + ``stamp_cache_bust_marker`` mutation). The bytes path must produce a body + whose ``orjson.loads`` equals that dict EXACTLY -- proving every outer field + (token mapping, stream override, extra, include_usage, warmup cap) is folded + into the pre-serialized tail identically. + """ + pool, known_path = two_trace_corpus + client = await _build_unified_client( + tmp_path, + pool, + known_path, + "bench", + dispatch_overrides={"model": "m", "max_output_tokens": 7}, + stream=streaming, + ) + endpoint = _endpoint( + streaming=streaming, + extra=extra, + use_server_token_count=use_server_token_count, + use_legacy_max_tokens=use_legacy_max_tokens, + ) + + try: + dict_payload = materialize_graph_request_unified( + client, + _TRACE_ID, + _NODE_ORDINAL, + phase_variant, + use_legacy_max_tokens=use_legacy_max_tokens, + ) + assert dict_payload is not None + apply_run_level_payload_options(dict_payload, endpoint) + + built = materialize_graph_request_unified_bytes( + client, + _TRACE_ID, + _NODE_ORDINAL, + phase_variant, + use_legacy_max_tokens=use_legacy_max_tokens, + endpoint=endpoint, + ) + finally: + client.close() + + assert built is not None + body, model, effective_stream = built + # Parsed-JSON parity is the correctness-relevant property (not raw-byte order). + assert orjson.loads(body) == dict_payload + # The per-node model is surfaced so the worker can stamp Turn.model = the same + # value the dict path's payload.get("model") would. + assert model == dict_payload.get("model") + # The effective wire stream is surfaced so the worker carries the recorded + # per-node mode onto RequestInfo; it matches the FINAL stamped dict value. + assert effective_stream == dict_payload["stream"] + + +@pytest.mark.asyncio +async def test_bytes_path_returns_none_without_handles(tmp_path, two_trace_corpus): + """A node manifest WITHOUT ``handles`` is a miss on both unified paths (None).""" + pool, known_path = two_trace_corpus + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id="bench") + for segment in pool._by_id.values(): + store.put_segment(segment.id, segment.role, segment.content) + # A handle-less envelope (no interned path) -- the unified fns must not + # guess; they return None and the worker reports a genuine miss. + store.add_node_manifest( + _TRACE_ID, + _NODE_ORDINAL, + "profiling", + orjson.dumps({"dispatch_overrides": {"model": "m"}, "stream": True}), + ) + await store.finalize() + + with GraphSegmentUnifiedClient(tmp_path, "bench").open() as client: + built = materialize_graph_request_unified_bytes( + client, + _TRACE_ID, + _NODE_ORDINAL, + "profiling", + endpoint=_endpoint(), + ) + payload = materialize_graph_request_unified( + client, _TRACE_ID, _NODE_ORDINAL, "profiling" + ) + assert built is None + assert payload is None + + +@pytest.mark.asyncio +async def test_cache_bust_only_dict_path_can_inject_marker(tmp_path, two_trace_corpus): + """Cache-bust mutates message CONTENT -> only the dict path can carry it. + + The bytes body is built once and cannot be mutated, which is exactly WHY the + worker gates the bytes path on ``cache_bust == NONE``. This proves the dict + path injects the marker into the first user message and the bytes body does + NOT contain it (so taking the bytes path under cache-bust would silently drop + the marker -- the bug the gate prevents). + """ + pool, known_path = two_trace_corpus + client = await _build_unified_client( + tmp_path, + pool, + known_path, + "bench", + dispatch_overrides={"model": "m", "max_output_tokens": 7}, + ) + endpoint = _endpoint(cache_bust=CacheBustTarget.FIRST_TURN_PREFIX) + try: + dict_payload = materialize_graph_request_unified( + client, _TRACE_ID, _NODE_ORDINAL, "profiling" + ) + assert dict_payload is not None + apply_run_level_payload_options(dict_payload, endpoint) + stamp_cache_bust_marker( + dict_payload, + benchmark_id="bench", + trace_instance_id=_TRACE_ID, + target=endpoint.cache_bust, + ) + built = materialize_graph_request_unified_bytes( + client, + _TRACE_ID, + _NODE_ORDINAL, + "profiling", + endpoint=endpoint, + ) + finally: + client.close() + + first_user = next(m for m in dict_payload["messages"] if m["role"] == "user") + assert first_user["content"].startswith("[rid:"), "dict path must inject the marker" + # The bytes body carries the ORIGINAL content (no marker) -- proving it cannot + # represent the cache-bust mutation, hence the worker's cache_bust==NONE gate. + assert built is not None + body, _, _ = built + assert b"[rid:" not in body + + +@pytest.mark.asyncio +async def test_bytes_path_preserves_per_node_model_in_wire_body( + tmp_path, two_trace_corpus +): + """A node whose model != endpoint primary keeps its OWN model in the wire body. + + The bytes path surfaces the node's ``dispatch_overrides["model"]`` (a real + multi-model weka pattern, e.g. Haiku WebFetch sidecars under a non-Haiku + primary) and folds it into the body bytes, matching the dict path's + ``payload.get("model")``. The worker does NOT stamp ``Turn.model`` with it: + ``record.model_name`` falls back to the run ``--model`` in + ``_finalize_request_record`` so tokenizer selection behaves like plain + aiperf (recorded deployment ids are usually not resolvable tokenizer + repos); the recorded model reaches the server via the verbatim body only. + """ + pool, known_path = two_trace_corpus + # Endpoint primary is "test-model"; the node overrides to a DIFFERENT model. + per_node_model = "anthropic/claude-haiku" + client = await _build_unified_client( + tmp_path, + pool, + known_path, + "bench", + dispatch_overrides={"model": per_node_model, "max_output_tokens": 7}, + ) + endpoint = _endpoint() + try: + built = materialize_graph_request_unified_bytes( + client, + _TRACE_ID, + _NODE_ORDINAL, + "profiling", + endpoint=endpoint, + ) + finally: + client.close() + + assert built is not None + body, model, _ = built + # The model is surfaced for dict-path parity checks. + assert model == per_node_model + # The wire body carries it (folded into the overrides tail). + assert orjson.loads(body)["model"] == per_node_model + + +def test_encode_overrides_inner(): + assert encode_overrides_inner({"max_tokens": 256}) == b'"max_tokens":256' + assert encode_overrides_inner({}) == b"" + assert encode_overrides_inner(None) == b"" + + +# --- payload_bytes record correctness (inference_client selection) ---------- + + +@pytest.fixture +def _http_transport_entry(): + entry = MagicMock() + entry.name = TransportType.HTTP.value + entry.metadata = {"url_schemes": ["http", "https"]} + return entry + + +@pytest.fixture +def inference_client(_http_transport_entry): + model_endpoint = ModelEndpointInfo( + models=ModelListInfo( + models=[ModelInfo(name="test-model")], + model_selection_strategy=ModelSelectionStrategy.ROUND_ROBIN, + ), + endpoint=EndpointInfo( + type=EndpointType.CHAT, base_url="http://localhost:8000/v1/chat" + ), + ) + mock_transport = MagicMock() + mock_endpoint = MagicMock() + mock_endpoint.get_endpoint_headers.return_value = {} + mock_endpoint.get_endpoint_params.return_value = {} + mock_endpoint.format_payload.return_value = {"rewritten": True} + + def mock_get_class(protocol, name): + if protocol == "endpoint": + return lambda **kwargs: mock_endpoint + if protocol == "transport": + return lambda **kwargs: mock_transport + raise ValueError(f"Unknown protocol: {protocol}") + + with ( + patch( + "aiperf.workers.inference_client.plugins.get_class", + side_effect=mock_get_class, + ), + patch( + "aiperf.workers.inference_client.plugins.list_entries", + return_value=[_http_transport_entry], + ), + ): + return InferenceClient( + model_endpoint=model_endpoint, service_id="test-service-id" + ) + + +@pytest.mark.asyncio +async def test_inference_client_records_bytes_payload_verbatim( + inference_client, sample_request_info +): + """A ``raw_payload_bytes`` turn records ``payload_bytes`` VERBATIM (no dumps).""" + body = b'{"messages":[]}' + request_info = sample_request_info + request_info.turns = [Turn(role="user", raw_payload_bytes=body)] + inference_client.transport.send_request = AsyncMock( + return_value=RequestRecord(request_info=request_info) + ) + + await inference_client.send_request(request_info) + + inference_client.endpoint.format_payload.assert_not_called() + # Recorded VERBATIM: the exact bytes, which round-trip back to valid JSON. + assert request_info.payload_bytes == body + assert orjson.loads(request_info.payload_bytes) == {"messages": []} + call_args = inference_client.transport.send_request.call_args + assert call_args.kwargs["payload"] == body + + +@pytest.mark.asyncio +async def test_inference_client_records_dict_payload_as_dumps( + inference_client, sample_request_info +): + """A dict ``raw_payload`` turn still records ``orjson.dumps(dict)`` (unchanged).""" + payload = {"messages": [{"role": "user", "content": "hi"}], "stream": True} + request_info = sample_request_info + request_info.turns = [Turn(role="user", raw_payload=payload)] + inference_client.transport.send_request = AsyncMock( + return_value=RequestRecord(request_info=request_info) + ) + + await inference_client.send_request(request_info) + + inference_client.endpoint.format_payload.assert_not_called() + assert request_info.payload_bytes == orjson.dumps(payload) + call_args = inference_client.transport.send_request.call_args + assert call_args.kwargs["payload"] == payload + + +@pytest.mark.asyncio +async def test_graph_record_model_name_falls_back_to_run_model( + inference_client, sample_request_info +): + """A graph turn carries no ``Turn.model``, so ``record.model_name`` resolves + to the run ``--model`` (primary), never the recorded body model — recorded + deployment ids (e.g. ``dynamo/org/model-fp8``) are not tokenizer repos. + """ + body = b'{"messages":[],"model":"dynamo/org/model-fp8"}' + request_info = sample_request_info + request_info.turns = [Turn(role="user", raw_payload_bytes=body)] + inference_client.transport.send_request = AsyncMock( + return_value=RequestRecord(request_info=request_info) + ) + + record = await inference_client.send_request(request_info) + + assert record.model_name == "test-model" diff --git a/tests/unit/graph/test_worker_graph_branch.py b/tests/unit/graph/test_worker_graph_branch.py new file mode 100644 index 0000000000..1711d57922 --- /dev/null +++ b/tests/unit/graph/test_worker_graph_branch.py @@ -0,0 +1,256 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Worker routing of a graph credit to the mmap materializer (D1). + +A credit carrying ``trace_id``/``node_ordinal`` must NOT take the linear +session-cache path: the worker opens the unified store reader from the same +(base_path, benchmark_id) the build wrote to, materializes the request, and sends +the materialized payload verbatim through the existing inference client. +""" + +import time +from pathlib import Path +from unittest.mock import AsyncMock + +import orjson +import pytest + +from aiperf.common.enums import CreditPhase +from aiperf.common.models import RequestRecord +from aiperf.common.models.dataset_models import GraphSegmentClientMetadata +from aiperf.credit.structs import Credit, CreditContext +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.segment_ir.store_builder import ( + build_unified_trie_store_interned, +) +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) +from aiperf.graph.worker_materialize import materialize_graph_request_unified +from aiperf.workers.worker import Worker +from tests.harness.fake_tokenizer import FakeTokenizer + +FIX_MIN = Path(__file__).parent / "fixtures" / "weka_min.json" + + +def _graph_client_metadata( + tmp_path: Path, benchmark_id: str +) -> GraphSegmentClientMetadata: + """Broadcast-shaped store location the worker opens (Task 5: no env fallback).""" + return GraphSegmentClientMetadata( + store_base_path=tmp_path, + benchmark_id=benchmark_id, + sidecar_path=tmp_path + / f"aiperf_graph_meta_{benchmark_id}" + / "graph_meta.msgpack", + ) + + +@pytest.fixture +async def mock_worker( + benchmark_run, fake_tokenizer: FakeTokenizer, skip_service_registration +): + worker = Worker(run=benchmark_run, service_id="mock-service-id") + await worker.initialize() + await worker.start() + yield worker + await worker.stop() + + +async def _build_store_for_worker(worker: Worker, tmp_path: Path): + """Build the interned unified trie store the worker will open. + + Mirrors ``dataset_manager._build_graph_trie_stores``: the content-addressed + pool AND every node's ``prompt_segment_ids`` manifest land in the ONE + unified store. The worker opens it from the graph-typed dataset broadcast + (``GraphSegmentClientMetadata``), which we set directly here. + """ + parsed = from_weka_trace(str(FIX_MIN)) + store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id=worker.run.benchmark_id + ) + addr = await build_unified_trie_store_interned(parsed, store) + worker._graph_client_metadata = _graph_client_metadata( + tmp_path, worker.run.benchmark_id + ) + return parsed, addr + + +@pytest.mark.asyncio +async def test_graph_credit_routes_to_materializer(mock_worker, tmp_path, monkeypatch): + """A graph credit sends the materialized payload via the inference client.""" + parsed, addr = await _build_store_for_worker(mock_worker, tmp_path) + t0 = parsed.traces[0].id + leaf_ordinal = addr[t0]["trace_03_n3:2"] + + captured = {} + + async def fake_send_request(request_info, first_token_callback=None): + captured["request_info"] = request_info + return RequestRecord( + request_info=request_info, + timestamp_ns=time.time_ns(), + start_perf_ns=time.perf_counter_ns(), + end_perf_ns=time.perf_counter_ns(), + ) + + monkeypatch.setattr( + mock_worker.inference_client, + "send_request", + AsyncMock(side_effect=fake_send_request), + ) + mock_worker._send_inference_result_message = AsyncMock() + + credit = Credit( + id=7, + phase=CreditPhase.PROFILING, + conversation_id=t0, + x_correlation_id="x-7", + turn_index=0, + num_turns=1, + issued_at_ns=time.time_ns(), + trace_id=t0, + node_ordinal=leaf_ordinal, + phase_variant="profiling", + ) + ctx = CreditContext(credit=credit, drop_perf_ns=time.perf_counter_ns()) + + await mock_worker._process_credit(ctx) + + # Inference client was called exactly once with the materialized payload. + mock_worker.inference_client.send_request.assert_awaited_once() + request_info = captured["request_info"] + assert len(request_info.turns) == 1 + # The interned unified store default: the graph credit takes the BYTES + # path -- raw_payload is None and raw_payload_bytes carries the + # materialized body verbatim. + raw_payload_bytes = request_info.turns[0].raw_payload_bytes + assert raw_payload_bytes is not None + wire = orjson.loads(raw_payload_bytes) + + # Reconstruct the expected materialized payload through the same unified + # reader the worker opened; the dict-path materializer walks the same + # interned handles, so the messages match the bytes path's wire body. + expected_client = GraphSegmentUnifiedClient( + base_path=tmp_path, benchmark_id=mock_worker.run.benchmark_id + ).open() + expected = materialize_graph_request_unified( + expected_client, + t0, + leaf_ordinal, + "profiling", + use_legacy_max_tokens=mock_worker.model_endpoint.endpoint.use_legacy_max_tokens, + ) + expected_client.close() + assert wire["messages"] == expected["messages"] + assert wire["messages"] + # No linear session was created for the graph credit. + assert mock_worker.session_manager.get("x-7") is None + + +@pytest.mark.asyncio +async def test_missing_graph_envelope_records_error_no_send( + mock_worker, tmp_path, monkeypatch +): + """A graph credit with no stored delta records an error and sends nothing.""" + parsed, addr = await _build_store_for_worker(mock_worker, tmp_path) + t0 = parsed.traces[0].id + + monkeypatch.setattr(mock_worker.inference_client, "send_request", AsyncMock()) + mock_worker._send_inference_result_message = AsyncMock() + + credit = Credit( + id=8, + phase=CreditPhase.PROFILING, + conversation_id=t0, + x_correlation_id="x-8", + turn_index=0, + num_turns=1, + issued_at_ns=time.time_ns(), + trace_id=t0, + node_ordinal=9999, + phase_variant="profiling", + ) + ctx = CreditContext(credit=credit, drop_perf_ns=time.perf_counter_ns()) + + await mock_worker._process_credit(ctx) + + mock_worker.inference_client.send_request.assert_not_awaited() + assert ctx.error is not None + assert ctx.error.type == "GraphEnvelopeMissing" + + +@pytest.mark.asyncio +async def test_pre_v3_unified_store_rejection_surfaces_in_fatal_error( + mock_worker, tmp_path, monkeypatch +): + """An on-disk pre-v3 (A1 JSON index) unified store surfaces its A2-strict + rejection in the fatal error instead of claiming no store exists.""" + mock_worker._graph_client_metadata = _graph_client_metadata( + tmp_path, mock_worker.run.benchmark_id + ) + store_dir = tmp_path / f"aiperf_graph_segments_{mock_worker.run.benchmark_id}" + store_dir.mkdir(parents=True) + # Legacy A1 index: a JSON object (starts with b'{') that the A2-strict + # reader rejects with "re-parse required". + (store_dir / "content.idx").write_bytes(b'{"ab12": [0, 4]}') + + monkeypatch.setattr(mock_worker.inference_client, "send_request", AsyncMock()) + mock_worker._send_inference_result_message = AsyncMock() + + credit = Credit( + id=10, + phase=CreditPhase.PROFILING, + conversation_id="t-x", + x_correlation_id="x-10", + turn_index=0, + num_turns=1, + issued_at_ns=time.time_ns(), + trace_id="t-x", + node_ordinal=0, + phase_variant="profiling", + ) + ctx = CreditContext(credit=credit, drop_perf_ns=time.perf_counter_ns()) + + await mock_worker._process_credit(ctx) + + mock_worker.inference_client.send_request.assert_not_awaited() + assert ctx.error is not None + assert ctx.error.type == "GraphStoreUnavailable" + assert "rejected" in ctx.error.message + assert "re-parse required" in ctx.error.message + + +@pytest.mark.asyncio +async def test_non_graph_credit_still_uses_session_path(mock_worker, monkeypatch): + """A non-graph credit (no trace_id) must NOT touch the graph materializer.""" + # Force the session path to a controlled early return: no dataset client and + # not stopping triggers the dataset-manager fallback we can assert against. + called = {"graph": False} + orig = mock_worker._process_graph_credit + + async def spy(*a, **k): + called["graph"] = True + return await orig(*a, **k) + + monkeypatch.setattr(mock_worker, "_process_graph_credit", spy) + monkeypatch.setattr( + mock_worker, + "_retrieve_conversation", + AsyncMock(side_effect=RuntimeError("stop")), + ) + + credit = Credit( + id=9, + phase=CreditPhase.PROFILING, + conversation_id="conv-9", + x_correlation_id="x-9", + turn_index=0, + num_turns=1, + issued_at_ns=time.time_ns(), + ) + ctx = CreditContext(credit=credit, drop_perf_ns=time.perf_counter_ns()) + + await mock_worker._process_credit(ctx) + assert called["graph"] is False diff --git a/tests/unit/graph/test_worker_graph_capture.py b/tests/unit/graph/test_worker_graph_capture.py new file mode 100644 index 0000000000..2e683d5b67 --- /dev/null +++ b/tests/unit/graph/test_worker_graph_capture.py @@ -0,0 +1,377 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Worker-side graph response capture. + +Drives ``Worker._process_graph_credit`` unbound against a mock self with a +REAL unified store and a REAL ``GraphDynamicPool``: a ``capture: true`` node +envelope (written through the raw ``add_node_manifest`` seam — the native +lowering stamps it in production) pools a structured ``GraphCapturedReply`` +(text, plus the verbatim assistant message JSON for tool_calls replies) before +the credit returns; error/cancel paths pool ``FAILED``; genuinely-empty +responses pool ``EMPTY``; envelopes without the flag never touch the pool. weka/dynamo +byte-parity follows from that last case: their envelopes never carry the flag. +""" + +import types +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +import orjson +import pytest + +from aiperf.common.enums import CacheBustTarget, CreditPhase +from aiperf.credit.structs import Credit, CreditContext +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) +from aiperf.graph.dynamic_pool import ( + GraphCapturedReply, + GraphDynamicPool, + GraphPoolSentinel, +) +from aiperf.workers.worker import Worker + +TRACE = "t-1" +# Instance id is ``{template}::{nonce}``; the worker strips the ``::{nonce}`` +# back to the base template id for catalog / store lookups. +INSTANCE = "t-1::inst0" + + +def _credit_context() -> CreditContext: + return CreditContext( + credit=Credit( + id=1, + phase=CreditPhase.PROFILING, + conversation_id=TRACE, + x_correlation_id="t-1::corr0", + turn_index=0, + num_turns=1, + issued_at_ns=0, + trace_id=INSTANCE, + node_ordinal=0, + ), + drop_perf_ns=0, + ) + + +def _mock_worker( + client: GraphSegmentUnifiedClient, + *, + record: MagicMock, + assistant_turn, +) -> MagicMock: + self = MagicMock() + self._graph_store_reader = MagicMock(return_value=client) + self._graph_dynamic_pool = GraphDynamicPool(max_bytes=1024 * 1024) + self.model_endpoint.primary_model_name = "mock-model" + self.model_endpoint.endpoint.cache_bust = CacheBustTarget.NONE + self.model_endpoint.endpoint.use_legacy_max_tokens = False + self._build_graph_request_info = MagicMock(return_value=MagicMock()) + self.inference_client.send_request = AsyncMock(return_value=record) + self.inference_client.endpoint.build_assistant_turn = MagicMock( + return_value=assistant_turn + ) + self._send_inference_result_message = AsyncMock() + self._set_graph_envelope_missing = MagicMock() + self.task_stats = MagicMock() + # Bind the real methods under test onto the mock self. + self._process_graph_credit = types.MethodType(Worker._process_graph_credit, self) + self._dispatch_graph_request = types.MethodType( + Worker._dispatch_graph_request, self + ) + self._graph_capture_value = types.MethodType(Worker._graph_capture_value, self) + return self + + +def _text_turn(text: str): + return types.SimpleNamespace( + texts=[types.SimpleNamespace(contents=[text])], + raw_messages=None, + ) + + +TOOL_CALLS = [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}, + } +] + + +def _tool_calls_turn(content: str | None): + """Chat-endpoint shape: ``build_assistant_turn`` returned ``raw_messages``.""" + return types.SimpleNamespace( + texts=None, + raw_messages=[ + {"role": "assistant", "content": content, "tool_calls": TOOL_CALLS} + ], + ) + + +def _multi_entry_turn(): + """Endpoint returned more than one assistant message (e.g. openai_responses).""" + return types.SimpleNamespace( + texts=None, + raw_messages=[ + {"role": "assistant", "content": "first"}, + {"role": "assistant", "content": "second"}, + ], + ) + + +def _unserializable_turn(): + """Single assistant message whose content cannot be orjson-serialized.""" + return types.SimpleNamespace( + texts=None, + raw_messages=[{"role": "assistant", "content": object()}], + ) + + +def _record(error: str | None = None) -> MagicMock: + record = MagicMock() + record.error = error + return record + + +async def _run(tmp_path: Path, envelope: dict, **worker_kw) -> MagicMock: + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id="bench") + handle = store.put_segment("seg0", "user", "hello") + store.add_node_manifest( + TRACE, 0, "profiling", orjson.dumps({"handles": [handle], **envelope}) + ) + await store.finalize() + client = GraphSegmentUnifiedClient(tmp_path, "bench").open() + self = _mock_worker(client, **worker_kw) + ctx = _credit_context() + await self._process_graph_credit(ctx, "x-req-1", None) + self._ctx = ctx + return self + + +@pytest.mark.asyncio +async def test_captured_node_pools_response_text(tmp_path: Path) -> None: + self = await _run( + tmp_path, + {"capture": True}, + record=_record(), + assistant_turn=_text_turn("the answer"), + ) + assert self._graph_dynamic_pool.get(INSTANCE, "profiling", 0) == GraphCapturedReply( + text="the answer", message_json=None + ) + assert self._ctx.error is None + + +@pytest.mark.asyncio +async def test_tool_calls_only_reply_pools_structured_capture(tmp_path: Path) -> None: + """A tool_calls-only reply is a structured capture, NOT EMPTY: its + ``message_json`` splices the verbatim assistant message into child seeds.""" + self = await _run( + tmp_path, + {"capture": True}, + record=_record(), + assistant_turn=_tool_calls_turn(None), + ) + value = self._graph_dynamic_pool.get(INSTANCE, "profiling", 0) + assert isinstance(value, GraphCapturedReply) + assert value.text == "" + assert orjson.loads(value.message_json) == { + "role": "assistant", + "content": None, + "tool_calls": TOOL_CALLS, + } + + +@pytest.mark.asyncio +async def test_mixed_text_and_tool_calls_reply_captures_both(tmp_path: Path) -> None: + turn = _tool_calls_turn("let me check") + self = await _run( + tmp_path, + {"capture": True}, + record=_record(), + assistant_turn=turn, + ) + value = self._graph_dynamic_pool.get(INSTANCE, "profiling", 0) + assert isinstance(value, GraphCapturedReply) + assert value.text == "let me check" + assert value.message_json == orjson.dumps(turn.raw_messages[0]).decode() + + +@pytest.mark.asyncio +async def test_text_turn_joining_to_empty_pools_empty(tmp_path: Path) -> None: + self = await _run( + tmp_path, + {"capture": True}, + record=_record(), + assistant_turn=_text_turn(""), + ) + assert ( + self._graph_dynamic_pool.get(INSTANCE, "profiling", 0) + is GraphPoolSentinel.EMPTY + ) + + +@pytest.mark.asyncio +async def test_uncaptured_node_never_touches_pool(tmp_path: Path) -> None: + self = await _run( + tmp_path, + {}, + record=_record(), + assistant_turn=_text_turn("ignored"), + ) + assert self._graph_dynamic_pool.get(INSTANCE, "profiling", 0) is None + assert self._graph_dynamic_pool.total_bytes == 0 + + +@pytest.mark.asyncio +async def test_error_record_pools_failed(tmp_path: Path) -> None: + self = await _run( + tmp_path, + {"capture": True}, + record=_record(error="boom"), + assistant_turn=None, + ) + assert ( + self._graph_dynamic_pool.get(INSTANCE, "profiling", 0) + is GraphPoolSentinel.FAILED + ) + + +@pytest.mark.asyncio +async def test_empty_response_pools_empty(tmp_path: Path) -> None: + self = await _run( + tmp_path, + {"capture": True}, + record=_record(), + assistant_turn=None, + ) + assert ( + self._graph_dynamic_pool.get(INSTANCE, "profiling", 0) + is GraphPoolSentinel.EMPTY + ) + + +@pytest.mark.asyncio +async def test_send_exception_pools_failed_and_attributes_error(tmp_path: Path) -> None: + """A raising send must not escape the credit task (the outer handler only + logs, so the credit would be counted as a SUCCESS with no record and the + RecordsManager barrier would starve): FAILED is pooled, the error is + attributed on the context, and a synthetic error record is emitted.""" + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id="bench") + handle = store.put_segment("seg0", "user", "hello") + store.add_node_manifest( + TRACE, 0, "profiling", orjson.dumps({"handles": [handle], "capture": True}) + ) + await store.finalize() + client = GraphSegmentUnifiedClient(tmp_path, "bench").open() + self = _mock_worker(client, record=_record(), assistant_turn=None) + self.inference_client.send_request = AsyncMock(side_effect=TimeoutError("dead")) + self._send_graph_error_record = AsyncMock() + + ctx = _credit_context() + await self._process_graph_credit(ctx, "x-req-1", None) + + assert ( + self._graph_dynamic_pool.get(INSTANCE, "profiling", 0) + is GraphPoolSentinel.FAILED + ) + assert ctx.error is not None and ctx.error.type == "TimeoutError" + self._send_graph_error_record.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_capture_extraction_failure_sets_error_and_pools_failed( + tmp_path: Path, +) -> None: + self = await _run( + tmp_path, + {"capture": True}, + record=_record(), + assistant_turn=None, + ) + # Re-run with a raising extractor against the same harness shape. + self.inference_client.endpoint.build_assistant_turn = MagicMock( + side_effect=RuntimeError("bad parse") + ) + ctx = _credit_context() + await self._process_graph_credit(ctx, "x-req-2", None) + + assert ( + self._graph_dynamic_pool.get(INSTANCE, "profiling", 0) + is GraphPoolSentinel.FAILED + ) + assert ctx.error is not None and "aiperf.graph.capture_failed" in ctx.error + + +@pytest.mark.asyncio +async def test_multi_entry_raw_messages_pools_failed_with_pointed_error( + tmp_path: Path, +) -> None: + """A single-message capture cannot represent >1 assistant message: FAILED + is pooled and a pointed capture_failed error names the entry count so the + truncation-to-first is never silent.""" + self = await _run( + tmp_path, + {"capture": True}, + record=_record(), + assistant_turn=_multi_entry_turn(), + ) + assert ( + self._graph_dynamic_pool.get(INSTANCE, "profiling", 0) + is GraphPoolSentinel.FAILED + ) + assert self._ctx.error == ( + "aiperf.graph.capture_failed: multi-entry raw_messages (2 entries) is " + "not representable in a single-message capture; single-message " + "endpoints only" + ) + + +@pytest.mark.asyncio +async def test_unserializable_reply_pools_failed_without_escape( + tmp_path: Path, +) -> None: + """An assistant message that orjson cannot serialize must map to FAILED + with capture_failed attribution -- never escape as an unhandled exception + that would drop the pool write.""" + self = await _run( + tmp_path, + {"capture": True}, + record=_record(), + assistant_turn=_unserializable_turn(), + ) + assert ( + self._graph_dynamic_pool.get(INSTANCE, "profiling", 0) + is GraphPoolSentinel.FAILED + ) + assert self._ctx.error is not None + assert "aiperf.graph.capture_failed" in self._ctx.error + + +@pytest.mark.asyncio +async def test_trace_end_mid_flight_defers_eviction(tmp_path: Path) -> None: + """GraphTraceEnd during processing evicts only after the bracket closes.""" + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id="bench") + handle = store.put_segment("seg0", "user", "hello") + store.add_node_manifest( + TRACE, 0, "profiling", orjson.dumps({"handles": [handle], "capture": True}) + ) + await store.finalize() + client = GraphSegmentUnifiedClient(tmp_path, "bench").open() + self = _mock_worker(client, record=_record(), assistant_turn=_text_turn("late")) + + pool = self._graph_dynamic_pool + + async def _send(request_info, first_token_callback=None): + # TraceEnd lands while the request is in flight. + pool.trace_end(INSTANCE, "profiling") + return _record() + + self.inference_client.send_request = AsyncMock(side_effect=_send) + await self._process_graph_credit(_credit_context(), "x-req-1", None) + + # The capture write landed in a live entry, then the deferred end evicted. + assert pool.get(INSTANCE, "profiling", 0) is None + assert pool.total_bytes == 0 diff --git a/tests/unit/graph/test_worker_materialize.py b/tests/unit/graph/test_worker_materialize.py new file mode 100644 index 0000000000..4c7aed6008 --- /dev/null +++ b/tests/unit/graph/test_worker_materialize.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Worker-side materialization of a graph node's request on the segment-trie IR. + +A trie node's interned manifest carries ``handles`` (an int-handle path into the +unified content pool); the materializer walks that path against the unified +store client to produce ``messages`` -- NO ancestor accumulation, NO reset +handling -- then applies the node's ``dispatch_overrides`` (mapping the +recorded token cap to the wire field) and ``stream`` flag. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.segment_ir.store_builder import ( + build_unified_trie_store_interned, +) +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) +from aiperf.graph.worker_materialize import ( + materialize_graph_request_unified, +) + +FIX_MIN = Path(__file__).parent / "fixtures" / "weka_min.json" + + +@pytest.mark.asyncio +async def test_materialize_maps_max_output_tokens_to_wire_field(tmp_path): + """``max_output_tokens`` -> ``max_tokens`` (legacy) / ``max_completion_tokens`` + (modern) on the trie materialization path. + + A trie node carries its prompt as interned ``handles``; its + ``dispatch_overrides`` token cap must map to the endpoint-appropriate wire + token field exactly as on the linear path (the mapping the worker applies + after walking the handle path). + """ + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id="b1") + handle = store.put_segment("s_u1", "user", "hi") + store.add_node_manifest_interned( + "t-1#0", 0, "profiling", [handle], {"max_output_tokens": 25}, False + ) + await store.finalize() + + client = GraphSegmentUnifiedClient(base_path=tmp_path, benchmark_id="b1").open() + try: + legacy = materialize_graph_request_unified( + client, "t-1#0", 0, "profiling", use_legacy_max_tokens=True + ) + modern = materialize_graph_request_unified( + client, "t-1#0", 0, "profiling", use_legacy_max_tokens=False + ) + finally: + client.close() + + assert legacy["messages"] == [{"role": "user", "content": "hi"}] + assert legacy.get("max_tokens") == 25 + assert "max_completion_tokens" not in legacy + assert "max_output_tokens" not in legacy + + assert modern.get("max_completion_tokens") == 25 + assert "max_tokens" not in modern + + +@pytest.mark.asyncio +async def test_materialize_missing_node_returns_none(tmp_path): + """A node ordinal not in the store materializes to None (graceful miss).""" + parsed = from_weka_trace(str(FIX_MIN)) + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id="b1") + await build_unified_trie_store_interned(parsed, store) + + t0 = parsed.traces[0].id + # The unified client duck-types BOTH the addressing and segment + # (content) faces, so one client serves both materializers. + client = GraphSegmentUnifiedClient(base_path=tmp_path, benchmark_id="b1").open() + try: + assert materialize_graph_request_unified(client, t0, 9999, "profiling") is None + finally: + client.close() diff --git a/tests/unit/graph/test_worker_payload_features.py b/tests/unit/graph/test_worker_payload_features.py new file mode 100644 index 0000000000..51c9ae59ad --- /dev/null +++ b/tests/unit/graph/test_worker_payload_features.py @@ -0,0 +1,700 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""F3/F4 regression: graph credits must honor run-level payload options and +must surface a clear, fatal error (not a silent retry-loop) when the +graph store is absent or corrupt. + +F3 (payload features dropped): the worker must layer ``endpoint.extra`` +(``--extra-inputs``) and ``stream_options.include_usage`` (``--server-token-count`` +on a streaming request) onto the materialized payload. The wire ``stream`` is +stamped from the RECORDED per-node mode (weka ``"n"``/``"s"``): a recorded +streaming node streams even when the global run is non-streaming, and a recorded +non-streaming node stays non-streaming even when the global run streams; the +global ``endpoint.streaming`` is only the fallback for a mode-less node. +``stream_options.include_usage`` keys on that FINAL stamped ``stream``. User +``extra`` CLOBBERS any colliding per-node dispatch key +(``payload.update(endpoint.extra)``); non-colliding per-node ``dispatch_overrides`` +still pass through. + +F4 (missing-store hang): when the unified store open raises because the +store files are missing, the worker must set ``credit_context.error`` (so the +run reports a real error, not zero-success/zero-error), and must not retry the +failed open on every subsequent credit. +""" + +import time +from pathlib import Path +from unittest.mock import AsyncMock + +import orjson +import pytest +from pytest import param + +from aiperf.common.enums import CreditPhase +from aiperf.common.models import RequestRecord +from aiperf.common.models.dataset_models import GraphSegmentClientMetadata +from aiperf.credit.structs import Credit, CreditContext +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.dataset.graph.segment_ir.store_builder import ( + build_unified_trie_store_interned, +) +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) +from aiperf.graph.worker_materialize import ( + materialize_graph_request_unified, + strip_dynamo_session_headers, + uniquify_dynamo_session_headers, +) +from aiperf.workers.worker import Worker +from tests.harness.fake_tokenizer import FakeTokenizer + +FIX_MIN = Path(__file__).parent / "fixtures" / "weka_min.json" +# Mixed-mode fixture: node trace_first_token_anchor:0 is recorded streaming ("s"); the top-level "n" node is trace_first_token_anchor:1. +FIX_MIXED = Path(__file__).parent / "fixtures" / "weka_first_token_anchor.json" + + +def _graph_client_metadata( + tmp_path: Path, benchmark_id: str +) -> GraphSegmentClientMetadata: + """Broadcast-shaped store location the worker opens (Task 5: no env fallback).""" + return GraphSegmentClientMetadata( + store_base_path=tmp_path, + benchmark_id=benchmark_id, + sidecar_path=tmp_path + / f"aiperf_graph_meta_{benchmark_id}" + / "graph_meta.msgpack", + ) + + +@pytest.fixture +async def mock_worker( + benchmark_run, fake_tokenizer: FakeTokenizer, skip_service_registration +): + worker = Worker(run=benchmark_run, service_id="mock-service-id") + await worker.initialize() + await worker.start() + yield worker + await worker.stop() + + +async def _build_store_for_worker( + worker: Worker, tmp_path: Path, *, fixture: Path = FIX_MIN +): + """Build the interned unified trie store the worker will open. + + Mirrors ``dataset_manager._build_graph_trie_stores``: the content-addressed + pool AND every node's ``prompt_segment_ids`` manifest land in the ONE + unified store. The worker opens it from the graph-typed dataset broadcast + (``GraphSegmentClientMetadata``), which we set directly here. + """ + parsed = from_weka_trace(str(fixture)) + store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id=worker.run.benchmark_id + ) + addr = await build_unified_trie_store_interned(parsed, store) + worker._graph_client_metadata = _graph_client_metadata( + tmp_path, worker.run.benchmark_id + ) + return parsed, addr + + +# --------------------------------------------------------------------------- +# F3 -- payload features carried onto the wire +# --------------------------------------------------------------------------- + + +async def _capture_graph_wire_payload( + mock_worker, + tmp_path, + monkeypatch, + *, + global_streaming: bool, + fixture: Path = FIX_MIN, + node_key: str = "trace_03_n3:2", +) -> dict: + """Drive one graph credit through the worker, formatting the final wire + payload exactly as the inference client would, and return it. + + ``global_streaming`` sets the run-level ``endpoint.streaming`` -- the FALLBACK + for a mode-less node; the recorded per-node mode (``node_key`` in ``fixture``) + wins for the wire ``stream`` byte. + """ + parsed, addr = await _build_store_for_worker(mock_worker, tmp_path, fixture=fixture) + t0 = parsed.traces[0].id + leaf_ordinal = addr[t0][node_key] + + endpoint = mock_worker.model_endpoint.endpoint + monkeypatch.setattr(endpoint, "streaming", global_streaming) + monkeypatch.setattr(endpoint, "use_server_token_count", True) + monkeypatch.setattr(endpoint, "extra", [("guided_decoding_backend", "outlines")]) + + captured: dict = {} + + async def capture_wire(request_info, first_token_callback=None): + # Mirror the real three-way select in _send_request_to_transport: + # raw_payload_bytes (verbatim) -> raw_payload (dict) -> format_payload. + rpb = request_info.turns[-1].raw_payload_bytes + raw_payload = request_info.turns[-1].raw_payload + if rpb is not None: + payload = orjson.loads(rpb) + elif raw_payload is not None: + payload = raw_payload + else: + payload = mock_worker.inference_client.endpoint.format_payload(request_info) + request_info.payload_bytes = orjson.dumps(payload) + captured["wire_payload"] = payload + return RequestRecord( + request_info=request_info, + timestamp_ns=time.time_ns(), + start_perf_ns=time.perf_counter_ns(), + end_perf_ns=time.perf_counter_ns(), + ) + + monkeypatch.setattr( + mock_worker.inference_client, + "send_request", + AsyncMock(side_effect=capture_wire), + ) + mock_worker._send_inference_result_message = AsyncMock() + + credit = Credit( + id=21, + phase=CreditPhase.PROFILING, + conversation_id=t0, + x_correlation_id="x-21", + turn_index=0, + num_turns=1, + issued_at_ns=time.time_ns(), + trace_id=t0, + node_ordinal=leaf_ordinal, + phase_variant="profiling", + ) + ctx = CreditContext(credit=credit, drop_perf_ns=time.perf_counter_ns()) + await mock_worker._process_credit(ctx) + assert ctx.error is None + return captured["wire_payload"] + + +@pytest.mark.asyncio +async def test_graph_payload_recorded_streaming_node_wins_over_global_off( + mock_worker, tmp_path, monkeypatch +): + """A recorded streaming (``"s"``) node streams on the wire even when the + global run is NON-streaming; ``extra`` is merged and ``include_usage`` is + layered because the FINAL stamped ``stream`` is True (server-token-count on). + Asserted on the FINAL payload, not the intermediate materialized dict.""" + wire = await _capture_graph_wire_payload( + mock_worker, + tmp_path, + monkeypatch, + global_streaming=False, + fixture=FIX_MIXED, + node_key="trace_first_token_anchor:0", + ) + assert wire["messages"], "materialized messages preserved" + assert wire["guided_decoding_backend"] == "outlines", "endpoint.extra merged" + assert wire["stream"] is True, "recorded 's' node streams despite global off" + assert wire["stream_options"] == {"include_usage": True}, ( + "include_usage layered when the FINAL stream is on + server-token-count" + ) + + +@pytest.mark.asyncio +async def test_graph_payload_recorded_non_streaming_node_wins_over_global_on( + mock_worker, tmp_path, monkeypatch +): + """A recorded non-streaming (``"n"``) node stays non-streaming even when the + global run streams; no ``stream_options`` is added even with + server-token-count on (the FINAL stream is False), and ``extra`` is still + merged.""" + wire = await _capture_graph_wire_payload( + mock_worker, tmp_path, monkeypatch, global_streaming=True + ) + assert wire["guided_decoding_backend"] == "outlines", "endpoint.extra still merged" + assert wire["stream"] is False, ( + "recorded 'n' node stays non-streaming despite global on" + ) + assert "stream_options" not in wire, "no usage opt-in when the FINAL stream is off" + + +@pytest.mark.asyncio +async def test_user_extra_clobbers_per_node_overrides(tmp_path): + """User ``--extra-inputs`` CLOBBERS colliding per-node dispatch keys (agentx + ``payload.update(endpoint.extra)`` precedence); non-colliding per-node keys + pass through. + + On the trie path the per-node ``dispatch_overrides`` carry the recorded + ``max_output_tokens`` cap (endpoint-mapped to ``max_completion_tokens`` + here) and ``model`` verbatim, so a run-level ``extra`` naming the mapped + wire key collides; the user must win. + """ + parsed = from_weka_trace(str(FIX_MIN)) + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id="ov") + addr = await build_unified_trie_store_interned(parsed, store) + + t0 = parsed.traces[0].id + ordinal = addr[t0]["trace_03_n3:0"] # recorded out cap 25 + client = GraphSegmentUnifiedClient(base_path=tmp_path, benchmark_id="ov").open() + + from aiperf.common.models import EndpointInfo + from aiperf.graph.worker_materialize import apply_run_level_payload_options + + payload = materialize_graph_request_unified( + client, + t0, + ordinal, + "profiling", + use_legacy_max_tokens=False, + ) + assert payload["max_completion_tokens"] == 25 + assert payload["model"] == "claude-opus-4-5-20251101" + + # Run-level extra collides with the per-node token cap + model; user wins. + endpoint = EndpointInfo( + type="chat", + streaming=True, + use_server_token_count=False, + extra=[("max_completion_tokens", 999), ("model", "run-level-model")], + ) + apply_run_level_payload_options(payload, endpoint) + assert payload["max_completion_tokens"] == 999, "user extra clobbers node cap" + assert payload["model"] == "run-level-model", "user extra clobbers node model" + client.close() + + +@pytest.mark.asyncio +async def test_include_usage_skipped_when_global_not_streaming(tmp_path): + """include_usage is only layered when the GLOBAL ``endpoint.streaming`` is on; + a False per-node materialized stream does not suppress it once global is on, + and a False global suppresses it regardless of the per-node value.""" + from aiperf.common.models import EndpointInfo + from aiperf.graph.worker_materialize import apply_run_level_payload_options + + # Global off -> wire stream forced False -> no include_usage. + endpoint_off = EndpointInfo( + type="chat", streaming=False, use_server_token_count=True, extra=[] + ) + payload = {"messages": [{"role": "user", "content": "hi"}], "stream": True} + apply_run_level_payload_options(payload, endpoint_off) + assert payload["stream"] is False, "global streaming=False stamps false" + assert "stream_options" not in payload + + # Global on -> wire stream forced True even though per-node was False -> + # include_usage layered. + endpoint_on = EndpointInfo( + type="chat", streaming=True, use_server_token_count=True, extra=[] + ) + payload2 = {"messages": [{"role": "user", "content": "hi"}], "stream": False} + apply_run_level_payload_options(payload2, endpoint_on) + assert payload2["stream"] is True, "global streaming=True stamps true" + assert payload2["stream_options"] == {"include_usage": True} + + +def test_skip_endpoint_extra_leaves_adapter_owned_key_untouched() -> None: + """``skip_endpoint_extra=True`` suppresses the ``endpoint.extra`` merge so an + adapter that already folded ``--extra-inputs`` into ``dispatch_overrides`` at + parse keeps its own value; the ``stream`` stamp and ``include_usage`` forcing + are UNCHANGED by the flag.""" + from aiperf.common.models import EndpointInfo + from aiperf.graph.worker_materialize import apply_run_level_payload_options + + endpoint = EndpointInfo( + type="chat", + streaming=True, + use_server_token_count=True, + extra=[("guided_decoding_backend", "run-level")], + ) + payload = { + "messages": [{"role": "user", "content": "hi"}], + "guided_decoding_backend": "adapter-owned", + } + apply_run_level_payload_options(payload, endpoint, skip_endpoint_extra=True) + assert payload["guided_decoding_backend"] == "adapter-owned", ( + "skip_endpoint_extra leaves the adapter-owned key un-clobbered" + ) + assert payload["stream"] is True, "stream stamp is unaffected by the flag" + assert payload["stream_options"] == {"include_usage": True}, ( + "include_usage forcing is unaffected by the flag" + ) + + +def test_skip_endpoint_extra_default_false_still_clobbers() -> None: + """Default ``skip_endpoint_extra=False`` is byte-identical to today: the + user's ``endpoint.extra`` still clobbers a colliding per-node key.""" + from aiperf.common.models import EndpointInfo + from aiperf.graph.worker_materialize import apply_run_level_payload_options + + endpoint = EndpointInfo( + type="chat", + streaming=False, + use_server_token_count=False, + extra=[("guided_decoding_backend", "run-level")], + ) + payload = { + "messages": [{"role": "user", "content": "hi"}], + "guided_decoding_backend": "adapter-owned", + } + apply_run_level_payload_options(payload, endpoint) + assert payload["guided_decoding_backend"] == "run-level", ( + "default behavior: user extra clobbers the per-node key" + ) + + +# --------------------------------------------------------------------------- +# F4 -- missing/corrupt store surfaces a clear fatal error (no silent retry) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_missing_store_surfaces_error_not_silent( + mock_worker, tmp_path, monkeypatch +): + """When the graph store is absent, the credit error is set with an + actionable message; the inference client is NOT called.""" + # Broadcast points at an empty dir -- no store written -> open() will fail. + mock_worker._graph_client_metadata = _graph_client_metadata( + tmp_path, mock_worker.run.benchmark_id + ) + monkeypatch.setattr(mock_worker.inference_client, "send_request", AsyncMock()) + mock_worker._send_inference_result_message = AsyncMock() + + credit = Credit( + id=41, + phase=CreditPhase.PROFILING, + conversation_id="t-missing", + x_correlation_id="x-41", + turn_index=0, + num_turns=1, + issued_at_ns=time.time_ns(), + trace_id="t-missing", + node_ordinal=0, + phase_variant="profiling", + ) + ctx = CreditContext(credit=credit, drop_perf_ns=time.perf_counter_ns()) + await mock_worker._process_credit(ctx) + + mock_worker.inference_client.send_request.assert_not_awaited() + assert ctx.error is not None, "missing store must set a fatal error, not swallow" + # Actionable: names the store directory / base path so the operator can fix + # the shared-FS / MMAP_BASE_PATH config. + assert str(tmp_path) in ctx.error.message or "graph store" in ctx.error.message + + +@pytest.mark.asyncio +async def test_missing_store_does_not_retry_open_each_credit( + mock_worker, tmp_path, monkeypatch +): + """A failed store open is cached: subsequent credits do not re-attempt the + open (no silent retry-loop).""" + mock_worker._graph_client_metadata = _graph_client_metadata( + tmp_path, mock_worker.run.benchmark_id + ) + monkeypatch.setattr(mock_worker.inference_client, "send_request", AsyncMock()) + mock_worker._send_inference_result_message = AsyncMock() + + open_calls = {"n": 0} + real_open = GraphSegmentUnifiedClient.open + + def counting_open(self): + open_calls["n"] += 1 + return real_open(self) + + monkeypatch.setattr(GraphSegmentUnifiedClient, "open", counting_open) + + def make_credit(cid): + return CreditContext( + credit=Credit( + id=cid, + phase=CreditPhase.PROFILING, + conversation_id="t-missing", + x_correlation_id=f"x-{cid}", + turn_index=0, + num_turns=1, + issued_at_ns=time.time_ns(), + trace_id="t-missing", + node_ordinal=0, + phase_variant="profiling", + ), + drop_perf_ns=time.perf_counter_ns(), + ) + + await mock_worker._process_credit(make_credit(50)) + await mock_worker._process_credit(make_credit(51)) + await mock_worker._process_credit(make_credit(52)) + + assert open_calls["n"] == 1, ( + "store open must be attempted once and the failure cached, " + f"not retried every credit (got {open_calls['n']} opens)" + ) + + +@pytest.mark.asyncio +async def test_graph_envelope_extra_headers_reach_request_turn(mock_worker): + """Envelope ``extra_headers`` (dynamo ``x-dynamo-*`` session identity) land + on the synthetic Turn -- the transport merges the LAST turn's extra_headers + into the request headers -- and never leak into the body payload.""" + headers = { + "x-dynamo-session-id": "sess-1", + "x-dynamo-parent-session-id": "root-1", + "x-dynamo-session-final": "true", + } + credit = Credit( + id=7, + phase=CreditPhase.PROFILING, + conversation_id="t-1#0", + x_correlation_id="x-7", + turn_index=0, + num_turns=1, + issued_at_ns=time.time_ns(), + trace_id="t-1#0", + node_ordinal=0, + phase_variant="profiling", + ) + ctx = CreditContext(credit=credit, drop_perf_ns=time.perf_counter_ns()) + + payload = {"messages": [{"role": "user", "content": "hi"}], "model": "m"} + info = mock_worker._build_graph_request_info( + ctx, payload, "x-req-7", extra_headers=headers + ) + assert info.turns[-1].extra_headers == headers + assert "extra_headers" not in payload and "nvext" not in payload + + info_bytes = mock_worker._build_graph_request_info( + ctx, + None, + "x-req-7", + raw_payload_bytes=b"{}", + extra_headers=headers, + ) + assert info_bytes.turns[-1].extra_headers == headers + + info_none = mock_worker._build_graph_request_info(ctx, payload, "x-req-7") + assert info_none.turns[-1].extra_headers is None + + +# --------------------------------------------------------------------------- +# N1 -- dynamo session identity uniquified per replay instance +# --------------------------------------------------------------------------- + +RECORDED_SESSION_HEADERS = { + "x-dynamo-session-id": "sess-X", + "x-dynamo-parent-session-id": "root-X", + "x-dynamo-session-final": "true", +} + + +def test_uniquify_two_instances_get_distinct_linked_session_ids() -> None: + """Two replay instances of ONE trace must open DISTINCT server sessions + (the first finisher's session-final would otherwise evict KV under the + still-running sibling); parent linkage transforms with the SAME suffix and + session-final is forwarded untouched.""" + h0 = uniquify_dynamo_session_headers( + dict(RECORDED_SESSION_HEADERS), + trace_instance_id="t-1::nonceA", + phase_variant="profiling", + ) + h1 = uniquify_dynamo_session_headers( + dict(RECORDED_SESSION_HEADERS), + trace_instance_id="t-1::nonceB", + phase_variant="profiling", + ) + assert h0["x-dynamo-session-id"] != h1["x-dynamo-session-id"] + for h in (h0, h1): + assert h["x-dynamo-session-id"] != "sess-X" + suffix = h["x-dynamo-session-id"].removeprefix("sess-X") + assert suffix, "instance suffix must be appended to the recorded id" + # Parent gets the SAME suffix: intra-instance parent-child linkage holds. + assert h["x-dynamo-parent-session-id"] == f"root-X{suffix}" + # session-final is a per-turn flag, not an identity: untouched, so each + # instance still closes (only) its own session on its final turn. + assert h["x-dynamo-session-final"] == "true" + + +def test_uniquify_deterministic_within_instance_and_phase() -> None: + """Every dispatch of one instance must agree on the transformed ids (the + stateless worker recomputes per credit), and a warmup instance must not + collide with the profiling instance of the same (lane, pass) slot.""" + kwargs = {"trace_instance_id": "t-1::nonce0", "phase_variant": "profiling"} + first = uniquify_dynamo_session_headers(dict(RECORDED_SESSION_HEADERS), **kwargs) + second = uniquify_dynamo_session_headers(dict(RECORDED_SESSION_HEADERS), **kwargs) + assert first == second + + warmup = uniquify_dynamo_session_headers( + dict(RECORDED_SESSION_HEADERS), + trace_instance_id="t-1::nonce0", + phase_variant="warmup", + ) + assert warmup["x-dynamo-session-id"] != first["x-dynamo-session-id"] + + +@pytest.mark.parametrize( + "extra_headers, trace_instance_id", + [ + param(None, "t-1::n", id="no_headers"), + param({}, "t-1::n", id="empty_headers"), + param({"x-other": "v"}, "t-1::n", id="no_dynamo_identity_header"), + param(dict(RECORDED_SESSION_HEADERS), "t-1", id="no_instance_suffix"), + param(dict(RECORDED_SESSION_HEADERS), None, id="no_trace_id"), + ], +) # fmt: skip +def test_uniquify_noop_paths_return_input_unchanged( + extra_headers: dict | None, trace_instance_id: str | None +) -> None: + """Plain (non-instanced) replay and header-less nodes are unaffected.""" + result = uniquify_dynamo_session_headers( + extra_headers, + trace_instance_id=trace_instance_id, + phase_variant="profiling", + ) + assert result is extra_headers + + +def test_strip_dynamo_session_headers_removes_identity_keeps_rest() -> None: + """With a live --session-routing plugin owning session identity, the + RECORDED identity headers (session id / parent / session-final) are stale + replay artifacts and must not ride the wire beside the plugin's live + headers; unrelated recorded headers survive.""" + headers = {**RECORDED_SESSION_HEADERS, "x-custom": "keep-me"} + stripped = strip_dynamo_session_headers(headers) + assert stripped == {"x-custom": "keep-me"} + # Identity-only header sets strip to None (no empty-dict envelope noise). + assert strip_dynamo_session_headers(dict(RECORDED_SESSION_HEADERS)) is None + # No-op paths return the input unchanged. + assert strip_dynamo_session_headers(None) is None + plain = {"x-custom": "keep-me"} + assert strip_dynamo_session_headers(plain) is plain + + +async def _build_session_headers_store(worker: Worker, tmp_path: Path): + """One-node unified store whose envelope carries recorded dynamo headers.""" + store = GraphSegmentUnifiedBackingStore( + base_path=tmp_path, benchmark_id=worker.run.benchmark_id + ) + handle = store.put_segment("seg0", "user", "hello") + store.add_node_manifest( + "t-1", + 0, + "profiling", + orjson.dumps({"handles": [handle], "extra_headers": RECORDED_SESSION_HEADERS}), + ) + await store.finalize() + worker._graph_client_metadata = _graph_client_metadata( + tmp_path, worker.run.benchmark_id + ) + + +@pytest.mark.asyncio +async def test_graph_credit_flow_uniquifies_session_headers_per_instance( + mock_worker, tmp_path, monkeypatch +) -> None: + """End to end through ``_process_credit``: two credits addressing the SAME + trace under different instance ids reach the wire with DIFFERENT + x-dynamo-session-id values, consistently-transformed parent ids, and an + untouched per-instance session-final.""" + await _build_session_headers_store(mock_worker, tmp_path) + + captured: list[dict] = [] + + async def capture_headers(request_info, first_token_callback=None): + captured.append(request_info.turns[-1].extra_headers) + return RequestRecord( + request_info=request_info, + timestamp_ns=time.time_ns(), + start_perf_ns=time.perf_counter_ns(), + end_perf_ns=time.perf_counter_ns(), + ) + + monkeypatch.setattr( + mock_worker.inference_client, + "send_request", + AsyncMock(side_effect=capture_headers), + ) + mock_worker._send_inference_result_message = AsyncMock() + + for i, instance_id in enumerate(("t-1::inst0", "t-1::inst1")): + credit = Credit( + id=100 + i, + phase=CreditPhase.PROFILING, + conversation_id="t-1", + x_correlation_id=f"x-{100 + i}", + turn_index=0, + num_turns=1, + issued_at_ns=time.time_ns(), + trace_id=instance_id, + node_ordinal=0, + phase_variant="profiling", + ) + ctx = CreditContext(credit=credit, drop_perf_ns=time.perf_counter_ns()) + await mock_worker._process_credit(ctx) + assert ctx.error is None + + h0, h1 = captured + assert h0["x-dynamo-session-id"] != h1["x-dynamo-session-id"] + for h in (h0, h1): + assert h["x-dynamo-session-id"] != "sess-X" + suffix = h["x-dynamo-session-id"].removeprefix("sess-X") + assert h["x-dynamo-parent-session-id"] == f"root-X{suffix}" + assert h["x-dynamo-session-final"] == "true" + + +@pytest.mark.asyncio +async def test_graph_credit_flow_strips_recorded_headers_when_routing_active( + mock_worker, tmp_path, monkeypatch +) -> None: + """With a live --session-routing plugin, the plugin OWNS session identity: + the recorded x-dynamo-* identity headers are STRIPPED (not uniquified), + and the graph RequestInfo carries the live routing facts the chokepoint + stamps from (corr / root corr / recorded finality).""" + await _build_session_headers_store(mock_worker, tmp_path) + monkeypatch.setattr( + type(mock_worker.inference_client), + "session_routing_active", + property(lambda self: True), + ) + + captured: list = [] + + async def capture_request_info(request_info, first_token_callback=None): + captured.append(request_info) + return RequestRecord( + request_info=request_info, + timestamp_ns=time.time_ns(), + start_perf_ns=time.perf_counter_ns(), + end_perf_ns=time.perf_counter_ns(), + ) + + monkeypatch.setattr( + mock_worker.inference_client, + "send_request", + AsyncMock(side_effect=capture_request_info), + ) + mock_worker._send_inference_result_message = AsyncMock() + + credit = Credit( + id=300, + phase=CreditPhase.PROFILING, + conversation_id="t-1", + x_correlation_id="t-1::corr300", + turn_index=0, + num_turns=1, + issued_at_ns=time.time_ns(), + trace_id="t-1::inst0", + node_ordinal=0, + phase_variant="profiling", + root_correlation_id="t-1::root-corr", + ) + ctx = CreditContext(credit=credit, drop_perf_ns=time.perf_counter_ns()) + await mock_worker._process_credit(ctx) + assert ctx.error is None + + (request_info,) = captured + # Recorded dynamo identity headers are gone (identity-only set -> None). + assert request_info.turns[-1].extra_headers is None + # Routing identity facts ride the RequestInfo for the chokepoint. + assert request_info.x_correlation_id == "t-1::corr300" + assert request_info.root_correlation_id == "t-1::root-corr" + assert request_info.is_final_turn is True # num_turns=1: recorded final + assert request_info.is_parent_final is None + assert request_info.is_tree_final is False diff --git a/tests/unit/graph/test_workload_detect_predicates.py b/tests/unit/graph/test_workload_detect_predicates.py new file mode 100644 index 0000000000..370af91d80 --- /dev/null +++ b/tests/unit/graph/test_workload_detect_predicates.py @@ -0,0 +1,318 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tightened weka workload auto-detection predicate tests. + +Pins the false-positive surface called out in ``adv2-integration.md`` F3: + +* ``_is_weka_trace_object`` must require the genuine :class:`WekaTrace` + discriminator strictly and REJECT objects carrying foreign/contradictory + top-level keys (a mooncake/sharegpt-shaped object that happens to also carry + the five weka keys must NOT be detected), and +* ``_looks_like_hf_dataset_id`` must NOT treat an arbitrary ``org/name`` string + as a weka HuggingFace dataset id -- only repo ids that carry the weka corpus + marker are routed to ``datasets.load_dataset``. + +Genuine weka inputs (the existing fixtures and the canonical weka HF corpus id) +MUST still be detected so :func:`resolve_graph_workload` stays correct. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +from pytest import param + +from aiperf.dataset.graph.adapters.weka.trace import ( + WekaTraceAdapter, + _is_weka_trace_object, + _looks_like_hf_dataset_id, +) + +FIXTURES = Path(__file__).parent / "fixtures" +WEKA_MIN = FIXTURES / "weka_min.json" + +# A minimal but genuine weka trace object (the on-disk single-document shape). +_GENUINE_WEKA_OBJECT: dict = { + "id": "trace_genuine", + "models": ["claude-opus-4-5-20251101"], + "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "n", "model": "m", "in": 10, "out": 5}, + ], +} + + +# --------------------------------------------------------------------------- # +# _is_weka_trace_object: genuine still detected # +# --------------------------------------------------------------------------- # +def test_genuine_weka_object_detected(): + assert _is_weka_trace_object(_GENUINE_WEKA_OBJECT) is True + + +def test_genuine_weka_object_with_optional_fields_detected(): + # tool_tokens / system_tokens / totals are legitimate WekaTrace fields. + doc = dict(_GENUINE_WEKA_OBJECT) + doc["tool_tokens"] = 12 + doc["system_tokens"] = 8 + doc["totals"] = {"requests": 1} + assert _is_weka_trace_object(doc) is True + + +def test_genuine_weka_jsonl_kind_marker_detected(): + # The native JSONL writer stamps a ``kind`` marker on serialized rows; a + # weka object that also carries ``kind`` must still be accepted. + doc = dict(_GENUINE_WEKA_OBJECT) + doc["kind"] = "trace" + assert _is_weka_trace_object(doc) is True + + +def test_global_hash_scope_detected(): + # 'global' is a supported hash_id_scope (cross-trace shared hash + # namespace), so a global-scope trace must detect as weka. + doc = {**_GENUINE_WEKA_OBJECT, "hash_id_scope": "global"} + assert _is_weka_trace_object(doc) is True + + +def test_weka_min_fixture_file_detected(): + assert WekaTraceAdapter.can_load(WEKA_MIN) is True + + +# --------------------------------------------------------------------------- # +# _is_weka_trace_object: near-misses and foreign-shaped objects rejected # +# --------------------------------------------------------------------------- # +@pytest.mark.parametrize( + "doc", + [ + param( + {**_GENUINE_WEKA_OBJECT, "hash_id_scope": "cluster"}, + id="unrecognized_hash_scope", + ), + param( + {**_GENUINE_WEKA_OBJECT, "conversations": [{"from": "human"}]}, + id="sharegpt_foreign_key", + ), + param( + {**_GENUINE_WEKA_OBJECT, "mooncake_field": True, "timestamp": 1}, + id="mooncake_foreign_keys", + ), + param( + { + "id": "x", + "models": ["m"], + "block_size": 64, + "hash_id_scope": "local", + "requests": [], + "prompt": "synthetic-style", + "completion": "x", + }, + id="synthetic_foreign_keys", + ), + ], +) # fmt: skip +def test_near_miss_and_foreign_objects_not_detected(doc: dict): + assert _is_weka_trace_object(doc) is False + + +@pytest.mark.parametrize( + "doc", + [ + param({"id": "x", "models": ["m"], "block_size": 64}, id="missing_keys"), + param("not-a-dict", id="not_a_dict"), + param( + {**_GENUINE_WEKA_OBJECT, "block_size": True}, + id="bool_block_size", + ), + param( + {**_GENUINE_WEKA_OBJECT, "models": "claude"}, + id="models_not_list", + ), + ], +) # fmt: skip +def test_malformed_objects_not_detected(doc: object): + assert _is_weka_trace_object(doc) is False + + +def test_foreign_keyed_object_not_detected_via_file(tmp_path: Path): + # End-to-end through the file sniff: a JSON file carrying the five weka keys + # PLUS foreign keys must NOT route to the graph pipeline. + f = tmp_path / "mooncake.json" + doc = {**_GENUINE_WEKA_OBJECT, "mooncake_field": True, "conversations": []} + f.write_text(json.dumps(doc)) + assert WekaTraceAdapter.can_load(f) is False + + +# --------------------------------------------------------------------------- # +# _looks_like_hf_dataset_id: only weka-marked repo ids match # +# --------------------------------------------------------------------------- # +def test_genuine_weka_hf_id_detected(): + assert _looks_like_hf_dataset_id("semianalysisai/cc-traces-weka-061526") is True + + +@pytest.mark.parametrize( + "repo_id", + [ + param("meta-llama/Llama-3", id="meta_llama"), + param("my-team/notes", id="team_notes"), + param("a/b", id="short"), + param("openai/gpt", id="openai_gpt"), + param("anthropic/claude", id="anthropic_claude"), + ], +) # fmt: skip +def test_arbitrary_org_name_not_detected_as_hf_weka_id(repo_id: str): + assert _looks_like_hf_dataset_id(repo_id) is False + + +def test_existing_parent_dir_not_detected_as_hf_weka_id(tmp_path, monkeypatch): + """A typo'd relative path under an EXISTING local dir is not an HF id. + + ``traces/weka-061526`` (weka-marked, org/name-shaped, non-existent) under a + real ``traces/`` directory is a local-path mistake; routing it to + ``datasets.load_dataset`` buries the typo behind a confusing HF error. + """ + (tmp_path / "traces").mkdir() + monkeypatch.chdir(tmp_path) + assert _looks_like_hf_dataset_id("traces/weka-061526") is False + + +@pytest.mark.parametrize( + "candidate", + [ + param("./traces/weka-061526", id="dot_slash_prefix"), + param("traces/weka-061526/", id="trailing_slash"), + param("a/b/weka-061526", id="multi_component_path"), + ], +) # fmt: skip +def test_pathlike_markers_not_detected_as_hf_weka_id(candidate: str): + assert _looks_like_hf_dataset_id(candidate) is False + + +def test_hf_load_failure_error_names_both_interpretations(monkeypatch): + """When the weka HF heuristic fires and the load fails, the error must + present BOTH readings (missing local path vs bad HF repo id).""" + import sys + import types + + from aiperf.dataset.graph.adapters.weka.trace import ( + WekaTraceAdapterError, + _load_hf_rows, + ) + + def load_dataset(*_args, **_kwargs): # noqa: ANN001 + raise FileNotFoundError("dataset not found on the hub") + + monkeypatch.setitem( + sys.modules, "datasets", types.SimpleNamespace(load_dataset=load_dataset) + ) + + rows = _load_hf_rows("no-such-org/cc-weka-typo", split="train", revision=None) + with pytest.raises(WekaTraceAdapterError) as excinfo: + next(rows) + message = str(excinfo.value) + assert "no such local file or directory" in message + assert "HuggingFace" in message + + +# --------------------------------------------------------------------------- # +# _detect_graph_workload_format: registry-driven detection + native exclusion # +# --------------------------------------------------------------------------- # +def test_detect_graph_workload_format_dynamo(tmp_path) -> None: + from pathlib import Path + + from aiperf.dataset.graph.workload_detect import ( + _detect_graph_workload_format, + ) + + fixture = ( + Path(__file__).resolve().parents[1] + / "dataset/graph/adapters/fixtures/dynamo_nested/nested_2_level.jsonl.gz" + ) + assert _detect_graph_workload_format(fixture) == "dynamo_trace" + + +def test_detect_graph_workload_format_excludes_native(tmp_path) -> None: + # A plain native-looking .jsonl must NOT be auto-detected as a graph + # workload (native is explicit --graph only); returns None so the linear + # pipeline keeps it. + from aiperf.dataset.graph.workload_detect import ( + _detect_graph_workload_format, + ) + + f = tmp_path / "plain.jsonl" + f.write_text('{"messages": [{"role": "user", "content": "hi"}]}\n') + assert _detect_graph_workload_format(f) is None + + +# --------------------------------------------------------------------------- # +# graph_format override: forces graph classification + parse format # +# --------------------------------------------------------------------------- # +def _write_minimal_native_graph(tmp_path: Path) -> Path: + """Write a minimal valid native graph JSONL (one LLM node + one trace). + + Native is auto-detect-EXCLUDED, so this file is only ever treated as a graph + workload when ``--graph-format native`` forces it. + """ + f = tmp_path / "native_min.jsonl" + f.write_text( + '{"kind": "graph", "nodes": ' + '{"a": {"node_type": "llm", ' + '"prompt": [{"role": "user", "content": "hi"}], "output": "out"}}}\n' + '{"kind": "trace", "id": "t1"}\n' + ) + return f + + +@pytest.fixture +def make_run(): + """Build a ``BenchmarkRun`` with a default ``FileDataset`` for the given path. + + Wraps the shared ``make_run_from_cli`` helper so the override tests reuse the + production resolver (CLI ``--input-file`` / ``--graph-format`` -> resolved + ``FileDataset.path`` / ``.graph_format``) instead of hand-building config. + """ + from aiperf.config.flags.cli_config import CLIConfig + from tests.unit.conftest import make_run_from_cli + + def _make(*, path: str, graph_format: str | None = None): + cfg = CLIConfig( + model_names=["test-model"], + input_file=path, + graph_format=graph_format, + ) + return make_run_from_cli(cfg) + + return _make + + +def test_graph_format_override_forces_native(make_run, tmp_path: Path) -> None: + from aiperf.dataset.graph.workload_detect import ( + parse_graph_workload, + resolve_graph_workload, + ) + + native_file = _write_minimal_native_graph(tmp_path) + run = make_run(path=str(native_file), graph_format="native") + # native is auto-detect-EXCLUDED, so only the override makes this a graph + # workload: + ref = resolve_graph_workload(run) + assert ref is not None + assert ref.path == native_file + assert ref.format == "native" + pb = parse_graph_workload(run, native_file) + assert pb.graph is not None + + +def test_graph_format_override_forces_dynamo(make_run) -> None: + from aiperf.dataset.graph.workload_detect import resolve_graph_workload + + fixture = ( + Path(__file__).resolve().parents[1] + / "dataset/graph/adapters/fixtures/dynamo_nested/nested_2_level.jsonl.gz" + ) + run = make_run(path=str(fixture), graph_format="dynamo_trace") + ref = resolve_graph_workload(run) + assert ref is not None + assert ref.path == fixture + assert ref.format == "dynamo_trace" diff --git a/tests/unit/metrics/conftest.py b/tests/unit/metrics/conftest.py index 5068a3f850..84a8b3f6b3 100644 --- a/tests/unit/metrics/conftest.py +++ b/tests/unit/metrics/conftest.py @@ -5,6 +5,7 @@ """ +from aiperf.common.constants import STREAMED_REQUEST_TAG from aiperf.common.enums import ( AggregationKind, CreditPhase, @@ -32,6 +33,16 @@ from aiperf.plugin.enums import EndpointType +def streamed_record_metrics() -> MetricRecordDict: + """A record-metrics dict primed with the streamed-request predicate. + + Streaming record metrics gate on the presence of ``STREAMED_REQUEST_TAG``; + tests that call a streaming metric's ``parse_record`` directly (rather than via + ``run_simple_metrics_pipeline``) must seed it to reach the metric-specific logic. + """ + return MetricRecordDict({STREAMED_REQUEST_TAG: 1}) + + def _create_test_request_info(model_name: str = "test-model") -> RequestInfo: """Create a RequestInfo for testing metrics.""" return RequestInfo( @@ -61,6 +72,7 @@ def create_record( input_tokens: int | None = None, output_tokens_per_response: int = 1, error: ErrorDetails | None = None, + streamed: bool = True, ) -> ParsedResponseRecord: """ Simple helper to create test records with sensible defaults. @@ -70,6 +82,10 @@ def create_record( meaning that the total output token count will be the number of responses times the output tokens per response. The end_perf_ns is the last response timestamp, or the start_ns if no responses are provided. + + ``streamed`` stamps ``RequestRecord.streamed`` (defaults True since these fixtures + simulate streaming runs); set it False to exercise the per-record streaming-metric + skip predicate gated by ``StreamedRequestMetric``. """ responses = responses or [start_ns + 50] # Single response 50ns later @@ -80,6 +96,7 @@ def create_record( timestamp_ns=start_ns, end_perf_ns=responses[-1] if responses else start_ns, error=error, + streamed=streamed, ) response_data = [] diff --git a/tests/unit/metrics/test_inter_chunk_latency_metric.py b/tests/unit/metrics/test_inter_chunk_latency_metric.py index 2a67827ee7..9e4b532b89 100644 --- a/tests/unit/metrics/test_inter_chunk_latency_metric.py +++ b/tests/unit/metrics/test_inter_chunk_latency_metric.py @@ -4,9 +4,12 @@ import pytest from aiperf.common.exceptions import NoMetricValue -from aiperf.metrics.metric_dicts import MetricRecordDict from aiperf.metrics.types.inter_chunk_latency_metric import InterChunkLatencyMetric -from tests.unit.metrics.conftest import create_record, run_simple_metrics_pipeline +from tests.unit.metrics.conftest import ( + create_record, + run_simple_metrics_pipeline, + streamed_record_metrics, +) class TestInterChunkLatencyMetric: @@ -96,7 +99,7 @@ def test_inter_chunk_latency_insufficient_responses(self): NoMetricValue, match="Record must have at least two content responses to calculate Inter Chunk Latency", ): - metric.parse_record(record, MetricRecordDict()) + metric.parse_record(record, streamed_record_metrics()) def test_inter_chunk_latency_empty_responses(self): """Test error when no responses""" @@ -107,7 +110,7 @@ def test_inter_chunk_latency_empty_responses(self): NoMetricValue, match="Record must have at least two content responses to calculate Inter Chunk Latency", ): - metric.parse_record(record, MetricRecordDict()) + metric.parse_record(record, streamed_record_metrics()) def test_inter_chunk_latency_invalid_order(self): """Test error when response timestamps are not in chronological order""" @@ -120,7 +123,7 @@ def test_inter_chunk_latency_invalid_order(self): ValueError, match="Each inter chunk latency must be non-negative. got: .*", ): - metric.parse_record(record, MetricRecordDict()) + metric.parse_record(record, streamed_record_metrics()) @pytest.mark.parametrize( "responses", @@ -139,7 +142,7 @@ def test_inter_chunk_latency_invalid_timestamp_order(self, responses): ValueError, match="Each inter chunk latency must be non-negative. got: .*", ): - metric.parse_record(record, MetricRecordDict()) + metric.parse_record(record, streamed_record_metrics()) def test_inter_chunk_latency_streaming_scenario(self): """Test ICL in a realistic streaming scenario""" @@ -163,7 +166,7 @@ def test_inter_chunk_latency_direct_parse_record(self): """Test direct parse_record method call""" record = create_record(start_ns=100, responses=[110, 125, 140]) metric = InterChunkLatencyMetric() - metric_dict = MetricRecordDict() + metric_dict = streamed_record_metrics() result = metric.parse_record(record, metric_dict) diff --git a/tests/unit/metrics/test_streamed_request_count.py b/tests/unit/metrics/test_streamed_request_count.py new file mode 100644 index 0000000000..d09da62652 --- /dev/null +++ b/tests/unit/metrics/test_streamed_request_count.py @@ -0,0 +1,201 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from pytest import approx, param + +from aiperf.common.constants import STREAMED_REQUEST_COUNT_TAG, STREAMED_REQUEST_TAG +from aiperf.common.enums import MetricConsoleGroup, MetricFlags +from aiperf.common.exceptions import NoMetricValue +from aiperf.metrics.accumulator import MetricsAccumulator +from aiperf.metrics.metric_dicts import MetricRecordDict +from aiperf.metrics.metric_registry import MetricRegistry +from aiperf.metrics.types.inter_chunk_latency_metric import InterChunkLatencyMetric +from aiperf.metrics.types.inter_token_latency_metric import InterTokenLatencyMetric +from aiperf.metrics.types.stream_latency_metrics import StreamSetupLatencyMetric +from aiperf.metrics.types.streamed_request_count_metric import ( + StreamedRequestCountMetric, +) +from aiperf.metrics.types.streamed_request_metric import StreamedRequestMetric +from aiperf.metrics.types.time_to_first_output_token_metric import ( + TimeToFirstOutputTokenMetric, +) +from aiperf.metrics.types.ttft_metric import TTFTMetric +from aiperf.metrics.types.ttst_metric import TTSTMetric +from tests.unit.conftest import make_benchmark_run +from tests.unit.metrics.conftest import create_record, run_simple_metrics_pipeline +from tests.unit.post_processors.conftest import create_metric_records_data + +# Streaming metrics whose _parse_record reads response timing directly, so they +# carry the explicit membership guard + a required_metrics dependency on the +# per-record streaming predicate. +GUARDED_STREAMING_METRICS = [ + TTFTMetric, + TTSTMetric, + InterChunkLatencyMetric, + TimeToFirstOutputTokenMetric, + StreamSetupLatencyMetric, +] + + +class TestStreamedRequestPredicate: + """The hidden per-record predicate that gates streaming metrics.""" + + def test_predicate_parses_streamed_and_skips_non_streamed(self): + """A streamed record parses to 1; a non-streamed record raises NoMetricValue.""" + metric = StreamedRequestMetric() + + streamed = create_record(start_ns=100, responses=[110], streamed=True) + assert metric.parse_record(streamed, MetricRecordDict()) == 1 + + non_streamed = create_record(start_ns=100, responses=[110], streamed=False) + with pytest.raises(NoMetricValue): + metric.parse_record(non_streamed, MetricRecordDict()) + + def test_predicate_only_present_for_streamed_records(self): + """Only streamed records carry a per-record predicate value.""" + records = [ + create_record(start_ns=100, responses=[110], streamed=True), + create_record(start_ns=200, responses=[210], streamed=False), + create_record(start_ns=300, responses=[310], streamed=True), + ] + results = run_simple_metrics_pipeline(records, STREAMED_REQUEST_TAG) + assert results[STREAMED_REQUEST_TAG] == [1, 1] + + def test_predicate_is_console_hidden(self): + """The predicate is excluded from the console (constant-1 gate is not useful).""" + assert StreamedRequestMetric.console_group == MetricConsoleGroup.NONE + + @pytest.mark.asyncio + async def test_predicate_absent_from_summary_export(self): + """The predicate is INTERNAL, so summarize() drops it from exports. + + The accumulator keeps the hidden predicate available for dependent + metrics, then filters the meaningless constant-1 stat row from the + final summary while retaining the visible streamed_request_count. + """ + assert StreamedRequestMetric.has_flags(MetricFlags.INTERNAL) + + accumulator = MetricsAccumulator(make_benchmark_run()) + accumulator._derive_funcs = {} + await accumulator.process_record( + create_metric_records_data( + x_request_id="stream-1", + results=[ + { + StreamedRequestMetric.tag: 1, + StreamedRequestCountMetric.tag: 2, + } + ], + ) + ) + + summary_tags = set((await accumulator.summarize()).results) + assert StreamedRequestMetric.tag not in summary_tags + assert StreamedRequestCountMetric.tag in summary_tags + + +class TestStreamedRequestCountMetric: + """The visible aggregate denominator displayed beside Request Count.""" + + def test_count_no_records(self): + """No records means no aggregate value (not a 0 default).""" + results = run_simple_metrics_pipeline([], STREAMED_REQUEST_COUNT_TAG) + assert STREAMED_REQUEST_COUNT_TAG not in results + + def test_count_aggregates_streamed_and_skips_non_streamed(self): + """The aggregate counts only records that streamed on the wire.""" + records = [ + create_record(start_ns=100, responses=[110], streamed=True), + create_record(start_ns=200, responses=[210], streamed=False), + create_record(start_ns=300, responses=[310], streamed=True), + ] + results = run_simple_metrics_pipeline(records, STREAMED_REQUEST_COUNT_TAG) + assert results[STREAMED_REQUEST_COUNT_TAG] == approx(2) + + def test_count_parse_record_skips_non_streamed(self): + """The aggregate's per-record parse mirrors the streamed predicate.""" + metric = StreamedRequestCountMetric() + + streamed = create_record(start_ns=100, responses=[110], streamed=True) + assert metric.parse_record(streamed, MetricRecordDict()) == 1 + + non_streamed = create_record(start_ns=100, responses=[110], streamed=False) + with pytest.raises(NoMetricValue): + metric.parse_record(non_streamed, MetricRecordDict()) + + @pytest.mark.parametrize("num_streamed", [1, 3, 10, 100]) + def test_count_matches_streamed_record_count(self, num_streamed: int): + """The aggregate equals the number of streamed records.""" + records = [ + create_record(start_ns=100 * i, streamed=True) for i in range(num_streamed) + ] + results = run_simple_metrics_pipeline(records, STREAMED_REQUEST_COUNT_TAG) + assert results[STREAMED_REQUEST_COUNT_TAG] == approx(num_streamed) + + def test_count_is_not_console_hidden(self): + """The visible aggregate uses the default console group (not hidden).""" + assert StreamedRequestCountMetric.console_group == MetricConsoleGroup.DEFAULT + + +class TestStreamingMetricGating: + def test_ttft_skips_non_streamed_record(self): + """A non-streamed record must NOT produce a TTFT value (the pollution bug). + + Pre-fix, TTFT computed ``first content response - start`` even when the sole + TextResponse timestamp was the completion time, reporting full latency as TTFT. + """ + record = create_record(start_ns=100, responses=[150], streamed=False) + + # Full pipeline: TTFT is absent because the predicate never fires. + results = run_simple_metrics_pipeline([record], TTFTMetric.tag) + assert TTFTMetric.tag not in results + + # Inline guard: _parse_record raises when the predicate tag is absent. + with pytest.raises(NoMetricValue): + TTFTMetric()._parse_record(record, MetricRecordDict()) + + def test_ttft_computes_for_streamed_record(self): + """A streamed record yields the unchanged TTFT value.""" + record = create_record(start_ns=100, responses=[110, 120], streamed=True) + + # Full pipeline computes the predicate first, then TTFT (110 - 100 = 10). + results = run_simple_metrics_pipeline([record], TTFTMetric.tag) + assert results[TTFTMetric.tag] == [10] + + # Direct parse with the predicate primed (dependency order). + primed = MetricRecordDict() + primed[STREAMED_REQUEST_TAG] = 1 + assert TTFTMetric().parse_record(record, primed) == 10 + + def test_inter_token_latency_inherits_skip_for_non_streamed(self): + """ITL consumes only guarded metrics, so it inherits the skip transitively.""" + record = create_record(start_ns=100, responses=[110, 120, 130], streamed=False) + results = run_simple_metrics_pipeline([record], InterTokenLatencyMetric.tag) + assert InterTokenLatencyMetric.tag not in results + + @pytest.mark.parametrize( + "metric_cls", + [param(m, id=m.tag) for m in GUARDED_STREAMING_METRICS], + ) # fmt: skip + def test_guarded_metrics_declare_dependency(self, metric_cls): + """Every directly-guarded metric declares the predicate in required_metrics.""" + assert metric_cls.required_metrics is not None + assert STREAMED_REQUEST_TAG in metric_cls.required_metrics + + @pytest.mark.parametrize( + "metric_cls", + [param(m, id=m.tag) for m in GUARDED_STREAMING_METRICS], + ) # fmt: skip + def test_guarded_metric_parse_raises_without_predicate(self, metric_cls): + """Each guarded metric's parse path raises NoMetricValue when the predicate is + absent from record_metrics (the non-streamed record case).""" + record = create_record(start_ns=100, responses=[110, 120], streamed=False) + with pytest.raises(NoMetricValue): + metric_cls()._parse_record(record, MetricRecordDict()) + + def test_dependency_order_places_predicate_before_dependents(self): + """Topological sort computes the predicate before its dependents.""" + order = MetricRegistry.create_dependency_order_for([TTFTMetric.tag]) + assert STREAMED_REQUEST_TAG in order + assert order.index(STREAMED_REQUEST_TAG) < order.index(TTFTMetric.tag) diff --git a/tests/unit/metrics/test_theoretical_prefix_cache_metric.py b/tests/unit/metrics/test_theoretical_prefix_cache_metric.py new file mode 100644 index 0000000000..4171ff18de --- /dev/null +++ b/tests/unit/metrics/test_theoretical_prefix_cache_metric.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from aiperf.common.enums import GenericMetricUnit, MetricConsoleGroup +from aiperf.common.exceptions import NoMetricValue +from aiperf.metrics.metric_dicts import MetricResultsDict +from aiperf.metrics.metric_registry import MetricRegistry +from aiperf.metrics.types.theoretical_prefix_cache_metric import ( + TheoreticalPrefixCacheHitMetric, +) +from aiperf.post_processors.theoretical_prefix_cache import ( + THEORETICAL_PREFIX_CACHE_HIT_TAG, +) + + +class TestTheoreticalPrefixCacheHitMetricRegistration: + """The tag must resolve in the MetricRegistry with CACHE display metadata. + + The value is produced by the standalone TheoreticalPrefixCacheAccumulator, + but display consumers (RealtimeMetricsTable, ConsoleMetricsExporter) resolve + the tag via strict/grouped registry lookups: an unregistered tag used to + kill the realtime dashboard table and be silently omitted from the console + table on every graph-IR run. + """ + + def test_tag_resolves_in_registry(self) -> None: + metric_class = MetricRegistry.get_class("theoretical_prefix_cache_hit") + assert metric_class is TheoreticalPrefixCacheHitMetric + assert metric_class.tag == "theoretical_prefix_cache_hit" + + def test_accumulator_tag_constant_matches_registered_class(self) -> None: + assert TheoreticalPrefixCacheHitMetric.tag == THEORETICAL_PREFIX_CACHE_HIT_TAG + + def test_display_metadata_matches_accumulator_output(self) -> None: + assert TheoreticalPrefixCacheHitMetric.header == "Theoretical Prefix Cache Hit" + assert TheoreticalPrefixCacheHitMetric.unit == GenericMetricUnit.PERCENT + assert TheoreticalPrefixCacheHitMetric.console_group == MetricConsoleGroup.CACHE + assert TheoreticalPrefixCacheHitMetric.display_order is not None + + def test_derive_value_raises_no_metric_value(self) -> None: + """Externally-injected contract: the derivation walk must skip this tag.""" + with pytest.raises(NoMetricValue) as exc_info: + TheoreticalPrefixCacheHitMetric()._derive_value(MetricResultsDict()) + + msg = str(exc_info.value) + assert TheoreticalPrefixCacheHitMetric.tag in msg + assert "MetricResultsDict" in msg + assert "TheoreticalPrefixCacheAccumulator.summarize" in msg diff --git a/tests/unit/metrics/test_time_to_first_output_metric.py b/tests/unit/metrics/test_time_to_first_output_metric.py index 897d3b60a5..e575472259 100644 --- a/tests/unit/metrics/test_time_to_first_output_metric.py +++ b/tests/unit/metrics/test_time_to_first_output_metric.py @@ -15,7 +15,11 @@ from aiperf.metrics.types.time_to_first_output_token_metric import ( TimeToFirstOutputTokenMetric, ) -from tests.unit.metrics.conftest import create_record, run_simple_metrics_pipeline +from tests.unit.metrics.conftest import ( + create_record, + run_simple_metrics_pipeline, + streamed_record_metrics, +) def create_response_record( @@ -39,6 +43,7 @@ def create_response_record( start_perf_ns=start_ns, timestamp_ns=start_ns, end_perf_ns=responses[-1][0], + streamed=True, ) return ParsedResponseRecord( @@ -169,7 +174,7 @@ def test_ttfo_no_valid_content(self, responses): NoMetricValue, match="Record must have at least one non-reasoning token", ): - metric.parse_record(record, MetricRecordDict()) + metric.parse_record(record, streamed_record_metrics()) def test_ttfo_tool_call_response(self): """Test TTFO with a tool-call-only response""" @@ -250,6 +255,7 @@ def _create_record_with_responses( start_perf_ns=start_ns, timestamp_ns=start_ns, end_perf_ns=responses[-1].perf_ns, + streamed=True, ) return ParsedResponseRecord( request=request, diff --git a/tests/unit/metrics/test_ttst_metric.py b/tests/unit/metrics/test_ttst_metric.py index 72b0321fd2..66865a9cd6 100644 --- a/tests/unit/metrics/test_ttst_metric.py +++ b/tests/unit/metrics/test_ttst_metric.py @@ -4,9 +4,12 @@ import pytest from aiperf.common.exceptions import NoMetricValue -from aiperf.metrics.metric_dicts import MetricRecordDict from aiperf.metrics.types.ttst_metric import TTSTMetric -from tests.unit.metrics.conftest import create_record, run_simple_metrics_pipeline +from tests.unit.metrics.conftest import ( + create_record, + run_simple_metrics_pipeline, + streamed_record_metrics, +) class TestTTSTMetric: @@ -57,7 +60,7 @@ def test_ttst_invalid_order(self): ValueError, match="Second response timestamp must be greater than or equal to the first response timestamp", ): - metric.parse_record(record, MetricRecordDict()) + metric.parse_record(record, streamed_record_metrics()) def test_ttst_insufficient_responses(self): """Test error when less than two content responses""" @@ -68,4 +71,4 @@ def test_ttst_insufficient_responses(self): NoMetricValue, match="Record must have at least two content responses to calculate TTST", ): - metric.parse_record(record, MetricRecordDict()) + metric.parse_record(record, streamed_record_metrics()) diff --git a/tests/unit/orchestrator/search_planner/test_monotonic.py b/tests/unit/orchestrator/search_planner/test_monotonic.py index 861aa8d4b6..a3923e2438 100644 --- a/tests/unit/orchestrator/search_planner/test_monotonic.py +++ b/tests/unit/orchestrator/search_planner/test_monotonic.py @@ -467,7 +467,7 @@ def test_mutate_base_preserves_credentials() -> None: def test_mutate_base_preserves_url_userinfo() -> None: - """REGRESSION-LOCK (PR #982 dynamo-ops): URL userinfo survives + """REGRESSION-LOCK: URL userinfo survives ``_mutate_base`` via ``context={"include_secrets": True}``. See ``smooth_isotonic`` for the full rationale. """ diff --git a/tests/unit/orchestrator/search_planner/test_smooth_isotonic_v2_wiring.py b/tests/unit/orchestrator/search_planner/test_smooth_isotonic_v2_wiring.py index 9c8c21676a..df830025d2 100644 --- a/tests/unit/orchestrator/search_planner/test_smooth_isotonic_v2_wiring.py +++ b/tests/unit/orchestrator/search_planner/test_smooth_isotonic_v2_wiring.py @@ -338,9 +338,9 @@ def test_mutate_base_preserves_sensitive_headers() -> None: def test_mutate_base_preserves_url_userinfo() -> None: - """REGRESSION-LOCK (PR #982 dynamo-ops): ``EndpointConfig.urls`` has an + """REGRESSION-LOCK: ``EndpointConfig.urls`` has an unconditional ``_redact_urls`` serializer (no ``when_used="json"`` guard), - so even ``mode="python"`` dumps strip ``user:pass@`` userinfo. The fix + so even ``mode="python"`` dumps strip ``user:pass@`` userinfo. The planner pairs ``mode="python"`` with ``context={"include_secrets": True}`` so the urls serializer's context-aware bypass fires for the planner's in-pipeline dump too. URL-credentialed endpoints (e.g. database URIs, diff --git a/tests/unit/orchestrator/test_cell_callback.py b/tests/unit/orchestrator/test_cell_callback.py index 953ead5085..74ca5276f9 100644 --- a/tests/unit/orchestrator/test_cell_callback.py +++ b/tests/unit/orchestrator/test_cell_callback.py @@ -102,8 +102,8 @@ def _make_executor_returning(results: list[RunResult]) -> MagicMock: """Mocked RunExecutor whose execute() pops successive results.""" iterator = iter(results) executor = MagicMock() - executor.derive_id.side_effect = ( - lambda plan, var_idx, trial: f"id-{var_idx}-{trial}" + executor.derive_id.side_effect = lambda plan, var_idx, trial: ( + f"id-{var_idx}-{trial}" ) executor.execute = AsyncMock(side_effect=lambda _run: next(iterator)) return executor diff --git a/tests/unit/orchestrator/test_local_executor.py b/tests/unit/orchestrator/test_local_executor.py index 97f685d3fc..93bc6fec6c 100644 --- a/tests/unit/orchestrator/test_local_executor.py +++ b/tests/unit/orchestrator/test_local_executor.py @@ -404,7 +404,7 @@ def fake_run_single_benchmark(run: BenchmarkRun) -> None: # --------------------------------------------------------------------------- -# Stale-env isolation (PR #982 review feedback) +# Stale-env isolation # --------------------------------------------------------------------------- @@ -415,14 +415,14 @@ async def test_stale_parent_env_headers_not_forwarded_when_run_has_none( """REGRESSION-LOCK: a stale ``AIPERF_INJECTED_HEADERS`` in the parent shell must not leak into a child run that has no sensitive headers. - Reproduced on PR #982 by ajcasagrande: parent env carried + Repro: parent env carries ``AIPERF_INJECTED_HEADERS={"Authorization":"Bearer stale-parent-secret"}``; - the child run configured only ``X-Trace-Id`` (non-sensitive). Pre-fix - the subprocess overlay still applied the stale Authorization onto + the child run configures only ``X-Trace-Id`` (non-sensitive). Without the + pop, the subprocess overlay applies the stale Authorization onto ``run.cfg.endpoint.headers``, causing the benchmark to send an unintended credential-bearing header. - Fix: ``_run_benchmark_subprocess`` pops both internal injection vars + ``_run_benchmark_subprocess`` therefore pops both internal injection vars from the copied env before conditionally re-setting them. """ monkeypatch.setenv( @@ -457,7 +457,7 @@ async def test_stale_parent_env_headers_not_forwarded_when_run_has_none( # --------------------------------------------------------------------------- -# URL userinfo IPC (PR #982 review feedback — dynamo-ops) +# URL userinfo IPC # --------------------------------------------------------------------------- @@ -570,7 +570,7 @@ def fake_run_single_benchmark(run: BenchmarkRun) -> None: # --------------------------------------------------------------------------- -# Malformed-env-var hardening (PR #982 review feedback — coderabbitai) +# Malformed-env-var hardening # --------------------------------------------------------------------------- @@ -652,7 +652,7 @@ def test_parse_injected_dict_rejects_non_string_values() -> None: def test_parse_injected_dict_malformed_json_raises_value_error_not_decode_error() -> ( None ): - """REGRESSION-LOCK (coderabbitai): malformed env-var JSON must surface as a + """REGRESSION-LOCK: malformed env-var JSON must surface as a ``ValueError`` naming the env var, NOT a raw ``orjson.JSONDecodeError`` that main()'s ``except orjson.JSONDecodeError`` block would misreport as a config-file error. @@ -673,7 +673,7 @@ def test_parse_injected_dict_malformed_json_raises_value_error_not_decode_error( def test_parse_injected_str_list_malformed_json_raises_value_error() -> None: - """REGRESSION-LOCK (coderabbitai): malformed ``AIPERF_INJECTED_ENDPOINT_URLS`` + """REGRESSION-LOCK: malformed ``AIPERF_INJECTED_ENDPOINT_URLS`` JSON surfaces as a ``ValueError`` naming the env var, not a decode error. """ from aiperf.orchestrator.subprocess_runner import _parse_injected_str_list @@ -718,7 +718,7 @@ def test_subprocess_runner_rejects_non_string_header_values_env( def test_subprocess_runner_malformed_env_not_attributed_to_config_file( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] ) -> None: - """REGRESSION-LOCK (coderabbitai): a malformed ``AIPERF_INJECTED_HEADERS`` + """REGRESSION-LOCK: a malformed ``AIPERF_INJECTED_HEADERS`` must not be reported as ``Invalid JSON in config file`` — the config file here is well-formed. The error must name the offending env var instead. """ @@ -750,3 +750,130 @@ def test_subprocess_runner_malformed_env_not_attributed_to_config_file( "malformed env var was misattributed to the config file" ) assert "AIPERF_INJECTED_HEADERS" in err + + +# --------------------------------------------------------------------------- +# C1: scenario lock resolved in the PARENT before the subprocess dump. +# These exercise the REAL executor path (_execute_sync -> _resolve_scenario_in_parent +# -> _prepare_run_artifacts) on a graph-workload run, then prove the on-disk +# run_config.json carries the baked auto-fills so the in-subprocess re-resolution +# (which re-marks every non-None field as "set") is a clean no-op. +# --------------------------------------------------------------------------- + +_WEKA_FIXTURE = (Path(__file__).parents[1] / "graph/fixtures/weka_min.json").resolve() + + +def _graph_scenario_run(tmp_path: Path) -> BenchmarkRun: + cfg = BenchmarkConfig( + models=["claude-opus-4-5-20251101"], + endpoint={"urls": ["http://localhost:8000/v1/chat/completions"]}, + datasets=[{"name": "profiling", "type": "file", "path": str(_WEKA_FIXTURE)}], + phases=[ + { + "name": "profiling", + "type": "concurrency", + "concurrency": 1, + "sessions": 5, + } + ], + scenario="inferencex-agentx-mvp", + ) + return BenchmarkRun( + benchmark_id="scn-mp", + cfg=cfg, + artifact_dir=tmp_path, + label="scenario_mp", + ) + + +@pytest.mark.asyncio +async def test_executor_bakes_scenario_autofills_into_run_config( + tmp_path: Path, +) -> None: + """The executor resolves the scenario in the parent and the dumped + run_config.json carries streaming=True / cache_bust=first_turn_prefix / + duration=1800 -- with NO spurious ScenarioLockError raised.""" + import orjson as _orjson + + run = _graph_scenario_run(tmp_path) + executor = LocalSubprocessExecutor(base_dir=tmp_path) + with patch("aiperf.orchestrator.local_executor.subprocess.run") as mock_run: + mock_run.return_value.returncode = 0 + mock_run.return_value.stderr = "" + await executor.execute(run) # must not raise + + on_disk = _orjson.loads((tmp_path / "run_config.json").read_bytes()) + endpoint = on_disk["cfg"]["endpoint"] + assert endpoint["streaming"] is True + assert endpoint["cache_bust"] == "first_turn_prefix" + phases = on_disk["cfg"]["phases"] + assert any(p.get("duration") == 1800.0 for p in phases) + # The parent's outcome is baked onto resolved state too. + assert on_disk["resolved"]["scenario_outcome"]["submission_valid"] is True + + +@pytest.mark.asyncio +async def test_subprocess_reresolve_of_baked_config_is_clean(tmp_path: Path) -> None: + """REAL subprocess-path proof: load the on-disk run_config.json the executor + wrote and run the scenario validator exactly as the subprocess does -- it + must NOT raise a spurious ScenarioLockError and must stay submission_valid.""" + import orjson as _orjson + + from aiperf.common.scenario import apply_scenario + + run = _graph_scenario_run(tmp_path) + executor = LocalSubprocessExecutor(base_dir=tmp_path) + with patch("aiperf.orchestrator.local_executor.subprocess.run") as mock_run: + mock_run.return_value.returncode = 0 + mock_run.return_value.stderr = "" + await executor.execute(run) + + data = _orjson.loads((tmp_path / "run_config.json").read_bytes()) + child = BenchmarkRun.model_validate(data) + # The round-trip spuriously marks streaming/cache_bust as "set"... + assert "streaming" in child.cfg.endpoint.model_fields_set + # ...but the baked values satisfy the locks, so re-resolution is a no-op. + outcome = apply_scenario(child) + assert outcome.submission_valid is True + assert outcome.violations == [] + + +@pytest.mark.asyncio +async def test_scenario_resolution_performs_no_process_global_writes( + tmp_path: Path, +) -> None: + """C2 regression, config-native form: scenario resolution mutates ONLY the + run's own config (apply-or-lock on ``run.cfg``); it must perform no + ``os.environ`` writes, so nothing can leak into a later plain run. Run A's + own subprocess gets the window via its serialized run config, not env.""" + import os as _os + + executor = LocalSubprocessExecutor(base_dir=tmp_path) + + def _record(*args: Any, **kwargs: Any): + import subprocess as _sp + + return _sp.CompletedProcess(args=args, returncode=0, stdout="", stderr="") + + env_before = dict(_os.environ) + with patch( + "aiperf.orchestrator.local_executor.subprocess.run", side_effect=_record + ): + # Run A: REAL scenario resolution auto-applies the t* window onto the + # run's config before it is serialized for the subprocess. + run_a = _graph_scenario_run(tmp_path / "run_a") + await executor.execute(run_a) + assert run_a.cfg.trajectory_start_max_ratio == 1.0 + + # Run B: plain non-scenario run -- must NOT inherit run A's t* window. + run_b = BenchmarkRun( + benchmark_id="plain", + cfg=_benchmark_config(), + artifact_dir=tmp_path / "run_b", + label="plain", + ) + await executor.execute(run_b) + assert run_b.cfg.trajectory_start_max_ratio is None + + # Scenario application wrote no process-global env state. + assert dict(_os.environ) == env_before diff --git a/tests/unit/plot/test_dashboard_utils.py b/tests/unit/plot/test_dashboard_utils.py index c7b329f867..88dcf9a1fc 100644 --- a/tests/unit/plot/test_dashboard_utils.py +++ b/tests/unit/plot/test_dashboard_utils.py @@ -404,8 +404,8 @@ def test_excludes_runs_with_missing_metrics(self): run.metadata.experiment_group = "group_a" if i == 1: # Make middle run missing y metric - run.get_metric = ( - lambda name: {"p50": 100} if name == "throughput" else None + run.get_metric = lambda name: ( + {"p50": 100} if name == "throughput" else None ) else: run.get_metric = lambda name: {"p50": 100} diff --git a/tests/unit/plugin/test_graph_adapter_category.py b/tests/unit/plugin/test_graph_adapter_category.py new file mode 100644 index 0000000000..04e671acfd --- /dev/null +++ b/tests/unit/plugin/test_graph_adapter_category.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from aiperf.plugin import plugins +from aiperf.plugin.enums import PluginType + + +def test_graph_adapter_category_registered() -> None: + assert hasattr(PluginType, "GRAPH_ADAPTER") + names = {e.name for e in plugins.iter_entries(PluginType.GRAPH_ADAPTER)} + assert {"native", "weka_trace", "dynamo_trace"} <= names + + +def test_graph_adapter_classes_resolve() -> None: + from aiperf.dataset.graph.adapters.dynamo.trace import DynamoTraceAdapter + from aiperf.dataset.graph.adapters.weka.trace import WekaTraceAdapter + + assert ( + plugins.get_class(PluginType.GRAPH_ADAPTER, "dynamo_trace") + is DynamoTraceAdapter + ) + assert plugins.get_class(PluginType.GRAPH_ADAPTER, "weka_trace") is WekaTraceAdapter + + +def test_detect_format_resolves_dynamo_fixture() -> None: + from pathlib import Path + + from aiperf.dataset.graph.parser import detect_format + + fixture = ( + Path(__file__).resolve().parents[1] + / "dataset/graph/adapters/fixtures/dynamo_nested/nested_2_level.jsonl.gz" + ) + assert detect_format(fixture) == "dynamo_trace" diff --git a/tests/unit/plugin/test_session_routing_registry.py b/tests/unit/plugin/test_session_routing_registry.py new file mode 100644 index 0000000000..1d2a2a809c --- /dev/null +++ b/tests/unit/plugin/test_session_routing_registry.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from pytest import param + +from aiperf.plugin import plugins +from aiperf.plugin.enums import PluginType +from aiperf.workers.session_routing import SessionRoutingBase + + +@pytest.mark.parametrize( + "name", + [ + param("dynamo_headers", id="dynamo_headers"), + param("dynamo_nvext", id="dynamo_nvext"), + param("smg_routing_key", id="smg_routing_key"), + param("session_id_header", id="session_id_header"), + ], +) # fmt: skip +def test_session_routing_plugins_resolve(name): + cls = plugins.get_class(PluginType.SESSION_ROUTING, name) + assert issubclass(cls, SessionRoutingBase) + + +def test_session_routing_enum_generated(): + from aiperf.plugin.enums import SessionRoutingType + + assert SessionRoutingType.DYNAMO_HEADERS == "dynamo_headers" + assert SessionRoutingType.SESSION_ID_HEADER == "session_id_header" diff --git a/tests/unit/post_processors/test_streaming_gate.py b/tests/unit/post_processors/test_streaming_gate.py new file mode 100644 index 0000000000..fbedc652e0 --- /dev/null +++ b/tests/unit/post_processors/test_streaming_gate.py @@ -0,0 +1,123 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run-level STREAMING_ONLY gate relaxation for graph-IR workloads. + +``BaseMetricsProcessor.get_filters`` disables every ``STREAMING_ONLY`` metric +(TTFT / TTST / ICL / TTFOT / stream-setup) when the global ``endpoint.streaming`` +flag is off. Graph-IR replays (weka / dynamo / native graph) stream per-request +from recorded node modes even without the global flag, so the run-level gate +must NOT fire for them -- per-record applicability is enforced by the +``streamed_request`` predicate metric instead. + +Every run here is a REAL resolved ``BenchmarkRun`` (native ``BenchmarkConfig`` / +the production ``resolve_config`` path), never MagicMock, so a wrong attribute +path in the graph-workload predicate cannot be silently auto-created (the +repo's MagicMock-path-drift trap). +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from aiperf.common.enums import MetricFlags +from aiperf.post_processors.base_metrics_processor import BaseMetricsProcessor + + +def _write_minimal_native_graph(tmp_path: Path) -> Path: + """Write a minimal valid native graph JSONL (one LLM node + one trace). + + Native is auto-detect-EXCLUDED, so this file is only treated as a graph + workload when ``--graph-format native`` forces it -- exactly the resolved + ``FileDataset.graph_format`` the predicate reads. + """ + f = tmp_path / "native_min.jsonl" + f.write_text( + '{"kind": "graph", "nodes": ' + '{"a": {"node_type": "llm", ' + '"prompt": [{"role": "user", "content": "hi"}], "output": "out"}}}\n' + '{"kind": "trace", "id": "t1"}\n' + ) + return f + + +def _graph_run(tmp_path: Path): + """Real resolved graph-workload ``BenchmarkRun`` (native --graph-format).""" + from aiperf.config.flags.cli_config import CLIConfig + from tests.unit.conftest import make_run_from_cli + + cfg = CLIConfig( + model_names=["test-model"], + input_file=str(_write_minimal_native_graph(tmp_path)), + graph_format="native", + ) + return make_run_from_cli(cfg) + + +class TestStreamingOnlyGate: + """Gate matrix for MetricFlags.STREAMING_ONLY over (global flag, workload).""" + + def test_global_on_non_graph_allows_streaming(self, mock_run) -> None: + # Synthetic (non-graph) dataset with the global flag on: unchanged. + mock_run.cfg.endpoint.streaming = True + _, disallowed = BaseMetricsProcessor(mock_run).get_filters() + assert not (disallowed & MetricFlags.STREAMING_ONLY) + + def test_global_off_non_graph_disallows_streaming(self, mock_run) -> None: + # Synthetic (non-graph) dataset with the global flag off: the gate still + # fires (nothing in the run can stream). Unchanged behavior. + mock_run.cfg.endpoint.streaming = False + _, disallowed = BaseMetricsProcessor(mock_run).get_filters() + assert disallowed & MetricFlags.STREAMING_ONLY + + def test_global_on_graph_allows_streaming(self, tmp_path) -> None: + run = _graph_run(tmp_path) + run.cfg.endpoint.streaming = True + _, disallowed = BaseMetricsProcessor(run).get_filters() + assert not (disallowed & MetricFlags.STREAMING_ONLY) + + def test_global_off_graph_allows_streaming(self, tmp_path) -> None: + # NEW: graph-IR workloads stream per-request, so the run-level gate must + # not fire even with the global flag off. + run = _graph_run(tmp_path) + run.cfg.endpoint.streaming = False + _, disallowed = BaseMetricsProcessor(run).get_filters() + assert not (disallowed & MetricFlags.STREAMING_ONLY) + + def test_predicate_true_for_graph_false_for_synthetic( + self, mock_run, tmp_path + ) -> None: + # Guards against MagicMock-path drift: the predicate reads real resolved + # config, so it must distinguish an actual graph workload from a + # synthetic one rather than auto-creating a truthy attribute path. + assert BaseMetricsProcessor(mock_run)._is_graph_workload() is False + assert BaseMetricsProcessor(_graph_run(tmp_path))._is_graph_workload() is True + + +class TestEmptyStreamedSubset: + """A graph run whose records are all non-streamed leaves TTFT absent.""" + + @pytest.mark.asyncio + async def test_all_non_streamed_records_leave_ttft_absent(self, tmp_path) -> None: + """With the gate relaxed, TTFT is applicable but the per-record streamed + predicate never fires, so no record carries the TTFT tag; the accumulator + must leave TTFT absent without raising or warning.""" + from aiperf.metrics.accumulator import MetricsAccumulator + from aiperf.metrics.types.request_latency_metric import RequestLatencyMetric + from aiperf.metrics.types.ttft_metric import TTFTMetric + from tests.unit.post_processors.conftest import create_metric_records_data + + accumulator = MetricsAccumulator(_graph_run(tmp_path)) + + for i in range(3): + msg = create_metric_records_data( + x_request_id=f"r-{i}", + request_start_ns=1_000_000_000 + i, + results=[{RequestLatencyMetric.tag: 42.0}], + ) + await accumulator.process_record(msg) + + summary = await accumulator.summarize() + assert RequestLatencyMetric.tag in summary.results + assert TTFTMetric.tag not in summary.results diff --git a/tests/unit/post_processors/test_theoretical_prefix_cache.py b/tests/unit/post_processors/test_theoretical_prefix_cache.py new file mode 100644 index 0000000000..54dfdd05f4 --- /dev/null +++ b/tests/unit/post_processors/test_theoretical_prefix_cache.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for theoretical prefix-cache accounting. + +Covers the WEKA graph-IR pre-pass that stamps per-turn hit/total block counts +and the ``TheoreticalPrefixCacheAccumulator`` results processor that emits the +cumulative ``theoretical_prefix_cache_hit`` percent. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from aiperf.common.enums import CreditPhase, GenericMetricUnit +from aiperf.common.models.dataset_models import ( + ConversationMetadata, + DatasetMetadata, + GraphDatasetMetadata, +) +from aiperf.plugin.enums import DatasetSamplingStrategy +from aiperf.post_processors.theoretical_prefix_cache import ( + THEORETICAL_PREFIX_CACHE_HIT_TAG, + TheoreticalPrefixCacheAccumulator, +) +from tests.unit.post_processors.conftest import ( + create_metric_metadata, + create_metric_records_data, +) + +# weka_min.json: a 3-turn linear chain, hash_ids [1,2] / [1,2,3] / [1,2,3,4], +# times 0.0 / 1.5 / 3.0, block_size 64, hash_id_scope local. Over ONE shared +# per-trace seen-set consumed in time order: +# turn 0: hit 0 / total 2 (seen={} -> {1,2}) +# turn 1: hit 2 / total 3 ([1,2] hit, 3 miss -> {1,2,3}) +# turn 2: hit 3 / total 4 ([1,2,3] hit, 4 miss -> {1,2,3,4}) +# cumulative: hit 5 / total 9 -> 100 * 5/9 = 55.5555...% +_WEKA_MIN = Path(__file__).parents[1] / "graph" / "fixtures" / "weka_min.json" +_EXPECTED_PER_TURN = [(0, 2), (2, 3), (3, 4)] +_EXPECTED_HIT_PCT = 100.0 * 5 / 9 + + +@pytest.mark.asyncio +async def test_accumulator_emits_cumulative_hit_rate(mock_run) -> None: + """on_dataset_configured + process_record + summarize -> hit-rate percent.""" + acc = TheoreticalPrefixCacheAccumulator(run=mock_run, service_id="proc-1") + + conv_id = "trace_03_n3" + from aiperf.common.models.dataset_models import TurnMetadata + + conv_meta = ConversationMetadata(conversation_id=conv_id) + conv_meta.turns = [ + TurnMetadata( + theoretical_prefix_cache_hit_blocks=h, + theoretical_prefix_cache_total_blocks=t, + ) + for h, t in _EXPECTED_PER_TURN + ] + metadata = DatasetMetadata( + conversations=[conv_meta], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) + acc.on_dataset_configured(metadata) + + for turn_index in range(len(_EXPECTED_PER_TURN)): + record = create_metric_records_data( + metadata=create_metric_metadata( + conversation_id=conv_id, + turn_index=turn_index, + benchmark_phase=CreditPhase.PROFILING, + ) + ) + await acc.process_record(record) + + results = await acc.summarize() + assert len(results) == 1 + result = results[0] + assert result.tag == THEORETICAL_PREFIX_CACHE_HIT_TAG + assert result.unit == str(GenericMetricUnit.PERCENT) + assert result.avg == pytest.approx(_EXPECTED_HIT_PCT) + assert result.sum == 5 + assert result.count == 9 + # Finite-discipline: the emitted value must be a finite percent in [0, 100]. + assert 0.0 <= result.avg <= 100.0 + + +@pytest.mark.asyncio +async def test_accumulator_node_map_path_matches_graph_dispatch(mock_run) -> None: + """Graph path: per-node map + node_id recovered from (conversation_id, turn_index). + + Mirrors the real WEKA graph dispatch where the record carries LEGACY-shaped + identity: ``conversation_id`` is the trajectory TEMPLATE id (root scope == + the trace id) and ``turn_index`` is the node's 0-based turn, so the + ``{scope}:{turn}`` node id -- and thus the join into the per-node map -- is a + pure function of those two record fields (no correlation-id parsing). + """ + acc = TheoreticalPrefixCacheAccumulator(run=mock_run, service_id="proc-1") + base_trace = "trace_03_n3" + acc.on_dataset_configured( + DatasetMetadata( + conversations=[], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + graph=GraphDatasetMetadata( + trace_ids=[base_trace], + prefix_cache_by_trace={ + base_trace: { + f"{base_trace}:0": [0, 2], + f"{base_trace}:1": [2, 3], + f"{base_trace}:2": [3, 4], + } + }, + ), + ) + ) + for turn_index in range(3): + record = create_metric_records_data( + metadata=create_metric_metadata( + # conversation_id is the root-scope TEMPLATE id (instance + # identity rides x_correlation_id, which this join never reads). + conversation_id=base_trace, + turn_index=turn_index, + benchmark_phase=CreditPhase.PROFILING, + ) + ) + await acc.process_record(record) + results = await acc.summarize() + assert len(results) == 1 + assert results[0].avg == pytest.approx(_EXPECTED_HIT_PCT) + assert results[0].sum == 5 + assert results[0].count == 9 + + +@pytest.mark.asyncio +async def test_accumulator_no_metadata_emits_nothing(mock_run) -> None: + """No prefix-cache metadata (non-weka dataset) -> empty summary, no NaN.""" + acc = TheoreticalPrefixCacheAccumulator(run=mock_run, service_id="proc-1") + metadata = DatasetMetadata( + conversations=[ConversationMetadata(conversation_id="c0")], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) + acc.on_dataset_configured(metadata) + record = create_metric_records_data( + metadata=create_metric_metadata( + conversation_id="c0", turn_index=0, benchmark_phase=CreditPhase.PROFILING + ) + ) + await acc.process_record(record) + assert await acc.summarize() == [] + + +@pytest.mark.asyncio +async def test_accumulator_skips_context_overflow_records(mock_run) -> None: + """WK4: an overflow-skip record (valid=True, error trimmed by the + RecordsManager) must NOT fold its planned blocks into the hit rate. + + The trimmed carrier keeps ``metadata.context_overflow_skip=True``; the + accumulator keys off that flag. An identical record without the flag still + accumulates (control). + """ + acc = TheoreticalPrefixCacheAccumulator(run=mock_run, service_id="proc-1") + base_trace = "trace_03_n3" + acc.on_dataset_configured( + DatasetMetadata( + conversations=[], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + graph=GraphDatasetMetadata( + trace_ids=[base_trace], + prefix_cache_by_trace={base_trace: {f"{base_trace}:0": [3, 4]}}, + ), + ) + ) + + def _record(overflow_skip: bool): + metadata = create_metric_metadata( + # conversation_id is the root-scope TEMPLATE id; turn_index 0 recovers + # node id ``{base_trace}:0``. + conversation_id=base_trace, + turn_index=0, + benchmark_phase=CreditPhase.PROFILING, + ) + metadata.context_overflow_skip = overflow_skip + return create_metric_records_data(metadata=metadata) + + overflow_record = _record(overflow_skip=True) + assert overflow_record.valid # trimmed carrier arrives error-free + await acc.process_record(overflow_record) + assert await acc.summarize() == [] + + await acc.process_record(_record(overflow_skip=False)) + results = await acc.summarize() + assert results[0].sum == 3 + assert results[0].count == 4 + + +@pytest.mark.asyncio +async def test_accumulator_clamps_overcounted_hits(mock_run) -> None: + """A loader miscount (hit > total) is clamped so the rate stays <= 100%.""" + acc = TheoreticalPrefixCacheAccumulator(run=mock_run, service_id="proc-1") + from aiperf.common.models.dataset_models import TurnMetadata + + conv_meta = ConversationMetadata(conversation_id="c0") + conv_meta.turns = [ + TurnMetadata( + theoretical_prefix_cache_hit_blocks=10, + theoretical_prefix_cache_total_blocks=4, + ) + ] + acc.on_dataset_configured( + DatasetMetadata( + conversations=[conv_meta], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) + ) + record = create_metric_records_data( + metadata=create_metric_metadata( + conversation_id="c0", turn_index=0, benchmark_phase=CreditPhase.PROFILING + ) + ) + await acc.process_record(record) + results = await acc.summarize() + assert results[0].avg == pytest.approx(100.0) diff --git a/tests/unit/property/_numeric_bounds_baseline.txt b/tests/unit/property/_numeric_bounds_baseline.txt index 76d6c11669..80cb3ae21b 100644 --- a/tests/unit/property/_numeric_bounds_baseline.txt +++ b/tests/unit/property/_numeric_bounds_baseline.txt @@ -424,3 +424,24 @@ WorkerStatusSummaryMessage.request_ns: numeric field has no ge/gt/le/lt bound an WorkerTaskStats.completed: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. WorkerTaskStats.failed: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. WorkerTaskStats.total: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +AgentReplayMetrics.input_length: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +AgentReplayMetrics.input_sequence_hashes: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +AgentReplayMetrics.trace_block_size: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +AgentRequestMetrics.avg_itl_ms: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +AgentRequestMetrics.cached_tokens: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +AgentRequestMetrics.input_tokens: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +AgentRequestMetrics.kv_hit_rate: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +AgentRequestMetrics.kv_transfer_estimated_latency_ms: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +AgentRequestMetrics.output_tokens: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +AgentRequestMetrics.prefill_time_ms: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +AgentRequestMetrics.prefill_wait_time_ms: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +AgentRequestMetrics.queue_depth: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +AgentRequestMetrics.request_received_ms: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +AgentRequestMetrics.total_time_ms: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +AgentRequestMetrics.ttft_ms: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +AgentToolEvent.duration_ms: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +AgentTraceRecord.event_time_unix_ms: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +WorkerInfo.decode_dp_rank: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +WorkerInfo.decode_worker_id: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +WorkerInfo.prefill_dp_rank: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +WorkerInfo.prefill_worker_id: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. diff --git a/tests/unit/property/test_finite_invariants.py b/tests/unit/property/test_finite_invariants.py index 9c24e8b629..9cad70d2b7 100644 --- a/tests/unit/property/test_finite_invariants.py +++ b/tests/unit/property/test_finite_invariants.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """Mechanical "global invariant" tests for the NaN/inf discipline. -Three CI-enforceable contracts that codify the round-1 finite-float +Three CI-enforceable contracts that codify the finite-float remediation work into rules a future PR can't accidentally regress: 1. Every JSON exporter that calls ``orjson.dumps`` on metric-bearing @@ -353,6 +353,11 @@ def _load_or_init_metric_baseline(current: list[str]) -> list[str]: # AdaptiveSearchSweep.outcome_constraints: list[OutcomeConstraint], not a # numeric field. Per-element OutcomeConstraint.bound is already FiniteFloat. "AdaptiveSearchSweep.outcome_constraints", + # GraphDatasetMetadata.prefix_cache_by_trace: dict[str, dict[str, list[int]]] + # of [theoretical_hit_blocks, total_blocks] block counts, not a numeric + # field. The bare-int substring check trips on the nested list[int]; a + # field-level ge/gt/le/lt bound cannot apply to a container. + "GraphDatasetMetadata.prefix_cache_by_trace", # ServerMetricsResults.warmup_endpoint_summaries: dict of summary models, # not a numeric leaf — same shape as the baselined endpoint_summaries # sibling; per-summary numeric fields carry their own bounds. diff --git a/tests/unit/property/test_pydantic_field_fuzz.py b/tests/unit/property/test_pydantic_field_fuzz.py index bc306613dd..46d64e45fb 100644 --- a/tests/unit/property/test_pydantic_field_fuzz.py +++ b/tests/unit/property/test_pydantic_field_fuzz.py @@ -10,7 +10,7 @@ ``RecursionError``, ``IndexError`` -- means a validator crashed on the adversarial input rather than rejecting it cleanly. -This protects against the class of bugs found in round-2 (e.g. NaN +This protects against a recurring class of bugs (e.g. NaN silently passing a finite-bounds check, an unhashable choices entry crashing inside ``set()`` build, the ``mean``-only ambiguous-distribution path raising the wrong error type). diff --git a/tests/unit/records/test_records_manager.py b/tests/unit/records/test_records_manager.py index 5fcef15349..47f724fa69 100644 --- a/tests/unit/records/test_records_manager.py +++ b/tests/unit/records/test_records_manager.py @@ -835,3 +835,110 @@ async def _raise_timeout(coro, *args, **kwargs): published = manager.publish.await_args.args[0] assert isinstance(published, BaseServiceErrorMessage) manager._dispatch_record.assert_not_called() + + +# --------------------------------------------------------------------------- +# context-overflow record metric-exclusion (graph-IR) +# --------------------------------------------------------------------------- + + +def _overflow_skip_message( + overflow_count: int = 1, + phase: CreditPhase = CreditPhase.PROFILING, +) -> RecordsMessage: + metadata = MetricRecordMetadata( + session_num=3, + conversation_id="conv-overflow", + turn_index=2, + request_start_ns=1_000_000_000, + request_end_ns=1_010_000_000, + worker_id="worker-overflow", + record_processor_id="rp-1", + benchmark_phase=phase, + context_overflow_skip=True, + ) + return RecordsMessage( + service_id="rp-1", + metadata=metadata, + records=[ + MetricRecordsData( + metadata=metadata, + metrics={ + "context_overflow_count": overflow_count, + "request_latency": 9_999_000_000, + }, + ) + ], + error=ErrorDetails(message="context_length_exceeded", code=400), + ) + + +def _create_real_tracker_manager() -> RecordsManager: + manager = RecordsManager.__new__(RecordsManager) + manager._records_tracker = RecordsTracker() + manager._error_tracker = ErrorTracker() + manager._skipped_context_overflow_count = 0 + manager._complete_credit_phases = set() + manager._dataset_configured_event = asyncio.Event() + manager._dataset_configured_event.set() + manager.info = MagicMock() + manager.debug = MagicMock() + manager.trace = MagicMock() + manager.warning = MagicMock() + manager.is_enabled_for = MagicMock(return_value=False) + manager._handle_all_records_received = AsyncMock() + forwarded: list[MetricRecordsData] = [] + + async def _capture(record, **kwargs) -> list[BaseException]: + forwarded.append(record) + return [] + + manager._dispatch_record = AsyncMock(side_effect=_capture) + manager._forwarded = forwarded # type: ignore[attr-defined] + return manager + + +class TestRecordsManagerOverflowExclusion: + """Graph-IR context-overflow records skip perf accumulation + the error + tracker, but still advance the success counter and forward ONLY the + context_overflow_count metric (submission-rate gate).""" + + @pytest.mark.asyncio + async def test_overflow_skip_counts_success_not_error(self) -> None: + manager = _create_real_tracker_manager() + await manager._on_records(_overflow_skip_message()) + + tracker = manager._records_tracker._get_phase_tracker(CreditPhase.PROFILING) + assert tracker._success_records == 1 + assert tracker._error_records == 0 + assert manager._skipped_context_overflow_count == 1 + summary = manager._error_tracker.get_error_summary_for_phase( + CreditPhase.PROFILING + ) + assert sum(e.count for e in summary) == 0 + + @pytest.mark.asyncio + async def test_overflow_skip_forwards_only_overflow_count_metric(self) -> None: + manager = _create_real_tracker_manager() + await manager._on_records(_overflow_skip_message(overflow_count=1)) + + forwarded = manager._forwarded # type: ignore[attr-defined] + assert len(forwarded) == 1 + trimmed = forwarded[0] + assert trimmed.metrics == {"context_overflow_count": 1} + assert "request_latency" not in trimmed.metrics + assert trimmed.error is None + assert trimmed.metadata.context_overflow_skip is True + assert trimmed.valid is True + + @pytest.mark.asyncio + async def test_non_overflow_record_unchanged(self) -> None: + manager = _create_real_tracker_manager() + await manager._on_records(_metric_records_message()) + + forwarded = manager._forwarded # type: ignore[attr-defined] + assert len(forwarded) == 1 + assert forwarded[0].metrics == {"request_latency": 250_000_000} + assert manager._skipped_context_overflow_count == 0 + tracker = manager._records_tracker._get_phase_tracker(CreditPhase.PROFILING) + assert tracker._success_records == 1 diff --git a/tests/unit/scenario/test_context_overflow.py b/tests/unit/scenario/test_context_overflow.py new file mode 100644 index 0000000000..b6d99cd846 --- /dev/null +++ b/tests/unit/scenario/test_context_overflow.py @@ -0,0 +1,242 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the InferenceX AgentX context-overflow submission contract. + +Covers the runtime classifier (``is_context_overflow_response``), the +``context_overflow_count`` aggregate metric, and the +``compute_submission_outcome`` fold that combines the static scenario-lock +outcome with the runtime overflow rate. Uses REAL config objects (no MagicMock) +so attribute-path drift on Environment / RequestRecord / ScenarioOutcome cannot +hide behind auto-created mock attributes. +""" + +from __future__ import annotations + +import pytest +from pytest import param + +from aiperf.common.environment import Environment +from aiperf.common.scenario import ( + CONTEXT_OVERFLOW_REASON, + compute_submission_outcome, + is_context_overflow_response, +) + +# --------------------------------------------------------------------------- +# Classifier +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "body,expected", + [ + param("This exceeds the maximum context length of 8192", True, id="raw_substr"), + param("error: prompt is too long for the model", True, id="raw_prompt_long"), + param('{"error": {"message": "context_length_exceeded"}}', True, id="openai"), + param('{"error": "maximum context reached"}', True, id="openai_str_error"), + param('{"detail": "context length exceeded"}', True, id="vllm_detail_raw"), + param("Internal server error", False, id="normal_error"), + param('{"error": {"message": "rate limit exceeded"}}', False, id="other_429"), + param("", False, id="empty"), + param(None, False, id="none"), + ], +) # fmt: skip +def test_is_context_overflow_response_classifies(body, expected): + """Default allowlist matches overflow bodies and rejects normal errors.""" + assert is_context_overflow_response(body=body) is expected + + +def test_is_context_overflow_response_case_insensitive(): + """Matching is case-insensitive against both body and nested message.""" + assert is_context_overflow_response(body="MAXIMUM CONTEXT length") is True + assert ( + is_context_overflow_response(body='{"error": {"message": "Context Length"}}') + is True + ) + + +def test_is_context_overflow_response_bytes_body(): + """Bytes bodies are decoded (utf-8, replace) before matching.""" + assert is_context_overflow_response(body=b"maximum context exceeded") is True + + +def test_is_context_overflow_response_explicit_substrings_override(): + """An explicit substrings arg overrides the env allowlist entirely.""" + assert ( + is_context_overflow_response( + body="custom overflow marker", substrings=["custom overflow"] + ) + is True + ) + # Default allowlist would not match this body. + assert is_context_overflow_response(body="custom overflow marker") is False + + +def test_is_context_overflow_response_empty_allowlist_disables(monkeypatch): + """An empty allowlist disables runtime detection (always False).""" + monkeypatch.setattr(Environment.AGENTX, "CONTEXT_OVERFLOW_SUBSTRINGS", []) + assert is_context_overflow_response(body="maximum context length") is False + + +def test_is_context_overflow_response_env_extension(monkeypatch): + """Extending the env allowlist enables a new server vocabulary.""" + monkeypatch.setattr( + Environment.AGENTX, + "CONTEXT_OVERFLOW_SUBSTRINGS", + ["token limit reached"], + ) + assert is_context_overflow_response(body="token limit reached") is True + assert is_context_overflow_response(body="maximum context length") is False + + +# --------------------------------------------------------------------------- +# compute_submission_outcome -- runtime rate fold +# --------------------------------------------------------------------------- + + +def test_compute_submission_outcome_over_threshold_flips_invalid(): + """> threshold overflow rate flips submission_valid=False + reason.""" + valid, reasons = compute_submission_outcome( + scenario_name="s", + validator_submission_valid=True, + total_responses=100, + context_overflow_count=2, + ) + assert valid is False + assert CONTEXT_OVERFLOW_REASON in reasons + + +def test_compute_submission_outcome_under_threshold_unaffected(): + """< threshold overflow rate leaves the lock outcome unaffected.""" + valid, reasons = compute_submission_outcome( + scenario_name="s", + validator_submission_valid=True, + total_responses=1000, + context_overflow_count=5, + ) + assert valid is True + assert reasons == [] + + +def test_compute_submission_outcome_boundary_equal_accepted(): + """Rate exactly equal to the limit is accepted (strict greater-than).""" + # 1/100 == 0.01 == default limit -> accepted. + valid, reasons = compute_submission_outcome( + scenario_name="s", + validator_submission_valid=True, + total_responses=100, + context_overflow_count=1, + ) + assert valid is True + assert reasons == [] + + +def test_compute_submission_outcome_zero_responses_no_flip(): + """Zero total responses treats the rate as 0 -- no overflow flip.""" + valid, reasons = compute_submission_outcome( + scenario_name="s", + validator_submission_valid=True, + total_responses=0, + context_overflow_count=0, + ) + assert valid is True + assert reasons == [] + + +def test_compute_submission_outcome_preserves_lock_reasons(): + """The lock-only outcome (unsafe_override) is preserved and merged.""" + valid, reasons = compute_submission_outcome( + scenario_name="s", + validator_submission_valid=False, + validator_reasons=["unsafe_override"], + total_responses=100, + context_overflow_count=5, + ) + assert valid is False + assert "unsafe_override" in reasons + assert CONTEXT_OVERFLOW_REASON in reasons + + +def test_compute_submission_outcome_lock_only_under_threshold_stays_invalid(): + """Lock-invalid + under-threshold overflow stays invalid, no overflow reason.""" + valid, reasons = compute_submission_outcome( + scenario_name="s", + validator_submission_valid=False, + validator_reasons=["unsafe_override"], + total_responses=100, + context_overflow_count=0, + ) + assert valid is False + assert reasons == ["unsafe_override"] + + +def test_compute_submission_outcome_no_scenario_returns_none(): + """No scenario -> (None, []); callers drop the field.""" + valid, reasons = compute_submission_outcome( + scenario_name=None, + validator_submission_valid=None, + total_responses=100, + context_overflow_count=50, + ) + assert valid is None + assert reasons == [] + + +def test_compute_submission_outcome_cancelled_flips_invalid(): + """A cancelled run is never a valid submission.""" + valid, reasons = compute_submission_outcome( + scenario_name="s", + validator_submission_valid=True, + total_responses=100, + context_overflow_count=0, + was_cancelled=True, + ) + assert valid is False + assert "run_cancelled" in reasons + + +def test_compute_submission_outcome_threshold_override(monkeypatch): + """A raised env threshold tolerates a higher overflow rate.""" + monkeypatch.setattr(Environment.AGENTX, "CONTEXT_OVERFLOW_RATE_LIMIT", 0.5) + # 10% overflow now under the 50% limit. + valid, reasons = compute_submission_outcome( + scenario_name="s", + validator_submission_valid=True, + total_responses=100, + context_overflow_count=10, + ) + assert valid is True + assert reasons == [] + + +# --------------------------------------------------------------------------- +# context_overflow_count metric +# --------------------------------------------------------------------------- + + +def test_context_overflow_count_metric_registered(): + """The metric registers under its tag with ERROR_ONLY flag.""" + from aiperf.common.enums import MetricFlags + from aiperf.metrics import MetricRegistry + + cls = MetricRegistry.get_class("context_overflow_count") + assert cls is not None + assert cls.tag == "context_overflow_count" + assert MetricFlags.ERROR_ONLY in cls.flags + assert MetricFlags.NO_INDIVIDUAL_RECORDS in cls.flags + + +def test_context_overflow_count_metric_counts_flagged_records(): + """_parse_record returns 1 only when request.context_overflow is True.""" + from types import SimpleNamespace + + from aiperf.metrics.metric_dicts import MetricRecordDict + from aiperf.metrics.types.context_overflow_count_metric import ( + ContextOverflowCountMetric, + ) + + metric = ContextOverflowCountMetric() + overflow_rec = SimpleNamespace(request=SimpleNamespace(context_overflow=True)) + normal_rec = SimpleNamespace(request=SimpleNamespace(context_overflow=False)) + assert metric._parse_record(overflow_rec, MetricRecordDict()) == 1 + assert metric._parse_record(normal_rec, MetricRecordDict()) == 0 diff --git a/tests/unit/scenario/test_env_locks.py b/tests/unit/scenario/test_env_locks.py new file mode 100644 index 0000000000..6038672836 --- /dev/null +++ b/tests/unit/scenario/test_env_locks.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Trajectory-window apply-or-lock (config-native) and ScenarioSpec doc contract (C3). + +``apply_trajectory_ratios`` operates purely on the run config — no +process-global state: unset ``cfg.trajectory_start_min/max_ratio`` fields are +auto-applied from the spec; user-explicit values are locked with a violation +on mismatch. These tests build a REAL ``BenchmarkConfig`` (no MagicMock) so +the ``model_fields_set``-based explicitness check is exercised for real. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest +from pytest import param + +from aiperf.common.scenario import get_scenario +from aiperf.common.scenario._env_locks import apply_trajectory_ratios +from aiperf.common.scenario.base import ScenarioSpec, ScenarioViolation +from aiperf.config import BenchmarkConfig +from aiperf.plugin.enums import TimingMode + + +def _cfg(**explicit: Any) -> BenchmarkConfig: + """Minimal real BenchmarkConfig; ``explicit`` kwargs mark fields user-set.""" + return BenchmarkConfig( + models=["m"], + endpoint={"urls": ["http://localhost:8000/v1/chat/completions"]}, + datasets=[ + { + "name": "profiling", + "type": "synthetic", + "entries": 5, + "prompts": {"isl": 32}, + } + ], + phases=[ + { + "name": "profiling", + "type": "concurrency", + "concurrency": 1, + "sessions": 5, + } + ], + **explicit, + ) + + +def _run(**explicit: Any) -> SimpleNamespace: + # apply_trajectory_ratios reads only ``run.cfg``. + return SimpleNamespace(cfg=_cfg(**explicit)) + + +def _spec() -> ScenarioSpec: + """Minimal real spec carrying a t* window to auto-apply (no MagicMock).""" + return ScenarioSpec( + name="lock-test", + timing_mode=TimingMode.GRAPH_IR, + require_ignore_eos=False, + forbid_input_truncation=False, + require_loader="weka_trace", + min_benchmark_duration_seconds=1, + default_trajectory_start_min_ratio=0.0, + default_trajectory_start_max_ratio=1.0, + ) + + +def test_trajectory_ratios_auto_applied_when_unset() -> None: + """Unset window fields are auto-applied onto the live config.""" + run = _run() + violations: list[ScenarioViolation] = [] + applied: list[str] = [] + apply_trajectory_ratios(run, _spec(), violations, applied) + assert violations == [] + assert "trajectory_start_ratios" in applied + assert run.cfg.trajectory_start_min_ratio == 0.0 + assert run.cfg.trajectory_start_max_ratio == 1.0 + + +def test_trajectory_ratios_mvp_spec_values_via_registry() -> None: + """The real inferencex-agentx-mvp spec applies its 0.0..1.0 window.""" + run = _run() + violations: list[ScenarioViolation] = [] + applied: list[str] = [] + apply_trajectory_ratios( + run, get_scenario("inferencex-agentx-mvp"), violations, applied + ) + assert violations == [] + assert (run.cfg.trajectory_start_min_ratio, run.cfg.trajectory_start_max_ratio) == ( + 0.0, + 1.0, + ) + + +def test_trajectory_ratios_explicit_conflict_violates() -> None: + """A user-explicit flag value differing from the spec is a violation and + is never overwritten; the matching explicit sibling passes.""" + run = _run(trajectory_start_min_ratio=0.10, trajectory_start_max_ratio=1.0) + violations: list[ScenarioViolation] = [] + applied: list[str] = [] + apply_trajectory_ratios(run, _spec(), violations, applied) + assert [v.flag for v in violations] == ["--trajectory-start-min-ratio"] + assert "trajectory_start_ratios" not in applied + assert run.cfg.trajectory_start_min_ratio == 0.10 + + +def test_trajectory_ratio_violation_still_applies_unset_sibling() -> None: + """A violated bound never blocks its unset sibling's auto-apply: under + --unsafe-override the run proceeds with the mixed window, so the sibling + write is load-bearing to pin.""" + run = _run(trajectory_start_max_ratio=0.9) + violations: list[ScenarioViolation] = [] + applied: list[str] = [] + apply_trajectory_ratios(run, _spec(), violations, applied) + assert [v.flag for v in violations] == ["--trajectory-start-max-ratio"] + assert run.cfg.trajectory_start_max_ratio == 0.9 # explicit, never rewritten + assert run.cfg.trajectory_start_min_ratio == 0.0 # unset sibling applied + + +def test_trajectory_ratios_explicit_match_passes() -> None: + """User-explicit values equal to the spec pass as applied.""" + run = _run(trajectory_start_min_ratio=0.0, trajectory_start_max_ratio=1.0) + violations: list[ScenarioViolation] = [] + applied: list[str] = [] + apply_trajectory_ratios(run, _spec(), violations, applied) + assert violations == [] + assert "trajectory_start_ratios" in applied + + +def test_inverted_window_rejected_at_config_validation() -> None: + """min > max is a config-level error, before any scenario lock runs.""" + with pytest.raises(ValueError, match="trajectory_start_min_ratio"): + _cfg(trajectory_start_min_ratio=0.5) + + +# --------------------------------------------------------------------------- +# C3: trajectory-ratio Field descriptions must state the real lock contract +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "field_name, flag", + [ + param( + "default_trajectory_start_min_ratio", + "--trajectory-start-min-ratio", + id="min_ratio", + ), + param( + "default_trajectory_start_max_ratio", + "--trajectory-start-max-ratio", + id="max_ratio", + ), + ], +) # fmt: skip +def test_trajectory_ratio_description_states_lock_contract( + field_name: str, flag: str +) -> None: + """The docs must name the enforced flag and the lock error, and must not + reference the retired AIPERF_GRAPH_START_* env vars.""" + description = ScenarioSpec.model_fields[field_name].description + assert description is not None + assert flag in description + assert "ScenarioLockError" in description + assert "AIPERF_GRAPH_START" not in description diff --git a/tests/unit/scenario/test_scenario_registry.py b/tests/unit/scenario/test_scenario_registry.py new file mode 100644 index 0000000000..b6d5d89c95 --- /dev/null +++ b/tests/unit/scenario/test_scenario_registry.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the scenario invariant-lock registry and the inferencex-agentx-mvp spec.""" + +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from aiperf.common.enums import CacheBustTarget +from aiperf.common.scenario import ( + INFERENCEX_AGENTX_MVP, + SCENARIOS, + ScenarioSpec, + UnknownScenarioError, + get_scenario, +) +from aiperf.plugin.enums import TimingMode + + +def test_get_scenario_returns_inferencex_agentx_mvp_spec() -> None: + spec = get_scenario("inferencex-agentx-mvp") + assert spec is INFERENCEX_AGENTX_MVP + assert isinstance(spec, ScenarioSpec) + assert spec.name == "inferencex-agentx-mvp" + + +def test_inferencex_agentx_mvp_field_values() -> None: + spec = get_scenario("inferencex-agentx-mvp") + assert spec.timing_mode == TimingMode.GRAPH_IR + assert spec.require_ignore_eos is True + assert spec.require_streaming is True + assert spec.forbid_input_truncation is True + assert spec.require_loader == ( + "semianalysis_cc_traces_weka_with_subagents", + "semianalysis_cc_traces_weka_with_subagents_256k", + "semianalysis_cc_traces_weka_with_subagents_060226", + "semianalysis_cc_traces_weka_with_subagents_060226_256k", + "semianalysis_cc_traces_weka_with_subagents_060526", + "semianalysis_cc_traces_weka_with_subagents_060526_256k", + "semianalysis_cc_traces_weka_with_subagents_060826", + "semianalysis_cc_traces_weka_with_subagents_060826_256k", + "semianalysis_cc_traces_weka_061326", + "semianalysis_cc_traces_weka_061326_256k", + "semianalysis_cc_traces_weka_061526", + "semianalysis_cc_traces_weka_061526_256k", + "semianalysis_cc_traces_weka_062126", + "semianalysis_cc_traces_weka_062126_256k", + "weka_trace", + "weka_hf", + ) + assert spec.min_benchmark_duration_seconds == 900 + assert spec.default_benchmark_duration_seconds == 1800 + assert spec.default_trajectory_start_min_ratio == 0.0 + assert spec.default_trajectory_start_max_ratio == 1.0 + assert spec.trace_idle_gap_cap_seconds == 10.0 + assert spec.require_cache_bust == CacheBustTarget.FIRST_TURN_PREFIX + + +def test_get_scenario_unknown_name_raises_listing_valid_names() -> None: + with pytest.raises(UnknownScenarioError) as exc_info: + get_scenario("nope") + message = str(exc_info.value) + assert "nope" in message + assert "inferencex-agentx-mvp" in message + + +def test_registry_contains_inferencex_agentx_mvp() -> None: + assert SCENARIOS["inferencex-agentx-mvp"] is INFERENCEX_AGENTX_MVP + + +def test_scenario_spec_is_frozen() -> None: + spec = get_scenario("inferencex-agentx-mvp") + with pytest.raises(ValidationError): + spec.name = "mutated" diff --git a/tests/unit/scenario/test_scenario_validator.py b/tests/unit/scenario/test_scenario_validator.py new file mode 100644 index 0000000000..5350a6cd2c --- /dev/null +++ b/tests/unit/scenario/test_scenario_validator.py @@ -0,0 +1,429 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the graph-IR scenario invariant-lock validator (apply_scenario). + +These build REAL ``BenchmarkRun`` objects (NOT MagicMock) against a real weka +graph-workload fixture: MagicMock auto-creates any attribute path and would +hide real-config drift (e.g. a renamed ``endpoint.cache_bust`` field or a +non-graph dataset slipping past ``resolve_graph_workload``). +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from aiperf.common.enums import CacheBustTarget +from aiperf.common.scenario import ScenarioLockError, apply_scenario +from aiperf.config import BenchmarkConfig +from aiperf.config.resolution.plan import BenchmarkRun + +_WEKA_FIXTURE = ( + Path(__file__).parents[2] / "unit/graph/fixtures/weka_min.json" +).resolve() +_SCENARIO = "inferencex-agentx-mvp" +_MODEL = "claude-opus-4-5-20251101" + + +def _make_graph_run( + *, + scenario: str | None = _SCENARIO, + unsafe_override: bool = False, + dataset_synthesis: dict | None = None, + trajectory_start_min_ratio: float | None = None, + trajectory_start_max_ratio: float | None = None, + **endpoint, +) -> BenchmarkRun: + """Build a real graph-workload ``BenchmarkRun`` wrapping the weka fixture. + + The profiling phase uses ``sessions`` (not ``duration``) so the duration + auto-fill path is exercised. ``endpoint`` kwargs are passed through so a + test can set an explicit ``streaming`` / ``cache_bust`` to trigger a + violation. ``dataset_synthesis`` merges a ``synthesis`` block onto the weka + dataset (e.g. to set ``idle_gap_cap_seconds``). The trajectory kwargs mark + the top-level window fields user-explicit (only when passed). + """ + dataset: dict = {"name": "profiling", "type": "file", "path": str(_WEKA_FIXTURE)} + if dataset_synthesis is not None: + dataset["synthesis"] = dataset_synthesis + top_level: dict = {} + if trajectory_start_min_ratio is not None: + top_level["trajectory_start_min_ratio"] = trajectory_start_min_ratio + if trajectory_start_max_ratio is not None: + top_level["trajectory_start_max_ratio"] = trajectory_start_max_ratio + cfg = BenchmarkConfig( + models=[_MODEL], + endpoint={ + "urls": ["http://localhost:8000/v1/chat/completions"], + **endpoint, + }, + datasets=[dataset], + phases=[ + { + "name": "profiling", + "type": "concurrency", + "concurrency": 1, + "sessions": 5, + } + ], + scenario=scenario, + unsafe_override=unsafe_override, + **top_level, + ) + return BenchmarkRun(benchmark_id="test-run", cfg=cfg, artifact_dir=Path("/tmp/x")) + + +def _make_synthetic_run() -> BenchmarkRun: + """Build a real NON-graph (synthetic) run with the scenario set.""" + cfg = BenchmarkConfig( + models=["m"], + endpoint={"urls": ["http://localhost:8000/v1/chat/completions"]}, + datasets=[ + { + "name": "profiling", + "type": "synthetic", + "entries": 5, + "prompts": {"isl": 32}, + } + ], + phases=[ + { + "name": "profiling", + "type": "concurrency", + "concurrency": 1, + "sessions": 5, + } + ], + scenario=_SCENARIO, + unsafe_override=True, + ) + return BenchmarkRun(benchmark_id="test-run", cfg=cfg, artifact_dir=Path("/tmp/x")) + + +def _flags(run: BenchmarkRun) -> list[str]: + """Apply the scenario (override active) and return the violation flags.""" + outcome = apply_scenario(run) + return [v.flag for v in outcome.violations] + + +# Scenario application is config-native (apply-or-lock on run.cfg) and +# performs no process-global writes, so nothing can leak across tests. + + +# --------------------------------------------------------------------------- +# No-scenario short-circuit +# --------------------------------------------------------------------------- + + +def test_no_scenario_is_noop_outcome() -> None: + run = _make_graph_run(scenario=None) + outcome = apply_scenario(run) + assert outcome.scenario_name is None + assert outcome.submission_valid is None + assert run.resolved.scenario_outcome is outcome + + +# --------------------------------------------------------------------------- +# Per-_apply_* auto-fills (unset -> filled, no violation) +# --------------------------------------------------------------------------- + + +def test_streaming_autoenabled_when_unset() -> None: + run = _make_graph_run() + apply_scenario(run) + assert run.cfg.endpoint.streaming is True + + +def test_ignore_eos_injected_when_absent() -> None: + run = _make_graph_run() + apply_scenario(run) + assert run.cfg.endpoint.extra["ignore_eos"] is True + + +def test_cache_bust_autofilled_when_default() -> None: + run = _make_graph_run() + apply_scenario(run) + assert run.cfg.endpoint.cache_bust == CacheBustTarget.FIRST_TURN_PREFIX + + +def test_duration_autofilled_when_unset() -> None: + run = _make_graph_run() + apply_scenario(run) + assert run.cfg.get_profiling_phases()[0].duration == 1800.0 + + +def test_timing_mode_passes_for_graph_workload() -> None: + run = _make_graph_run() + outcome = apply_scenario(run) + assert "timing_mode" in outcome.applied_locks + + +# --------------------------------------------------------------------------- +# Per-_apply_* violations (explicit conflict) +# --------------------------------------------------------------------------- + + +def test_streaming_explicit_false_violates() -> None: + run = _make_graph_run(unsafe_override=True, streaming=False) + assert "--streaming" in _flags(run) + + +def test_cache_bust_explicit_conflict_violates() -> None: + run = _make_graph_run(unsafe_override=True, cache_bust="none") + assert "--cache-bust" in _flags(run) + + +def test_non_graph_workload_violates_timing_mode_and_loader() -> None: + run = _make_synthetic_run() + flags = _flags(run) + assert "--input-file (timing_mode)" in flags + assert "--input-file (loader)" in flags + + +# --------------------------------------------------------------------------- +# Config apply-or-lock (explicit flag values conflicting with the spec) +# --------------------------------------------------------------------------- + + +def test_trajectory_start_min_ratio_flag_override_violates() -> None: + # Under apply-or-lock semantics an override is "explicit" only when the + # config field is user-set; unset would simply be auto-applied. + run = _make_graph_run( + unsafe_override=True, + trajectory_start_min_ratio=0.10, + trajectory_start_max_ratio=1.0, + ) + assert "--trajectory-start-min-ratio" in _flags(run) + + +def test_trace_idle_gap_cap_config_override_violates() -> None: + # Scenario locks the cap to 10.0; a run setting --synthesis-idle-gap-cap=30 + # via config must violate. + run = _make_graph_run( + unsafe_override=True, dataset_synthesis={"idle_gap_cap_seconds": 30.0} + ) + assert "--synthesis-idle-gap-cap" in _flags(run) + + +def test_trace_idle_gap_cap_explicit_match_passes() -> None: + # An explicit --synthesis-idle-gap-cap equal to the spec (10.0) locks clean. + run = _make_graph_run(dataset_synthesis={"idle_gap_cap_seconds": 10.0}) + outcome = apply_scenario(run) + assert "trace_idle_gap_cap" in outcome.applied_locks + assert outcome.violations == [] + + +def test_env_defaults_satisfy_scenario_no_violation() -> None: + # Sanity: with a bare run the t* window (0.0/1.0) and idle-gap cap (10.0) + # are auto-applied, not violated. + run = _make_graph_run() + outcome = apply_scenario(run) + assert "trajectory_start_ratios" in outcome.applied_locks + assert "trace_idle_gap_cap" in outcome.applied_locks + # The cap landed on the live config (unset -> scenario auto-applies 10.0), + # overriding the bare 60s adapter default. + dataset = run.cfg.get_default_dataset() + assert dataset.synthesis.idle_gap_cap_seconds == 10.0 + + +# --------------------------------------------------------------------------- +# End-to-end: all auto-fills + submission_valid; conflict raises; override +# --------------------------------------------------------------------------- + + +def test_end_to_end_all_autofills_submission_valid() -> None: + run = _make_graph_run() + outcome = apply_scenario(run) + assert outcome.submission_valid is True + assert outcome.scenario_name == _SCENARIO + assert outcome.violations == [] + # Every auto-fill landed on the live config. + assert run.cfg.endpoint.cache_bust == CacheBustTarget.FIRST_TURN_PREFIX + assert run.cfg.endpoint.streaming is True + assert run.cfg.endpoint.extra["ignore_eos"] is True + assert run.cfg.get_profiling_phases()[0].duration == 1800.0 + assert run.resolved.scenario_outcome is outcome + + +def test_end_to_end_explicit_conflict_raises_lock_error() -> None: + run = _make_graph_run(streaming=False) + with pytest.raises(ScenarioLockError): + apply_scenario(run) + + +def test_end_to_end_unsafe_override_downgrades_and_stamps_invalid() -> None: + run = _make_graph_run(unsafe_override=True, streaming=False) + outcome = apply_scenario(run) + assert outcome.submission_valid is False + assert outcome.submission_invalid_reasons == ["unsafe_override"] + assert any(v.flag == "--streaming" for v in outcome.violations) + assert run.resolved.scenario_outcome is outcome + + +# --------------------------------------------------------------------------- +# C1 regression: model_fields_set is stale after the multi-run subprocess +# round-trip (model_dump(exclude_none=True) -> model_validate re-marks every +# non-None field as set). Resolving in the PARENT before the dump must bake the +# auto-fills so the in-subprocess re-resolution is a clean no-op (no spurious +# ScenarioLockError on streaming / cache_bust). +# --------------------------------------------------------------------------- + + +def _roundtrip(run: BenchmarkRun) -> BenchmarkRun: + """Mirror the orchestrator->subprocess hop (local_executor + subprocess_runner).""" + import orjson + + data = run.model_dump(mode="json", exclude_none=True) + return BenchmarkRun.model_validate(orjson.loads(orjson.dumps(data))) + + +def test_roundtrip_without_parent_resolve_marks_defaults_as_set() -> None: + """Documents the C1 trap: the round-trip makes endpoint.model_fields_set + spuriously contain streaming + cache_bust even when the user set neither.""" + run = _make_graph_run() + assert run.cfg.endpoint.model_fields_set == {"urls"} + run2 = _roundtrip(run) + assert "streaming" in run2.cfg.endpoint.model_fields_set + assert "cache_bust" in run2.cfg.endpoint.model_fields_set + + +def test_parent_resolve_then_roundtrip_then_resolve_no_spurious_violation() -> None: + """C1 FIX: resolve in the parent (faithful model_fields_set), dump+reload, + then re-resolve in the 'subprocess' -- must NOT raise and stay valid.""" + from aiperf.orchestrator.local_executor import _resolve_scenario_in_parent + + parent = _make_graph_run() + _resolve_scenario_in_parent(parent) + assert parent.resolved.scenario_outcome.submission_valid is True + + child = _roundtrip(parent) + # The parent baked the auto-fills; the child sees satisfying values. + assert child.cfg.endpoint.streaming is True + assert child.cfg.endpoint.cache_bust == CacheBustTarget.FIRST_TURN_PREFIX + assert child.cfg.get_profiling_phases()[0].duration == 1800.0 + + # Subprocess re-resolution is a clean no-op (no spurious ScenarioLockError). + outcome = apply_scenario(child) + assert outcome.submission_valid is True + assert outcome.violations == [] + + +def test_roundtrip_without_parent_resolve_raises_spurious_lock_error() -> None: + """Proves the bug exists absent the parent-resolve: a default scenario run + that round-trips and is THEN resolved (as the subprocess does) raises.""" + run = _make_graph_run() + child = _roundtrip(run) + with pytest.raises(ScenarioLockError): + apply_scenario(child) + + +def test_parent_resolve_noop_when_no_scenario() -> None: + from aiperf.orchestrator.local_executor import _resolve_scenario_in_parent + + run = _make_graph_run(scenario=None) + _resolve_scenario_in_parent(run) + assert run.resolved.scenario_outcome is None + # No auto-fill applied when no scenario. + assert run.cfg.endpoint.cache_bust == CacheBustTarget.NONE + + +# --------------------------------------------------------------------------- +# Gap I1: concurrency-sweep rejection (a swept concurrency multiplies the +# locked config into N diverging runs; AgentX rejects a list-shaped concurrency) +# --------------------------------------------------------------------------- + + +def _make_swept_graph_run(values: dict) -> BenchmarkRun: + """A graph run whose SweepVariation records a swept dotted-path key.""" + from aiperf.config.sweep import SweepVariation + + run = _make_graph_run(unsafe_override=True) + run.variation = SweepVariation(index=0, label="swept", values=values) + return run + + +def test_concurrency_sweep_under_scenario_violates() -> None: + run = _make_swept_graph_run({"phases.profiling.concurrency": 20}) + assert "--concurrency" in _flags(run) + + +def test_prefill_concurrency_sweep_under_scenario_violates() -> None: + run = _make_swept_graph_run({"phases.profiling.prefill_concurrency": 4}) + assert "--concurrency" in _flags(run) + + +def test_non_concurrency_sweep_key_does_not_violate() -> None: + # A request-rate sweep key is not a concurrency sweep -> no concurrency flag. + run = _make_swept_graph_run({"phases.profiling.request_rate": 5.0}) + assert "--concurrency" not in _flags(run) + + +def test_single_run_no_variation_does_not_violate_concurrency() -> None: + run = _make_graph_run() + outcome = apply_scenario(run) + assert all(v.flag != "--concurrency" for v in outcome.violations) + + +# --------------------------------------------------------------------------- +# Gap I3: require_loader full entry list + canonical weka HF-repo pin +# --------------------------------------------------------------------------- + + +def _make_hf_graph_run(hf_id: str, *, unsafe_override: bool = True) -> BenchmarkRun: + """A graph run whose --input-file is a HuggingFace dataset id string.""" + cfg = BenchmarkConfig( + models=[_MODEL], + endpoint={"urls": ["http://localhost:8000/v1/chat/completions"]}, + datasets=[{"name": "profiling", "type": "file", "path": hf_id}], + phases=[ + { + "name": "profiling", + "type": "concurrency", + "concurrency": 1, + "sessions": 5, + } + ], + scenario=_SCENARIO, + unsafe_override=unsafe_override, + ) + return BenchmarkRun(benchmark_id="test-run", cfg=cfg, artifact_dir=Path("/tmp/x")) + + +def test_canonical_weka_hf_repo_passes_loader_and_pin() -> None: + run = _make_hf_graph_run("semianalysisai/cc-traces-weka-062126") + outcome = apply_scenario(run) + assert "require_loader" in outcome.applied_locks + assert all( + v.flag not in ("--input-file (loader)", "--input-file (hf-repo)") + for v in outcome.violations + ) + + +def test_canonical_weka_hf_repo_256k_variant_passes() -> None: + run = _make_hf_graph_run("semianalysisai/cc-traces-weka-062126-256k") + outcome = apply_scenario(run) + assert "require_loader" in outcome.applied_locks + + +def test_local_weka_file_workload_passes_loader() -> None: + # The local-file fixture is a weka graph workload -> loader OK, no HF pin. + run = _make_graph_run() + outcome = apply_scenario(run) + assert "require_loader" in outcome.applied_locks + + +def test_foreign_hf_repo_with_weka_marker_violates_pin() -> None: + # Carries the "weka" marker (so it sniffs as a graph workload) but is NOT + # under the canonical SemiAnalysis prefix -> submission_invalid. + run = _make_hf_graph_run("someorg/my-weka-traces") + flags = _flags(run) + assert "--input-file (hf-repo)" in flags + assert "require_loader" not in apply_scenario(run).applied_locks + + +def test_non_weka_hf_repo_violates_loader() -> None: + # No "weka" marker -> not a graph workload at all -> loader violation. + run = _make_hf_graph_run("meta-llama/Llama-3") + flags = _flags(run) + assert "--input-file (loader)" in flags diff --git a/tests/unit/search_recipes/test_recipes_round3.py b/tests/unit/search_recipes/test_recipes_round3.py index eaaa48792f..912ffc8675 100644 --- a/tests/unit/search_recipes/test_recipes_round3.py +++ b/tests/unit/search_recipes/test_recipes_round3.py @@ -1,9 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Round-3 adversarial regressions: NaN/inf discipline in BO + search recipes. +"""Adversarial regressions: NaN/inf discipline in BO + search recipes. -Covers the behaviours hardened in this round: +Covers: - NaN metric values are treated as "missing" by the Optuna BO planner and logged once per planner instance. diff --git a/tests/unit/search_recipes/test_recipes_round3_followup.py b/tests/unit/search_recipes/test_recipes_round3_followup.py index 562d69cd89..5935bfdf24 100644 --- a/tests/unit/search_recipes/test_recipes_round3_followup.py +++ b/tests/unit/search_recipes/test_recipes_round3_followup.py @@ -1,12 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Round-3 follow-up regressions for the search-recipe post-processors. +"""Baseline-validation regressions for the search-recipe post-processors. -This file tracks the fixes applied after the initial round-3 cleanup landed. -The original round-3 commit ``ab5f07a16`` taught :class:`DegradationKneeDetect` -to reject negative and non-finite baselines, but the strictly-zero baseline was -still accepted. A zero baseline produces ``cutoff = 0 * (1 + threshold) = 0`` so +:class:`DegradationKneeDetect` +rejects negative and non-finite baselines; a strictly-zero baseline must be +rejected too. A zero baseline produces ``cutoff = 0 * (1 + threshold) = 0`` so any positive value is flagged as an "infinite" degradation, which is a meaningless threshold; reject it up front instead. """ diff --git a/tests/unit/server_metrics/test_server_metrics_data_collector.py b/tests/unit/server_metrics/test_server_metrics_data_collector.py index ea29e7ed4d..be8d20791b 100644 --- a/tests/unit/server_metrics/test_server_metrics_data_collector.py +++ b/tests/unit/server_metrics/test_server_metrics_data_collector.py @@ -554,7 +554,7 @@ def test_realistic_sglang_payload_with_nan_gauge_drops_only_offender(self): - The resulting record can be constructed without raising Regression test for the silent-loss bug observed against sglang --enable-metrics where sglang:fwd_occupancy emitted NaN, extended to - cover the parallel histogram-path filter (Task 1b).""" + cover the parallel histogram-path filter.""" metrics_text = """# HELP sglang:fwd_occupancy Forward pass GPU occupancy percentage. # TYPE sglang:fwd_occupancy gauge sglang:fwd_occupancy{engine_type="unified",model_name="Qwen/Qwen3-0.6B",moe_ep_rank="0",pp_rank="0",tp_rank="0"} NaN diff --git a/tests/unit/test_cli_runner_sweep_helpers.py b/tests/unit/test_cli_runner_sweep_helpers.py index bb0f33528d..0e49592655 100644 --- a/tests/unit/test_cli_runner_sweep_helpers.py +++ b/tests/unit/test_cli_runner_sweep_helpers.py @@ -386,11 +386,11 @@ async def test_aggregate_per_variation_writes_single_run_in_degraded_mode( ): """Single-trial cells get a degraded-mode aggregate, mirroring sweep aggregate. - Per round-2 R2-L10: per-variation and sweep aggregation paths must + Per-variation and sweep aggregation paths must use the same gating. ``ConfidenceAggregation`` has a documented single-run mode (std=0, CI collapsed to mean, ``single_run: True``); - skipping it here used to produce dangling references when the sweep - summary still listed the cell. + skipping it here would produce dangling references when the sweep + summary still lists the cell. """ plan = _make_plan_mode(SweepMode.INDEPENDENT) r = _result("c10", concurrency=10, throughput=100.0, ttft_p99=50.0) diff --git a/tests/unit/timing/conftest.py b/tests/unit/timing/conftest.py index 4a93fb624d..405d337070 100644 --- a/tests/unit/timing/conftest.py +++ b/tests/unit/timing/conftest.py @@ -51,10 +51,14 @@ async def _async_true(*args, **kwargs) -> bool: class MockCreditRouter: sent_credits: list[Credit] = field(default_factory=list) auto_return: bool = False + ended_graph_traces: list[tuple[str, str, str]] = field(default_factory=list) _return_cb: Callable[[str, CreditReturn], Awaitable[None]] | None = None _first_token_cb: Callable[[FirstToken], Awaitable[None]] | None = None _pending: list[asyncio.Task] = field(default_factory=list) + async def end_graph_trace(self, trace_id: str, phase_variant: str) -> None: + self.ended_graph_traces.append((trace_id, phase_variant)) + async def send_credit(self, credit: Credit) -> None: self.sent_credits.append(credit) if self.auto_return and self._return_cb: @@ -500,11 +504,15 @@ class MockCreditSender: def __init__(self) -> None: self.sent_credits: list[Credit] = [] self.cancelled = False + self.ended_graph_traces: list[tuple[str, str, str]] = [] self._cb: Callable[[str, CreditReturn], Awaitable[None]] | None = None async def send_credit(self, credit: Credit) -> None: self.sent_credits.append(credit) + async def end_graph_trace(self, trace_id: str, phase_variant: str) -> None: + self.ended_graph_traces.append((trace_id, phase_variant)) + async def cancel_all_credits(self) -> None: self.cancelled = True diff --git a/tests/unit/timing/phase/test_runner.py b/tests/unit/timing/phase/test_runner.py index 326614a781..ab22e46507 100644 --- a/tests/unit/timing/phase/test_runner.py +++ b/tests/unit/timing/phase/test_runner.py @@ -14,6 +14,7 @@ DatasetMetadata, TurnMetadata, ) +from aiperf.common.scenario import TrajectoryWarmupFailedError from aiperf.credit.sticky_router import StickyCreditRouter from aiperf.credit.structs import Credit from aiperf.plugin.enums import ArrivalPattern, DatasetSamplingStrategy, TimingMode @@ -583,6 +584,311 @@ async def test_waits_for_returns_on_final_phase( pub.publish_phase_complete.assert_called_once() +@dataclass +class TeardownStrategy(MockStrategy): + """MockStrategy exposing the optional ``teardown_phase`` hook (TC2).""" + + teardown_calls: int = 0 + + async def teardown_phase(self) -> None: + self.teardown_calls += 1 + + +class TestStrategyTeardown: + """TC2: PhaseRunner invokes the strategy's optional ``teardown_phase``.""" + + async def test_run_invokes_teardown_after_phase_completes( + self, + conv_src: MagicMock, + pub: MagicMock, + router: MagicMock, + conc: MagicMock, + cancel: MagicMock, + cb: MagicMock, + ) -> None: + r = make_runner(cfg(), conv_src, pub, router, conc, cancel, cb) + strategy = TeardownStrategy() + with patch( + "aiperf.timing.phase.runner.plugins.get_class", + return_value=lambda **kw: strategy, + ): + r._progress.all_credits_sent_event.set() + r._progress.all_credits_returned_event.set() + await r.run(is_final_phase=True) + assert strategy.teardown_calls == 1 + + async def test_run_invokes_teardown_on_failure_path( + self, + conv_src: MagicMock, + pub: MagicMock, + router: MagicMock, + conc: MagicMock, + cancel: MagicMock, + cb: MagicMock, + ) -> None: + """Teardown runs in the finally even when the phase setup blows up.""" + + @dataclass + class FailingSetupStrategy(TeardownStrategy): + async def setup_phase(self) -> None: + raise RuntimeError("setup failed") + + r = make_runner(cfg(), conv_src, pub, router, conc, cancel, cb) + strategy = FailingSetupStrategy() + with ( + patch( + "aiperf.timing.phase.runner.plugins.get_class", + return_value=lambda **kw: strategy, + ), + pytest.raises(RuntimeError, match="setup failed"), + ): + await r.run(is_final_phase=True) + assert strategy.teardown_calls == 1 + + async def test_run_skips_strategies_without_teardown_hook( + self, + conv_src: MagicMock, + pub: MagicMock, + router: MagicMock, + conc: MagicMock, + cancel: MagicMock, + cb: MagicMock, + ) -> None: + """Linear strategies (no ``teardown_phase``) are unaffected (no-op).""" + r = make_runner(cfg(), conv_src, pub, router, conc, cancel, cb) + strategy = MockStrategy() + with patch( + "aiperf.timing.phase.runner.plugins.get_class", + return_value=lambda **kw: strategy, + ): + r._progress.all_credits_sent_event.set() + r._progress.all_credits_returned_event.set() + await r.run(is_final_phase=True) # must not raise on the missing hook + + async def test_seamless_non_final_defers_teardown_until_returns_land( + self, + conv_src: MagicMock, + pub: MagicMock, + router: MagicMock, + conc: MagicMock, + cancel: MagicMock, + cb: MagicMock, + ) -> None: + """Seamless phases tear down only once the background return wait ends. + + Distinguishes DEFERRED from EAGER teardown: ``run()`` exits with one + credit still in flight (sent=1, returned=0), so an eager teardown in + ``run()``'s finally would fire while the return wait is still parked. + Teardown may only run after the returned event lands and the + background return-wait task's done-callback schedules it. + """ + r = make_runner(cfg(seamless=True), conv_src, pub, router, conc, cancel, cb) + strategy = TeardownStrategy() + with patch( + "aiperf.timing.phase.runner.plugins.get_class", + return_value=lambda **kw: strategy, + ): + # One credit in flight: run()'s sending-complete finally freezes + # final_requests_sent=1 with zero returned, so the background + # return wait parks on all_credits_returned_event. + r._progress._counter._requests_sent = 1 + r._progress.all_credits_sent_event.set() + await r.run(is_final_phase=False) + + assert r._return_wait_task is not None + assert not r._return_wait_task.done() + for _ in range(10): + await asyncio.sleep(0) + assert strategy.teardown_calls == 0, ( + "teardown fired while the phase's returns were still in " + "flight (eager-teardown regression)" + ) + + # The in-flight return lands -> the return-wait task finishes -> + # its done-callback schedules the deferred teardown. + r._progress.all_credits_returned_event.set() + for _ in range(10): + await asyncio.sleep(0) + assert strategy.teardown_calls == 1 + + +class _WarmupReportStrategy(MockStrategy): + """MockStrategy exposing the optional ``report_warmup_failures`` hook. + + ``raises`` toggles whether the hook aborts (agentx warmup-failure parity); + ``report_calls`` records whether the runner reached the hook at all. + """ + + report_calls: int = 0 + raises: bool = False + + def report_warmup_failures(self) -> None: + self.report_calls += 1 + if self.raises: + raise TrajectoryWarmupFailedError(["t-0#0.0[node_ordinal=0]: boom"]) + + +class TestWarmupFailureAbort: + """The runner aborts a non-cancelled phase whose strategy reports failures.""" + + async def test_run_strategy_propagates_warmup_failure( + self, + conv_src: MagicMock, + pub: MagicMock, + router: MagicMock, + conc: MagicMock, + cancel: MagicMock, + cb: MagicMock, + ) -> None: + """A raising ``report_warmup_failures`` propagates out of the non-seamless path.""" + r = make_runner( + cfg(phase=CreditPhase.WARMUP), conv_src, pub, router, conc, cancel, cb + ) + strategy = _WarmupReportStrategy() + strategy.raises = True + r._progress.all_credits_sent_event.set() + r._progress.all_credits_returned_event.set() + + with pytest.raises(TrajectoryWarmupFailedError): + await r._run_strategy(strategy, is_final_phase=True) + assert strategy.report_calls == 1 + + async def test_run_strategy_skips_warmup_report_on_cancel( + self, + conv_src: MagicMock, + pub: MagicMock, + router: MagicMock, + conc: MagicMock, + cancel: MagicMock, + cb: MagicMock, + ) -> None: + """A user-cancelled run returns early and never re-labels itself a warmup failure.""" + r = make_runner( + cfg(phase=CreditPhase.WARMUP), conv_src, pub, router, conc, cancel, cb + ) + strategy = _WarmupReportStrategy() + strategy.raises = True + r._was_cancelled = True + r._progress.all_credits_sent_event.set() + + result = await r._run_strategy(strategy, is_final_phase=True) + if r._progress_task is not None: + r._progress_task.cancel() + + assert strategy.report_calls == 0 + assert isinstance(result, CreditPhaseStats) + + +async def _idle_forever() -> None: + """Timer-free stand-in for the progress-report loop. + + The real loop sleeps on a repeating ``asyncio.sleep`` timer, and a + concurrent task holding an active timer defeats looptime's virtual-time + fast-forward of the returning-wait / drain ``wait_for``s. A pure event wait + (no timer) keeps the loop idle-detectable so looptime fast-forwards. + """ + await asyncio.Event().wait() + + +class TestWarmupDrainGrace: + """A finite warmup grace bounds the returning wait for a duration-less + WARMUP phase (agentx accelerated-warmup drain parity); infinite grace (the + boundary-priming default) keeps today's unbounded wait. + + These pin the BLOCKER fix: ``time_left_in_seconds(include_grace_period=True)`` + returns None when ``expected_duration_sec is None``, so a finite + ``grace_period_sec`` was DEAD config -- the returning wait was unbounded and + a lost pressure return hung the run. The runner now bounds a duration-less + WARMUP drain by the finite grace directly. + """ + + async def test_finite_warmup_grace_bounds_returning_wait( + self, + conv_src: MagicMock, + pub: MagicMock, + router: MagicMock, + conc: MagicMock, + cancel: MagicMock, + cb: MagicMock, + time_traveler: MagicMock, + ) -> None: + """RED before the fix: with ``expected_duration_sec=None`` the returning + wait never sees the finite grace, so a never-returning credit hangs + ``_run_strategy`` (the in-test ``wait_for`` bound then times out). The + fix bounds the WARMUP drain by the finite grace, so the credit is + cancelled via the timeout branch and the phase completes. + """ + r = make_runner( + cfg(phase=CreditPhase.WARMUP, dur=None, grace=0.2), + conv_src, + pub, + router, + conc, + cancel, + cb, + ) + r._progress_report_loop = _idle_forever + strategy = MockStrategy() + # One credit in flight that never returns: sending-complete freezes + # final_requests_sent=1 with zero returned, so the returning wait parks. + r._progress._counter._requests_sent = 1 + r._progress.all_credits_sent_event.set() + + result = await asyncio.wait_for( + r._run_strategy(strategy, is_final_phase=True), timeout=30.0 + ) + router.cancel_all_credits.assert_called_once() + assert isinstance(result, CreditPhaseStats) + + async def test_infinite_warmup_grace_leaves_returning_wait_unbounded( + self, + conv_src: MagicMock, + pub: MagicMock, + router: MagicMock, + conc: MagicMock, + cancel: MagicMock, + cb: MagicMock, + time_traveler: MagicMock, + ) -> None: + """The boundary-priming warmup (grace=inf) keeps the unbounded returning + wait: the finite-grace bound is WARMUP-only AND finite-only, so the + returning-wait timeout stays None and ``cancel_all_credits`` is never + reached. + """ + r = make_runner( + cfg(phase=CreditPhase.WARMUP, dur=None, grace=float("inf")), + conv_src, + pub, + router, + conc, + cancel, + cb, + ) + captured: dict[str, float | None] = {} + original = r._wait_for_event_with_timeout + + async def spy(*, name, event, timeout, **kwargs) -> bool: + if "returned" in name: + # Capture the timeout the runner computed for the returning wait, + # then short-circuit so the unbounded (None) wait does not hang. + captured["timeout"] = timeout + event.set() + return False + return await original(name=name, event=event, timeout=timeout, **kwargs) + + r._wait_for_event_with_timeout = spy + r._progress_report_loop = _idle_forever + strategy = MockStrategy() + r._progress._counter._requests_sent = 1 + r._progress.all_credits_sent_event.set() + + await asyncio.wait_for( + r._run_strategy(strategy, is_final_phase=True), timeout=30.0 + ) + assert captured["timeout"] is None + router.cancel_all_credits.assert_not_called() + + class TestComponentOwnership: def test_phase_property_returns_configured_phase(self, runner: PhaseRunner) -> None: assert runner.phase == CreditPhase.PROFILING diff --git a/tests/unit/timing/strategies/test_graph_observer_teardown.py b/tests/unit/timing/strategies/test_graph_observer_teardown.py new file mode 100644 index 0000000000..e087047a88 --- /dev/null +++ b/tests/unit/timing/strategies/test_graph_observer_teardown.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Deferred graph-phase teardown must not clear the NEXT phase's observers. + +One ``CreditCallbackHandler`` is shared by every ``PhaseRunner``. A seamless +non-final graph phase defers ``teardown_phase`` to its background return-wait +completion, which can fire AFTER the next phase's ``setup_phase`` installed ITS +observers on the same shared slots. Teardown must therefore compare-and-clear: +null a slot only when the observer installed there is the tearing-down +strategy's own. These tests drive the REAL ``CreditCallbackHandler`` (a +MagicMock handler hides exactly this bug) through the production +``PhaseRunner._build_graph_ir_strategy`` wiring. +""" + +from __future__ import annotations + +from pathlib import Path + +from aiperf.common.enums import CreditPhase +from aiperf.credit.callback_handler import CreditCallbackHandler +from aiperf.credit.messages import CreditReturn, FirstToken +from aiperf.credit.structs import Credit +from aiperf.dataset.graph.adapters.weka.trace import from_weka_trace +from aiperf.plugin.enums import TimingMode +from aiperf.timing.graph_channel import GraphPhaseChannel +from aiperf.timing.phase.runner import PhaseRunner +from aiperf.timing.strategies.graph_ir_replay import GraphIRReplayStrategy + +_WEKA_MIN = Path(__file__).parents[2] / "graph" / "fixtures" / "weka_min.json" + + +class _FakeConcurrency: + def release_session_slot(self, phase: CreditPhase) -> None: ... + + def release_prefill_slot(self, phase: CreditPhase) -> None: ... + + +class _FakeAdapter: + """Records the de-mux calls the strategy routes to the owning adapter.""" + + def __init__(self) -> None: + self.resolved: list[tuple[str | None, bool]] = [] + self.first_tokens: list[tuple[str | None, int | None]] = [] + + def resolve(self, credit: Credit, error: str | None, cancelled: bool) -> None: + self.resolved.append((error, cancelled)) + + def on_first_token( + self, x_correlation_id: str | None, turn_index: int | None + ) -> None: + self.first_tokens.append((x_correlation_id, turn_index)) + + +def _build_strategy(handler: CreditCallbackHandler) -> GraphIRReplayStrategy: + """Build through the production runner seam so the register/unregister + wiring under test is the real one (not a hand-rolled approximation).""" + + class _Config: + timing_mode = TimingMode.GRAPH_IR + phase = None + concurrency = None + expected_num_sessions = None + + channel = GraphPhaseChannel(parsed_graph=from_weka_trace(str(_WEKA_MIN))) + runner = PhaseRunner.__new__(PhaseRunner) + runner._config = _Config() + runner._conversation_source = None + runner._graph_channel = channel + runner._scheduler = None + runner._stop_checker = None + runner._credit_issuer = object() + runner._lifecycle = None + runner._callback_handler = handler + return runner._build_graph_ir_strategy(GraphIRReplayStrategy) + + +def _graph_credit(trace_id: str) -> Credit: + return Credit( + id=1, + phase=CreditPhase.PROFILING, + conversation_id="c", + x_correlation_id="x0", + turn_index=0, + num_turns=1, + issued_at_ns=0, + trace_id=trace_id, + node_ordinal=0, + phase_variant="profiling", + ) + + +def _first_token(trace_id: str) -> FirstToken: + return FirstToken( + credit_id=1, + phase=CreditPhase.PROFILING, + ttft_ns=5, + trace_id=trace_id, + x_correlation_id="x0", + turn_index=0, + ) + + +async def test_stale_deferred_teardown_preserves_next_phase_observers() -> None: + """Phase A's late teardown must not null phase B's live observers.""" + handler = CreditCallbackHandler(_FakeConcurrency()) + strategy_a = _build_strategy(handler) + strategy_b = _build_strategy(handler) + + await strategy_a.setup_phase() + await strategy_b.setup_phase() # the next phase takes over the shared slots + + # Phase A's deferred (return-wait done-callback) teardown fires late. + await strategy_a.teardown_phase() + + # B's return observer must still dispatch into B's adapter registry. + adapter = _FakeAdapter() + strategy_b._adapters["t-live"] = adapter + await handler.on_credit_return( + "w0", CreditReturn(credit=_graph_credit("t-live"), cancelled=False, error=None) + ) + assert adapter.resolved == [(None, False)], ( + "phase B's graph return was dropped: the stale phase A teardown " + "cleared the shared graph-return observer slot" + ) + + # B's first-token observer must survive too (post-TTFT anchoring). + await handler.on_first_token(_first_token("t-live")) + assert adapter.first_tokens == [("x0", 0)], ( + "phase B's first-token event was dropped: the stale phase A teardown " + "cleared the shared first-token observer slot" + ) + + +async def test_teardown_clears_own_still_installed_observers() -> None: + """A strategy whose observers ARE still installed must clear both slots.""" + handler = CreditCallbackHandler(_FakeConcurrency()) + strategy = _build_strategy(handler) + + await strategy.setup_phase() + assert handler._graph_return_observer is not None + assert handler._graph_first_token_observer is not None + + await strategy.teardown_phase() + assert handler._graph_return_observer is None + assert handler._graph_first_token_observer is None diff --git a/tests/unit/timing/strategies/test_user_centric_rate.py b/tests/unit/timing/strategies/test_user_centric_rate.py index 34d7677d51..06c2e5a1cd 100644 --- a/tests/unit/timing/strategies/test_user_centric_rate.py +++ b/tests/unit/timing/strategies/test_user_centric_rate.py @@ -3,8 +3,11 @@ from unittest.mock import MagicMock import pytest +from pytest import param from aiperf.common.enums import CreditPhase +from aiperf.common.models import TurnMetadata +from aiperf.credit.structs import Credit, TurnToSend from aiperf.plugin.enums import TimingMode from aiperf.timing.config import CreditPhaseConfig from aiperf.timing.strategies.user_centric_rate import User, UserCentricStrategy @@ -257,3 +260,71 @@ def test_dataclass_fields(self) -> None: assert u.next_send_time == 1000 assert u.max_turns == 3 assert u.order == 5 + + +class TestContinuationBranchPlumb: + """Regression: the continuation turn built in ``handle_credit_return`` must + carry the next turn's branch/fork facts. The strategy has to pass the next + turn's metadata into ``TurnToSend.from_previous_credit`` -- dropping it + silently zeroes ``has_branches`` (and ``has_forks``), so a branching turn + would be wrongly treated as a leaf by downstream finality stamping. + """ + + def _make_strategy( + self, next_meta: TurnMetadata + ) -> tuple[UserCentricStrategy, MagicMock]: + cfg = CreditPhaseConfig( + phase=CreditPhase.PROFILING, + timing_mode=TimingMode.USER_CENTRIC_RATE, + request_rate=10.0, + num_users=5, + total_expected_requests=10, + ) + conversation_source = MagicMock() + conversation_source.get_next_turn_metadata.return_value = next_meta + credit_issuer = MagicMock() + strategy = UserCentricStrategy( + config=cfg, + conversation_source=conversation_source, + scheduler=MagicMock(), + stop_checker=MagicMock(), + credit_issuer=credit_issuer, + lifecycle=MagicMock(), + ) + sampled = MagicMock() + sampled.x_correlation_id = "sess-1" + strategy._session_to_user["sess-1"] = User(user_id=1, sampled=sampled) + return strategy, credit_issuer + + @pytest.mark.parametrize( + "branch_ids, expected", + [ + param(["b1"], True, id="declares-branch"), + param([], False, id="no-branch"), + ], + ) # fmt: skip + async def test_continuation_turn_carries_has_branches( + self, branch_ids, expected + ) -> None: + next_meta = TurnMetadata(branch_ids=branch_ids) + strategy, credit_issuer = self._make_strategy(next_meta) + credit = Credit( + id=1, + phase=CreditPhase.PROFILING, + conversation_id="conv-1", + x_correlation_id="sess-1", + turn_index=0, + num_turns=3, + issued_at_ns=0, + ) + + await strategy.handle_credit_return(credit) + + strategy._conversation_source.get_next_turn_metadata.assert_called_once_with( + credit + ) + credit_issuer.issue_credit.assert_called_once() + turn = credit_issuer.issue_credit.call_args.args[0] + assert isinstance(turn, TurnToSend) + assert turn.turn_index == 1 + assert turn.has_branches is expected diff --git a/tests/unit/timing/test_graph_phase_channel.py b/tests/unit/timing/test_graph_phase_channel.py new file mode 100644 index 0000000000..5d83f6a692 --- /dev/null +++ b/tests/unit/timing/test_graph_phase_channel.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GraphPhaseChannel: first-class graph-phase state, no sampler for graph runs. + +Graph runs used to fabricate stub conversations so PhaseOrchestrator's +unconditional sampler build would not crash, and smuggled the ParsedGraph + +warmup handoff through attributes bolted onto the generic ConversationSource. +These tests pin the native shape: with a parsed graph the orchestrator builds +NO sampler and NO conversation source, and phase state travels on a typed +channel. +""" + +from unittest.mock import MagicMock + +from aiperf.common.models import DatasetMetadata +from aiperf.common.models.dataset_models import GraphDatasetMetadata +from aiperf.dataset.graph.models import GraphRecord, ParsedGraph, TraceRecord +from aiperf.plugin.enums import DatasetSamplingStrategy, TimingMode +from aiperf.timing.graph_channel import GraphPhaseChannel +from aiperf.timing.phase_orchestrator import PhaseOrchestrator +from tests.unit.timing.conftest import make_timing_config + + +def _parsed() -> ParsedGraph: + return ParsedGraph(graph=GraphRecord(), traces=[TraceRecord(id="t-1", tags=["x"])]) + + +def _make_router() -> MagicMock: + router = MagicMock() + router.set_return_callback = MagicMock() + router.set_first_token_callback = MagicMock() + return router + + +def test_channel_holds_parsed_graph_and_handoff_slot(): + channel = GraphPhaseChannel(parsed_graph=_parsed()) + assert channel.parsed_graph.traces[0].id == "t-1" + assert channel.warmup_handoff is None + + +def test_conversation_source_has_no_graph_attributes(): + """The R6 attribute-injection channel is gone.""" + from aiperf.timing.conversation_source import ConversationSource + + assert "parsed_graph" not in ConversationSource.__init__.__code__.co_names + src = ConversationSource.__new__(ConversationSource) + assert not hasattr(src, "parsed_graph") + assert not hasattr(src, "graph_warmup_handoff") + + +def test_graph_orchestrator_builds_no_sampler_and_owns_channel(): + """A graph run builds NO sampler and NO conversation source. + + ``DatasetMetadata.conversations`` is empty for graph runs, so the + conversation-shaped sampler would crash on empty ids; the orchestrator must + skip both and own a ``GraphPhaseChannel`` instead. + """ + orchestrator = PhaseOrchestrator( + config=make_timing_config(TimingMode.GRAPH_IR), + phase_publisher=MagicMock(), + credit_router=_make_router(), + dataset_metadata=DatasetMetadata( + conversations=[], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + graph=GraphDatasetMetadata(trace_ids=["t-1"]), + ), + parsed_graph=_parsed(), + ) + assert orchestrator._dataset_sampler is None + assert orchestrator._conversation_source is None + assert orchestrator._graph_channel is not None + assert orchestrator._graph_channel.parsed_graph.traces[0].id == "t-1" + + +def test_non_graph_orchestrator_builds_sampler_and_source(): + """A non-graph run is unchanged: sampler + ConversationSource built, no channel.""" + from aiperf.common.models import ConversationMetadata, TurnMetadata + + orchestrator = PhaseOrchestrator( + config=make_timing_config( + TimingMode.REQUEST_RATE, request_count=5, request_rate=10.0 + ), + phase_publisher=MagicMock(), + credit_router=_make_router(), + dataset_metadata=DatasetMetadata( + conversations=[ + ConversationMetadata(conversation_id="conv-0", turns=[TurnMetadata()]) + ], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ), + ) + assert orchestrator._dataset_sampler is not None + assert orchestrator._conversation_source is not None + assert orchestrator._graph_channel is None diff --git a/tests/unit/timing/test_session_tree_finality.py b/tests/unit/timing/test_session_tree_finality.py new file mode 100644 index 0000000000..b2f25844df --- /dev/null +++ b/tests/unit/timing/test_session_tree_finality.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Finality-query unit tests for the trimmed main ``SessionTreeRegistry``. + +The main registry is finality bookkeeping only: it holds no concurrency +manager and releases no slot (unlike the agentic-replay variant it was ported +from). These tests exercise the state map + the two finality queries directly. +""" + +from aiperf.common.enums import CreditPhase +from aiperf.timing.session_tree import SessionTreeRegistry + + +def _make_registry() -> SessionTreeRegistry: + return SessionTreeRegistry() + + +def _registry_with_tree( + root: str = "root-1", descendants: int = 0 +) -> SessionTreeRegistry: + registry = _make_registry() + registry.open_tree(root, phase=CreditPhase.PROFILING, root_pending=True) + if descendants: + registry.register_descendants(root, n=descendants) + return registry + + +def test_root_terminal_unknown_tree_is_none(): + assert _make_registry().root_terminal("nope") is None + + +def test_root_terminal_false_while_pending_true_after(): + # A live descendant keeps the tree open past on_root_terminal so the + # post-terminal state is observable; a drained tree would be retired + # (popped) and root_terminal would then read as unknown/None. + registry = _registry_with_tree(descendants=1) + assert registry.root_terminal("root-1") is False + registry.on_root_terminal("root-1") + assert registry.root_terminal("root-1") is True + + +def test_last_tree_request_root_with_no_descendants(): + registry = _registry_with_tree() + assert registry.is_last_tree_request( + "root-1", is_final_turn=True, is_root_credit=True, has_branches=False + ) + + +def test_not_last_when_descendants_outstanding_or_branches_pending(): + registry = _registry_with_tree(descendants=1) + assert not registry.is_last_tree_request( + "root-1", is_final_turn=True, is_root_credit=True, has_branches=False + ) + registry_no_desc = _registry_with_tree(root="root-2") + assert not registry_no_desc.is_last_tree_request( + "root-2", is_final_turn=True, is_root_credit=True, has_branches=True + ) + assert not registry_no_desc.is_last_tree_request( + "root-2", is_final_turn=False, is_root_credit=True, has_branches=False + ) + + +def test_last_tree_request_final_child_after_root_done(): + registry = _registry_with_tree(descendants=1) + registry.on_root_terminal("root-1") + assert registry.is_last_tree_request( + "root-1", is_final_turn=True, is_root_credit=False, has_branches=False + ) + + +def test_unknown_tree_is_conservative_false(): + assert not _make_registry().is_last_tree_request( + "nope", is_final_turn=True, is_root_credit=True, has_branches=False + ) + + +def test_descendant_done_retires_drained_tree(): + # Sole child completes after root terminal -> tree drains and is retired. + registry = _registry_with_tree(descendants=1) + registry.on_root_terminal("root-1") + assert registry.has_tree("root-1") + retired = registry.on_descendant_done("root-1") + assert retired is True + assert not registry.has_tree("root-1") + assert registry.root_terminal("root-1") is None + + +def test_release_all_retires_open_trees_without_slot_release(): + registry = _make_registry() + registry.open_tree("a", CreditPhase.PROFILING, root_pending=True) + registry.open_tree("b", CreditPhase.PROFILING, root_pending=True) + assert registry.open_count() == 2 + assert registry.release_all() == 2 + assert registry.open_count() == 0 + + +def test_final_turn_spawn_resurrects_retired_tree_for_grandchild_finality(): + """Regression: root done -> last-outstanding child C's final turn declares a + SPAWN grandchild. The callback order retires the tree on C's own + on_descendant_done (step 4b) BEFORE the return-intercept registers C's + grandchildren (step 5); register_descendants must RESURRECT the retired tree + so the grandchild's genuinely-last credit can still stamp is_tree_final=True. + """ + registry = _registry_with_tree(descendants=1) # root + one live child C + registry.on_root_terminal("root-1") # root's terminal turn returns; C still live + + # C is the last outstanding descendant; its final-turn return decrements it, + # draining and retiring the tree (step 4b, before C's spawn intercept). + assert registry.on_descendant_done("root-1") is True + assert not registry.has_tree("root-1") + + # Step 5: C's final-turn SPAWN registers the grandchild AFTER that retire. + # Old behavior buffered it into a retired root nothing drains; now it + # resurrects the tree root-terminal with the grandchild outstanding. + registry.register_descendants("root-1", n=1) + assert registry.has_tree("root-1") + assert registry.root_terminal("root-1") is True + + # The grandchild's genuinely-last (non-branching final) credit CAN now be + # stamped tree-final -- root terminal, sole outstanding descendant. + assert registry.is_last_tree_request( + "root-1", is_final_turn=True, is_root_credit=False, has_branches=False + ) + + # The grandchild finishing re-drains and re-retires the tree coherently. + assert registry.on_descendant_done("root-1") is True + assert not registry.has_tree("root-1") + assert registry.late_events == 0 + + +def test_release_all_drains_pending_descendants(): + """release_all must clear the pre-open descendant buffer, not just _trees.""" + registry = _make_registry() + # No open_tree first: register_descendants buffers into _pending_descendants + # (the defensive pre-open path). + registry.register_descendants("orphan-root", n=2) + assert registry._pending_descendants # buffered, not yet folded into a tree + + registry.release_all() + + assert registry._pending_descendants == {} + assert registry._retired_roots == {} + + +def test_finality_flows_credit_to_request_info(): + """Schema guard: the three lineage-finality fields exist on BOTH the Credit + struct and the RequestInfo model, so the worker has fields to copy between. + + This asserts field-NAME presence only -- it does NOT verify a value is + actually copied (deleting the plumb kwargs in ``worker._create_request_info`` + keeps this green). The value-level plumb guard is + ``test_worker.py::test_create_request_info_plumbs_finality_from_credit``, + which stamps a REAL Credit and asserts the values surface on the RequestInfo. + """ + from aiperf.common.models.record_models import RequestInfo + from aiperf.credit.structs import Credit + + credit_fields = {"is_parent_final", "is_tree_final", "root_correlation_id"} + assert credit_fields <= set(Credit.__struct_fields__) + assert credit_fields <= set(RequestInfo.model_fields) diff --git a/tests/unit/timing/test_session_tree_wiring.py b/tests/unit/timing/test_session_tree_wiring.py new file mode 100644 index 0000000000..9eb12cd990 --- /dev/null +++ b/tests/unit/timing/test_session_tree_wiring.py @@ -0,0 +1,487 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end wiring test at the DAG orchestrator seam. + +Drives the REAL CreditIssuer + REAL BranchOrchestrator + REAL +SessionTreeRegistry + REAL ConversationSource through the full tree lifecycle +(open -> register -> root terminal -> descendant done) and asserts the +transitions produce correct lineage finality on the ISSUED credits. No +MagicMock for the registry or any credit -- only the phase scalars +(progress/lifecycle/stop_checker) and the concurrency slots are faked. +""" + +import time +from unittest.mock import MagicMock + +import pytest + +from aiperf.common.enums import ConversationBranchMode, CreditPhase +from aiperf.common.models import ConversationMetadata, DatasetMetadata, TurnMetadata +from aiperf.common.models.branch import ConversationBranchInfo +from aiperf.credit.issuer import CreditIssuer +from aiperf.credit.structs import Credit, TurnToSend +from aiperf.plugin import plugins +from aiperf.plugin.enums import DatasetSamplingStrategy, PluginType +from aiperf.timing.branch_orchestrator import BranchOrchestrator +from aiperf.timing.conversation_source import ConversationSource +from aiperf.timing.session_tree import SessionTreeRegistry + + +class _FakeConcurrency: + """Slots always granted; releases are no-ops.""" + + async def acquire_session_slot(self, phase, can_proceed) -> bool: + return True + + async def acquire_prefill_slot(self, phase, can_proceed) -> bool: + return True + + def release_session_slot(self, phase) -> None: + pass + + +class _CapturingRouter: + def __init__(self) -> None: + self.sent: list[Credit] = [] + + async def send_credit(self, *, credit: Credit) -> None: + self.sent.append(credit) + + +def _mk_source() -> ConversationSource: + """Root (2 turns, SPAWN branch on turn 0) -> one child (2 turns).""" + ds = DatasetMetadata( + conversations=[ + ConversationMetadata( + conversation_id="root-conv", + turns=[ + TurnMetadata(timestamp_ms=0.0, branch_ids=["root-conv:0"]), + TurnMetadata(timestamp_ms=0.0), + ], + branches=[ + ConversationBranchInfo( + branch_id="root-conv:0", + child_conversation_ids=["child-conv"], + mode=ConversationBranchMode.SPAWN, + dispatch_timing="post", + ), + ], + agent_depth=0, + ), + ConversationMetadata( + conversation_id="child-conv", + turns=[TurnMetadata(timestamp_ms=0.0), TurnMetadata(timestamp_ms=0.0)], + agent_depth=1, + parent_conversation_id="root-conv", + is_root=False, + ), + ], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) + SamplerClass = plugins.get_class(PluginType.DATASET_SAMPLER, ds.sampling_strategy) + sampler = SamplerClass( + conversation_ids=[c.conversation_id for c in ds.conversations] + ) + return ConversationSource(ds, sampler) + + +def _make_issuer( + registry: SessionTreeRegistry, router: _CapturingRouter +) -> CreditIssuer: + progress = MagicMock() + progress.increment_sent = MagicMock(return_value=(1, False)) + progress.freeze_sent_counts = MagicMock() + progress.all_credits_sent_event = MagicMock() + + stop_checker = MagicMock() + stop_checker.can_send_any_turn = MagicMock(return_value=True) + stop_checker.can_start_new_session = MagicMock(return_value=True) + stop_checker.can_send_dag_child_turn = MagicMock(return_value=True) + + cancellation = MagicMock() + cancellation.next_cancellation_delay_ns = MagicMock(return_value=None) + + lifecycle = MagicMock() + lifecycle.started_at_ns = time.time_ns() + lifecycle.started_at_perf_ns = time.perf_counter_ns() + + return CreditIssuer( + phase=CreditPhase.PROFILING, + stop_checker=stop_checker, + progress=progress, + concurrency_manager=_FakeConcurrency(), + credit_router=router, + cancellation_policy=cancellation, + lifecycle=lifecycle, + session_tree_registry=registry, + ) + + +@pytest.mark.asyncio +async def test_orchestrator_seam_drives_tree_finality_transitions(): + registry = SessionTreeRegistry() + router = _CapturingRouter() + source = _mk_source() + issuer = _make_issuer(registry, router) + orch = BranchOrchestrator( + conversation_source=source, + credit_issuer=issuer, + session_tree_registry=registry, + ) + + # 1. OPEN: issuing the root's turn-0 credit opens the tree (issuer seam). + root_turn0 = TurnToSend( + conversation_id="root-conv", + x_correlation_id="root-x", + turn_index=0, + num_turns=2, + ) + await issuer.issue_credit(root_turn0) + assert registry.has_tree("root-x") + root_credit0 = router.sent[0] + assert root_credit0.agent_depth == 0 + # Turn 0 of 2 is not final -> conservative False. + assert root_credit0.is_tree_final is False + + # 2. REGISTER: intercepting the root turn-0 return spawns the child and + # registers it as a descendant (orchestrator seam). + await orch.intercept(root_credit0) + child_credits = [c for c in router.sent if c.agent_depth == 1] + assert len(child_credits) == 1 + child_credit0 = child_credits[0] + assert child_credit0.root_correlation_id == "root-x" + # Root still pending + descendant outstanding -> child not tree-final, and + # parent (== root) not yet final. + assert child_credit0.is_parent_final is False + assert child_credit0.is_tree_final is False + assert registry.open_count() == 1 + child_x = child_credit0.x_correlation_id + + # 3. ROOT TERMINAL: issue the root's final turn, then intercept it so the + # orchestrator clears root_pending (orchestrator seam via intercept). + root_turn1 = TurnToSend( + conversation_id="root-conv", + x_correlation_id="root-x", + turn_index=1, + num_turns=2, + ) + await issuer.issue_credit(root_turn1) + root_credit1 = router.sent[-1] + # Child still outstanding when the root's final turn is issued. + assert root_credit1.is_tree_final is False + await orch.intercept(root_credit1) + assert registry.root_terminal("root-x") is True + + # 4. FINALITY ON ISSUE: the child's final continuation, issued after the + # root terminal with the child as the sole outstanding descendant, is + # provably the tree's last request AND its parent (the root) is final. + child_turn1 = TurnToSend( + conversation_id="child-conv", + x_correlation_id=child_x, + turn_index=1, + num_turns=2, + agent_depth=1, + parent_correlation_id="root-x", + root_correlation_id="root-x", + ) + await issuer.dispatch_child_turn(child_turn1) + child_credit1 = router.sent[-1] + assert child_credit1.x_correlation_id == child_x + assert child_credit1.is_parent_final is True + assert child_credit1.is_tree_final is True + + # 5. DESCENDANT DONE: the child's terminal return drains + retires the tree + # (orchestrator seam). + await orch.on_child_leaf_reached(child_x) + assert not registry.has_tree("root-x") + assert registry.open_count() == 0 + + +# ============================================================================= +# Conservative-contract regressions: SPAWN branches must gate finality +# ============================================================================= + + +def _mk_single_turn_spawn_source() -> ConversationSource: + """Single-turn root whose ONLY turn declares a terminal SPAWN branch.""" + ds = DatasetMetadata( + conversations=[ + ConversationMetadata( + conversation_id="root-conv", + turns=[TurnMetadata(timestamp_ms=0.0, branch_ids=["root-conv:0"])], + branches=[ + ConversationBranchInfo( + branch_id="root-conv:0", + child_conversation_ids=["child-conv"], + mode=ConversationBranchMode.SPAWN, + dispatch_timing="post", + ), + ], + agent_depth=0, + ), + ConversationMetadata( + conversation_id="child-conv", + turns=[TurnMetadata(timestamp_ms=0.0)], + agent_depth=1, + parent_conversation_id="root-conv", + is_root=False, + ), + ], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) + SamplerClass = plugins.get_class(PluginType.DATASET_SAMPLER, ds.sampling_strategy) + sampler = SamplerClass( + conversation_ids=[c.conversation_id for c in ds.conversations] + ) + return ConversationSource(ds, sampler) + + +@pytest.mark.asyncio +async def test_single_turn_root_with_terminal_spawn_branch_not_tree_final(): + """(a) A single-turn root declaring a terminal SPAWN branch must stamp + is_tree_final=False on the root credit: its children spawn at + return-intercept, AFTER issue-time stamping, so the registry shows nothing + outstanding yet. Previously stamped a wrong True (has_forks is FORK-only).""" + from aiperf.timing.conversation_source import SampledSession + + registry = SessionTreeRegistry() + router = _CapturingRouter() + source = _mk_single_turn_spawn_source() + issuer = _make_issuer(registry, router) + orch = BranchOrchestrator( + conversation_source=source, + credit_issuer=issuer, + session_tree_registry=registry, + ) + + session = SampledSession( + conversation_id="root-conv", + metadata=source.get_metadata("root-conv"), + x_correlation_id="root-a", + ) + root_turn0 = session.build_first_turn() + # End-to-end stamp: SPAWN-only branch -> has_branches True, has_forks False. + assert root_turn0.has_branches is True + assert root_turn0.has_forks is False + + await issuer.issue_credit(root_turn0) + root_credit = router.sent[0] + assert root_credit.is_final_turn is True + assert root_credit.is_tree_final is False + + # The spawn then registers + the tree drains cleanly through the child. + await orch.intercept(root_credit) + child_credits = [c for c in router.sent if c.agent_depth == 1] + assert len(child_credits) == 1 + await orch.on_child_leaf_reached(child_credits[0].x_correlation_id) + assert not registry.has_tree("root-a") + assert registry.late_events == 0 + + +def _mk_grandchild_spawn_source() -> ConversationSource: + """Single-turn root -> child (2 turns; final turn declares a SPAWN + branch to a grandchild).""" + ds = DatasetMetadata( + conversations=[ + ConversationMetadata( + conversation_id="root-conv", + turns=[TurnMetadata(timestamp_ms=0.0, branch_ids=["root-conv:0"])], + branches=[ + ConversationBranchInfo( + branch_id="root-conv:0", + child_conversation_ids=["child-conv"], + mode=ConversationBranchMode.SPAWN, + dispatch_timing="post", + ), + ], + agent_depth=0, + ), + ConversationMetadata( + conversation_id="child-conv", + turns=[ + TurnMetadata(timestamp_ms=0.0), + TurnMetadata(timestamp_ms=0.0, branch_ids=["child-conv:1"]), + ], + branches=[ + ConversationBranchInfo( + branch_id="child-conv:1", + child_conversation_ids=["grandchild-conv"], + mode=ConversationBranchMode.SPAWN, + dispatch_timing="post", + ), + ], + agent_depth=1, + parent_conversation_id="root-conv", + is_root=False, + ), + ConversationMetadata( + conversation_id="grandchild-conv", + turns=[TurnMetadata(timestamp_ms=0.0)], + agent_depth=2, + parent_conversation_id="child-conv", + is_root=False, + ), + ], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) + SamplerClass = plugins.get_class(PluginType.DATASET_SAMPLER, ds.sampling_strategy) + sampler = SamplerClass( + conversation_ids=[c.conversation_id for c in ds.conversations] + ) + return ConversationSource(ds, sampler) + + +@pytest.mark.asyncio +async def test_child_final_turn_spawning_grandchild_not_tree_final(): + """(b) A child's final turn that declares a SPAWN branch (grandchild + pending) must stamp is_tree_final=False even when the child is the sole + outstanding descendant and the root is already terminal -- previously the + exact wrong-True scenario.""" + registry = SessionTreeRegistry() + router = _CapturingRouter() + source = _mk_grandchild_spawn_source() + issuer = _make_issuer(registry, router) + orch = BranchOrchestrator( + conversation_source=source, + credit_issuer=issuer, + session_tree_registry=registry, + ) + + # Root turn 0 (single turn, spawning): opens tree; intercept spawns the + # child, registers it, then marks the root terminal. + root_turn0 = TurnToSend( + conversation_id="root-conv", + x_correlation_id="root-b", + turn_index=0, + num_turns=1, + has_branches=True, + ) + await issuer.issue_credit(root_turn0) + root_credit = router.sent[0] + assert root_credit.is_tree_final is False + await orch.intercept(root_credit) + assert registry.root_terminal("root-b") is True + child_credit0 = next(c for c in router.sent if c.agent_depth == 1) + child_x = child_credit0.x_correlation_id + + # Child's final turn, built from real metadata (branch_ids on turn 1), + # issued while the child is the sole outstanding descendant. + next_meta = source.get_next_turn_metadata(child_credit0) + child_turn1 = TurnToSend.from_previous_credit(child_credit0, next_meta) + assert child_turn1.has_branches is True + assert child_turn1.has_forks is False + await issuer.dispatch_child_turn(child_turn1) + child_credit1 = router.sent[-1] + assert child_credit1.is_final_turn is True + assert child_credit1.is_parent_final is True + assert child_credit1.is_tree_final is False # grandchild pending + + # Production return order: leaf-reached fires BEFORE intercept spawns the + # grandchild (callback handler step 4b before step 5). The child is the last + # outstanding descendant, so its leaf-reached decrement drains and retires + # the tree; intercept's register_descendants then RESURRECTS it (root already + # terminal). The grandchild -- a single-turn, sole-remaining, genuinely-last + # request -- is therefore correctly stamped tree-final (previously the + # resurrect was missing and it was wrongly under-fired to False). + await orch.on_child_leaf_reached(child_x) + await orch.intercept(child_credit1) + grandchild_credits = [c for c in router.sent if c.agent_depth == 2] + assert len(grandchild_credits) == 1 + assert grandchild_credits[0].root_correlation_id == "root-b" + assert grandchild_credits[0].is_tree_final is True + await orch.on_child_leaf_reached(grandchild_credits[0].x_correlation_id) + assert not registry.has_tree("root-b") + assert registry.late_events == 0 + + +def _mk_pre_session_source() -> ConversationSource: + """Two-turn root with a pre-session SPAWN branch attached to turn 0.""" + ds = DatasetMetadata( + conversations=[ + ConversationMetadata( + conversation_id="root-conv", + turns=[ + TurnMetadata(timestamp_ms=0.0, branch_ids=["root-conv:pre"]), + TurnMetadata(timestamp_ms=0.0), + ], + branches=[ + ConversationBranchInfo( + branch_id="root-conv:pre", + child_conversation_ids=["pre-child-conv"], + mode=ConversationBranchMode.SPAWN, + dispatch_timing="pre", + ), + ], + agent_depth=0, + ), + ConversationMetadata( + conversation_id="pre-child-conv", + turns=[TurnMetadata(timestamp_ms=0.0)], + agent_depth=1, + parent_conversation_id="root-conv", + is_root=False, + ), + ], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) + SamplerClass = plugins.get_class(PluginType.DATASET_SAMPLER, ds.sampling_strategy) + sampler = SamplerClass( + conversation_ids=[c.conversation_id for c in ds.conversations] + ) + return ConversationSource(ds, sampler) + + +@pytest.mark.asyncio +async def test_pre_session_child_live_blocks_root_final_turn_tree_final(): + """(c) A live pre-session SPAWN child must keep the root's final turn from + stamping is_tree_final=True. Pre-dispatch runs before sampling (no root id + exists yet), so the orchestrator folds the live pre children into the root + instance's tree at its turn-0 return -- before any final-turn issue.""" + registry = SessionTreeRegistry() + router = _CapturingRouter() + source = _mk_pre_session_source() + issuer = _make_issuer(registry, router) + orch = BranchOrchestrator( + conversation_source=source, + credit_issuer=issuer, + session_tree_registry=registry, + ) + + # Pre-dispatch fires the background child before any root credit exists. + await orch.dispatch_pre_session_branches() + pre_credits = [c for c in router.sent if c.agent_depth == 1] + assert len(pre_credits) == 1 + pre_child_x = pre_credits[0].x_correlation_id + assert pre_credits[0].is_tree_final is False + + # Root turn 0: opens the tree; its return-intercept folds the live + # pre-session child into this root's tree. + root_turn0 = TurnToSend( + conversation_id="root-conv", + x_correlation_id="root-c", + turn_index=0, + num_turns=2, + has_branches=True, # turn 0 declares the pre branch + ) + await issuer.issue_credit(root_turn0) + root_credit0 = [c for c in router.sent if c.agent_depth == 0][0] + await orch.intercept(root_credit0) + + # Root's FINAL turn (declares no branches) with the pre child still live: + # must stamp False. Previously wrong-True (pre children were never + # registered anywhere). + next_meta = source.get_next_turn_metadata(root_credit0) + root_turn1 = TurnToSend.from_previous_credit(root_credit0, next_meta) + assert root_turn1.has_branches is False + await issuer.issue_credit(root_turn1) + root_credit1 = router.sent[-1] + assert root_credit1.is_final_turn is True + assert root_credit1.is_tree_final is False # pre-session child still live + + # Root terminal, then the pre child's terminal drains + retires the tree. + await orch.intercept(root_credit1) + assert registry.root_terminal("root-c") is True + assert registry.has_tree("root-c") + await orch.on_child_leaf_reached(pre_child_x) + assert not registry.has_tree("root-c") + assert registry.late_events == 0 diff --git a/tests/unit/timing/test_timing_manager_first_token_advisory.py b/tests/unit/timing/test_timing_manager_first_token_advisory.py new file mode 100644 index 0000000000..284f143c6c --- /dev/null +++ b/tests/unit/timing/test_timing_manager_first_token_advisory.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TimingManager first-token-anchoring source-streaming advisory. + +A post-TTFT-anchored ``StaticEdge`` releases its successor at the SOURCE node's +observed first token. That observation only exists when the source node itself +streams. Each graph node streams per its own recorded ``streaming`` mode +(per-request override), so the global ``--streaming`` flag does not govern +whether a first-token event is emitted. The advisory therefore warns iff a +first-token edge's SOURCE ``LlmNode`` carries ``streaming=False`` -- possible +only in hand-authored/degenerate graphs, since recorded corpora are consistent +by construction (the same recorded ttft drives both the edge and the node mode). +These tests exercise that source-node matrix; the advisory logs once via the +module logger (not ``self.warning``), so no lifecycle/logger init is needed. +""" + +import logging + +from aiperf.dataset.graph.models import ( + GraphRecord, + LlmNode, + ParsedGraph, + StaticEdge, + TraceRecord, +) +from aiperf.timing.manager import TimingManager + +_ADVISORY_LOGGER = "aiperf.timing.manager" +_ADVISORY_NEEDLE = "first-token-anchored" + + +def _make_tm() -> TimingManager: + """Bare TimingManager (bypassing service init). + + The advisory reads nothing off ``self`` -- it scans the passed graph and + logs via the module logger -- so a raw ``__new__`` instance suffices. + """ + return TimingManager.__new__(TimingManager) + + +def _graph_with_nodes(*, source_streaming: bool, first_token: bool) -> GraphRecord: + """One edge ``a -> b`` with both ends materialized as ``LlmNode``s. + + ``first_token`` toggles post-TTFT anchoring on the edge; ``source_streaming`` + sets the source node ``a``'s recorded streaming mode. + """ + edge = StaticEdge( + source="a", + target="b", + delay_after_predecessor_start_us=1000.0, + delay_after_predecessor_first_token_us=500.0 if first_token else None, + ) + return GraphRecord( + nodes={ + "a": LlmNode(prompt=[], output="a_out", streaming=source_streaming), + "b": LlmNode(prompt=[], output="b_out"), + }, + edges=[edge], + ) + + +def _advisory_records(caplog) -> list[logging.LogRecord]: + return [r for r in caplog.records if _ADVISORY_NEEDLE in r.getMessage()] + + +def test_recorded_streaming_sources_are_silent(caplog): + """Recorded corpus: first-token source streams -> no warning (global OFF).""" + parsed = ParsedGraph( + graph=_graph_with_nodes(source_streaming=True, first_token=True), + traces=[TraceRecord(id="t-1", tags=["x"])], + ) + tm = _make_tm() + with caplog.at_level(logging.WARNING, logger=_ADVISORY_LOGGER): + tm._advise_non_streaming_first_token_sources(parsed) + assert _advisory_records(caplog) == [] + + +def test_non_streaming_first_token_source_warns_once(caplog): + """Degenerate corpus: first-token source has streaming=False -> one warning.""" + parsed = ParsedGraph( + graph=_graph_with_nodes(source_streaming=False, first_token=True), + traces=[TraceRecord(id="t-1", tags=["x"])], + ) + tm = _make_tm() + with caplog.at_level(logging.WARNING, logger=_ADVISORY_LOGGER): + tm._advise_non_streaming_first_token_sources(parsed) + records = _advisory_records(caplog) + assert len(records) == 1, "expected exactly one source-streaming advisory warning" + message = records[0].getMessage() + assert "'a'" in message, "warning must name the offending source node id" + assert "COMPLETION latch" in message, ( + "warning must describe the completion-latch fallback" + ) + + +def test_no_first_token_edges_is_silent(caplog): + """No post-TTFT edge -> no first-token source -> silent even for streaming=False.""" + parsed = ParsedGraph( + graph=_graph_with_nodes(source_streaming=False, first_token=False), + traces=[TraceRecord(id="t-1", tags=["x"])], + ) + tm = _make_tm() + with caplog.at_level(logging.WARNING, logger=_ADVISORY_LOGGER): + tm._advise_non_streaming_first_token_sources(parsed) + assert _advisory_records(caplog) == [] + + +def test_non_streaming_source_in_subgraph_warns(caplog): + """Detection scans subgraph bodies too, not just the top-level graph.""" + parsed = ParsedGraph( + graph=GraphRecord(), + graphs={"body": _graph_with_nodes(source_streaming=False, first_token=True)}, + traces=[TraceRecord(id="t-1", tags=["x"])], + ) + tm = _make_tm() + with caplog.at_level(logging.WARNING, logger=_ADVISORY_LOGGER): + tm._advise_non_streaming_first_token_sources(parsed) + assert len(_advisory_records(caplog)) == 1 + + +def test_multiple_non_streaming_sources_warn_once(caplog): + """Once-per-run: multiple non-streaming first-token sources -> one warning.""" + graph = GraphRecord( + nodes={ + "a": LlmNode(prompt=[], output="a_out", streaming=False), + "b": LlmNode(prompt=[], output="b_out"), + "c": LlmNode(prompt=[], output="c_out", streaming=False), + "d": LlmNode(prompt=[], output="d_out"), + }, + edges=[ + StaticEdge( + source="a", + target="b", + delay_after_predecessor_start_us=1000.0, + delay_after_predecessor_first_token_us=500.0, + ), + StaticEdge( + source="c", + target="d", + delay_after_predecessor_start_us=1000.0, + delay_after_predecessor_first_token_us=600.0, + ), + ], + ) + parsed = ParsedGraph(graph=graph, traces=[TraceRecord(id="t-1", tags=["x"])]) + tm = _make_tm() + with caplog.at_level(logging.WARNING, logger=_ADVISORY_LOGGER): + tm._advise_non_streaming_first_token_sources(parsed) + assert len(_advisory_records(caplog)) == 1 diff --git a/tests/unit/timing/test_timing_manager_sidecar_load.py b/tests/unit/timing/test_timing_manager_sidecar_load.py new file mode 100644 index 0000000000..facae409e0 --- /dev/null +++ b/tests/unit/timing/test_timing_manager_sidecar_load.py @@ -0,0 +1,141 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``TimingManager._load_graph_sidecar`` mandatory broadcast-ingest contract. + +The DatasetManager writes the graph_meta sidecar on EVERY graph build route +and advertises its exact path on the graph-typed +``DatasetConfiguredNotification.client_metadata``; the schedule plane ingests +the sidecar from that broadcast path only. A graph run whose broadcast is not +graph-typed, or whose advertised file is missing, undecodable, or +store-divergent, is a hard configure-time failure. There is NO re-parse +fallback and NO env-convention path re-derivation. +""" + +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +from aiperf.common.exceptions import InvalidStateError +from aiperf.common.models.dataset_models import GraphSegmentClientMetadata +from aiperf.config.resolution.plan import GraphWorkloadRef +from aiperf.dataset.graph.graph_meta_sidecar import write_graph_meta_sidecar +from aiperf.dataset.graph.models import GraphRecord, ParsedGraph, TraceRecord +from aiperf.timing.manager import TimingManager + + +def _graph_meta(tmp_path: Path, sidecar: Path) -> GraphSegmentClientMetadata: + return GraphSegmentClientMetadata( + store_base_path=tmp_path, benchmark_id="b1", sidecar_path=sidecar + ) + + +def _make_tm(client_metadata: GraphSegmentClientMetadata | None) -> TimingManager: + tm = TimingManager.__new__(TimingManager) # bypass full service init + tm.run = MagicMock() + tm.run.benchmark_id = "b1" + tm._graph_client_metadata = client_metadata + tm.info = lambda *a, **k: None + tm.debug = lambda *a, **k: None + return tm + + +def _force_graph_run(monkeypatch) -> None: + # The single detection seam every consumer reads. Patched (not derived): + # this MagicMock run never carries a real memo, and the accessor's strict + # ``graph_workload_resolved is True`` check refuses the mock's + # auto-truthified attributes. + monkeypatch.setattr( + "aiperf.dataset.graph.workload_detect.resolve_graph_workload", + lambda run: GraphWorkloadRef( + path=Path("/does/not/matter.json"), format="weka_trace" + ), + ) + + +def _write_sidecar(tmp_path: Path) -> Path: + pg = ParsedGraph(graph=GraphRecord(), traces=[TraceRecord(id="t-7", tags=["x"])]) + return write_graph_meta_sidecar( + pg, + base_path=tmp_path, + benchmark_id="b1", + source_fingerprint={}, + schema_version=1, + ) + + +def test_loads_sidecar_from_broadcast_path(tmp_path, monkeypatch): + sidecar = _write_sidecar(tmp_path) + tm = _make_tm(_graph_meta(tmp_path, sidecar)) + _force_graph_run(monkeypatch) + result = tm._load_graph_sidecar() + assert result is not None + assert [t.id for t in result.traces] == ["t-7"] + + +def test_non_graph_run_returns_none(monkeypatch): + tm = _make_tm(None) + monkeypatch.setattr( + "aiperf.dataset.graph.workload_detect.resolve_graph_workload", + lambda run: None, + ) + assert tm._load_graph_sidecar() is None + + +def test_graph_run_without_graph_broadcast_raises(monkeypatch): + tm = _make_tm(None) + _force_graph_run(monkeypatch) + with pytest.raises(InvalidStateError, match="GraphSegmentClientMetadata"): + tm._load_graph_sidecar() + + +def test_advertised_but_missing_file_raises(tmp_path, monkeypatch): + missing = tmp_path / "aiperf_graph_meta_b1" / "graph_meta.msgpack" + tm = _make_tm(_graph_meta(tmp_path, missing)) + _force_graph_run(monkeypatch) + with pytest.raises(InvalidStateError, match="missing"): + tm._load_graph_sidecar() + + +def test_corrupt_sidecar_raises(tmp_path, monkeypatch): + sidecar = tmp_path / "aiperf_graph_meta_b1" / "graph_meta.msgpack" + sidecar.parent.mkdir(parents=True) + sidecar.write_bytes(b"not-a-msgpack-frame") + tm = _make_tm(_graph_meta(tmp_path, sidecar)) + _force_graph_run(monkeypatch) + with pytest.raises(InvalidStateError, match="unreadable"): + tm._load_graph_sidecar() + + +def test_index_mismatch_raises(tmp_path, monkeypatch): + sidecar = _write_sidecar(tmp_path) + tm = _make_tm(_graph_meta(tmp_path, sidecar)) + _force_graph_run(monkeypatch) + tm._sidecar_passes_index_check = lambda graph, sidecar: False + with pytest.raises(InvalidStateError, match="cross-check"): + tm._load_graph_sidecar() + + +@pytest.mark.asyncio +async def test_notification_handler_stores_graph_client_metadata(tmp_path): + import asyncio + + from aiperf.common.messages import DatasetConfiguredNotification + from aiperf.common.models.dataset_models import DatasetMetadata + from aiperf.plugin.enums import DatasetSamplingStrategy + + tm = _make_tm(None) + tm._dataset_metadata = None + tm._dataset_configured_event = asyncio.Event() + sidecar = tmp_path / "aiperf_graph_meta_b1" / "graph_meta.msgpack" + message = DatasetConfiguredNotification( + service_id="dataset_manager", + metadata=DatasetMetadata( + conversations=[], sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL + ), + client_metadata=_graph_meta(tmp_path, sidecar), + ) + await tm._on_dataset_configured_notification(message) + assert tm._graph_client_metadata is not None + assert tm._graph_client_metadata.sidecar_path == sidecar + assert tm._dataset_configured_event.is_set() diff --git a/tests/unit/transports/test_aiohttp_transport.py b/tests/unit/transports/test_aiohttp_transport.py index 8102edc917..0d74472291 100644 --- a/tests/unit/transports/test_aiohttp_transport.py +++ b/tests/unit/transports/test_aiohttp_transport.py @@ -6,11 +6,16 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from pytest import param -from aiperf.common.enums import ConnectionReuseStrategy, CreditPhase +from aiperf.common.enums import ( + ConnectionReuseStrategy, + CreditPhase, + ModelSelectionStrategy, +) from aiperf.common.models.record_models import RequestInfo, RequestRecord from aiperf.plugin import plugins -from aiperf.plugin.enums import TransportType +from aiperf.plugin.enums import EndpointType, TransportType from aiperf.transports.aiohttp_transport import ( AioHttpTransport, ConnectionLeaseManager, @@ -31,6 +36,7 @@ def create_request_info( is_final_turn: bool = True, turn_index: int = 0, credit_num: int = 1, + stream_override: bool | None = None, ) -> RequestInfo: """Create RequestInfo with sensible defaults for transport tests.""" return RequestInfo( @@ -46,6 +52,7 @@ def create_request_info( conversation_id=conversation_id, cancel_after_ns=cancel_after_ns, is_final_turn=is_final_turn, + stream_override=stream_override, ) @@ -142,6 +149,74 @@ def test_get_transport_headers(self, transport, streaming, expected_accept): assert headers["Content-Type"] == "application/json" assert headers["Accept"] == expected_accept + @pytest.mark.parametrize( + "endpoint_streaming,stream_override,expected_accept", + [ + param(False, True, "text/event-stream", id="override-on-beats-global-off"), + param(True, False, "application/json", id="override-off-beats-global-on"), + param(True, None, "text/event-stream", id="no-override-follows-global-on"), + param(False, None, "application/json", id="no-override-follows-global-off"), + ], + ) # fmt: skip + def test_transport_accept_header_honors_override( + self, transport, endpoint_streaming, stream_override, expected_accept + ): + """The Accept header follows the per-request stream override for graph credits. + + A recorded per-node override wins over the global ``endpoint.streaming``; + ``None`` follows the global (the non-graph default). + """ + model_endpoint = create_model_endpoint_info(streaming=endpoint_streaming) + request_info = create_request_info( + model_endpoint, stream_override=stream_override + ) + headers = transport.get_transport_headers(request_info) + assert headers["Accept"] == expected_accept + + @pytest.mark.parametrize( + "endpoint_streaming,stream_override,expected_path", + [ + param(False, True, "/generate_stream", id="override-on-selects-stream-path"), + param(True, False, "/generate", id="override-off-selects-base-path"), + param(True, None, "/generate_stream", id="no-override-follows-global-on"), + param(False, None, "/generate", id="no-override-follows-global-off"), + ], + ) # fmt: skip + def test_streaming_path_selection_honors_override( + self, endpoint_streaming, stream_override, expected_path + ): + """The streaming URL path follows the per-request override for graph credits. + + The huggingface_generate endpoint exposes a distinct ``streaming_path`` + (``/generate_stream``) vs its base ``endpoint_path`` (``/generate``); the + per-node override selects it independent of the global flag. + """ + from aiperf.common.models.model_endpoint_info import ( + EndpointInfo, + ModelEndpointInfo, + ModelInfo, + ModelListInfo, + ) + + model_endpoint = ModelEndpointInfo( + models=ModelListInfo( + models=[ModelInfo(name="test-model")], + model_selection_strategy=ModelSelectionStrategy.ROUND_ROBIN, + ), + endpoint=EndpointInfo( + type=EndpointType.HUGGINGFACE_GENERATE, + base_urls=["http://localhost:8000"], + custom_endpoint=None, + streaming=endpoint_streaming, + ), + ) + transport = AioHttpTransport(model_endpoint=model_endpoint) + request_info = create_request_info( + model_endpoint, stream_override=stream_override + ) + url = transport.get_url(request_info) + assert url == f"http://localhost:8000{expected_path}" + @pytest.mark.parametrize( "base_url,custom_endpoint,expected_url", [ diff --git a/tests/unit/transports/test_aiohttp_transport_bytes_passthrough.py b/tests/unit/transports/test_aiohttp_transport_bytes_passthrough.py new file mode 100644 index 0000000000..f530969a7b --- /dev/null +++ b/tests/unit/transports/test_aiohttp_transport_bytes_passthrough.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from unittest.mock import AsyncMock + +import orjson +import pytest + +from aiperf.common.models.record_models import RequestRecord +from aiperf.transports.aiohttp_transport import AioHttpTransport +from tests.unit.transports.test_aiohttp_transport import create_request_info + + +class TestBytesPayloadPassthrough: + """A pre-serialized bytes payload must reach aiohttp ``data=`` verbatim. + + The transport double-encodes a bytes payload into a JSON string (``"..."``) + if it routes through ``orjson.dumps``; the bytes path must short-circuit. + """ + + @pytest.fixture + def transport(self, model_endpoint_non_streaming): + return AioHttpTransport(model_endpoint=model_endpoint_non_streaming) + + async def _setup(self, transport): + await transport.initialize() + transport.aiohttp_client.post_request = AsyncMock(return_value=RequestRecord()) + + @pytest.mark.asyncio + async def test_bytes_payload_sent_verbatim( + self, transport, model_endpoint_non_streaming + ): + """A bytes payload is forwarded byte-identical, not orjson.dumps'd.""" + await self._setup(transport) + + request_info = create_request_info(model_endpoint_non_streaming) + payload = b'{"model":"m","messages":[{"role":"user","content":"hi"}]}' + + await transport.send_request(request_info, payload) + + body = transport.aiohttp_client.post_request.call_args[0][1] + # Identity: the exact bytes object is forwarded, never re-encoded. + assert body is payload + + @pytest.mark.asyncio + async def test_bytearray_payload_sent_verbatim( + self, transport, model_endpoint_non_streaming + ): + """A bytearray payload is also short-circuited (no orjson.dumps).""" + await self._setup(transport) + + request_info = create_request_info(model_endpoint_non_streaming) + payload = bytearray(b'{"k":"v"}') + + await transport.send_request(request_info, payload) + + body = transport.aiohttp_client.post_request.call_args[0][1] + # Identity: the exact bytearray object is forwarded, never re-encoded. + assert body is payload + + @pytest.mark.asyncio + async def test_dict_payload_still_orjson_dumped( + self, transport, model_endpoint_non_streaming + ): + """A dict payload keeps its existing orjson.dumps behavior (unchanged).""" + await self._setup(transport) + + request_info = create_request_info(model_endpoint_non_streaming) + payload = {"model": "m", "temperature": 0.7} + + await transport.send_request(request_info, payload) + + body = transport.aiohttp_client.post_request.call_args[0][1] + assert body == orjson.dumps(payload) diff --git a/tests/unit/ui/test_realtime_metrics_dashboard.py b/tests/unit/ui/test_realtime_metrics_dashboard.py index 0f29470483..b2b5ff70dc 100644 --- a/tests/unit/ui/test_realtime_metrics_dashboard.py +++ b/tests/unit/ui/test_realtime_metrics_dashboard.py @@ -14,6 +14,9 @@ OutputTokenCountMetric, ) from aiperf.metrics.types.request_latency_metric import RequestLatencyMetric +from aiperf.metrics.types.theoretical_prefix_cache_metric import ( + TheoreticalPrefixCacheHitMetric, +) from aiperf.metrics.types.ttft_metric import TTFTMetric from aiperf.ui.dashboard.realtime_metrics_dashboard import RealtimeMetricsTable @@ -37,6 +40,9 @@ class TestRealtimeMetricsTable: (TTFTMetric.tag, True, False), (InterTokenLatencyMetric.tag, False, False), (InterTokenLatencyMetric.tag, True, False), + # Externally-injected accumulator metric (CACHE group) - always shown + (TheoreticalPrefixCacheHitMetric.tag, False, False), + (TheoreticalPrefixCacheHitMetric.tag, True, False), ], ) # fmt: skip def test_should_skip_logic_with_real_metrics( @@ -55,3 +61,67 @@ def test_should_skip_logic_with_real_metrics( ) assert table._should_skip(metric_result) is should_skip + + +class TestRealtimeMetricsTableExternallyInjectedTags: + """Registered accumulator tags render; unregistered tags cannot kill the table. + + The theoretical_prefix_cache_hit MetricResult is minted by a standalone + accumulator (not the record-metric pipeline), so it reaches the dashboard + on every graph-IR run. Before its display class was registered, the strict + MetricRegistry.get_class calls in the update path raised MetricTypeError on + every tick and no rows rendered at all. + """ + + def _mounted_table(self) -> RealtimeMetricsTable: + table = RealtimeMetricsTable(Mock()) + table.data_table = Mock() + table.data_table.is_mounted = True + return table + + def _rendered_headers(self, table: RealtimeMetricsTable) -> list[str]: + return [call.args[0].plain for call in table.data_table.add_row.call_args_list] + + def test_update_renders_theoretical_prefix_cache_hit_row(self): + table = self._mounted_table() + table.update( + [ + MetricResult( + tag=TheoreticalPrefixCacheHitMetric.tag, + header="ignored - display metadata comes from the registry", + unit="%", + avg=42.5, + ), + ] + ) + assert self._rendered_headers(table) == ["Theoretical Prefix Cache Hit (%)"] + + def test_update_with_unregistered_tag_renders_fallback_and_sorts_last(self): + table = self._mounted_table() + table.update( + [ + MetricResult( + tag="not_a_registered_metric", + header="Mystery", + unit="widgets", + avg=1.0, + ), + MetricResult( + tag=TheoreticalPrefixCacheHitMetric.tag, + header="ignored", + unit="%", + avg=42.5, + ), + ] + ) + assert self._rendered_headers(table) == [ + "Theoretical Prefix Cache Hit (%)", + "Mystery (widgets)", + ] + + def test_should_skip_returns_false_for_unregistered_tag(self): + table = RealtimeMetricsTable(Mock()) + metric_result = MetricResult( + tag="not_a_registered_metric", header="Mystery", unit="widgets", avg=1.0 + ) + assert table._should_skip(metric_result) is False diff --git a/tests/unit/workers/test_inference_client.py b/tests/unit/workers/test_inference_client.py index 4d23f72515..06b10c50ac 100644 --- a/tests/unit/workers/test_inference_client.py +++ b/tests/unit/workers/test_inference_client.py @@ -5,6 +5,7 @@ import warnings from unittest.mock import AsyncMock, MagicMock, patch +import orjson import pytest from pytest import param @@ -20,6 +21,19 @@ from aiperf.common.redact import REDACTED_VALUE from aiperf.plugin.enums import EndpointType, TransportType from aiperf.workers.inference_client import InferenceClient, detect_transport_from_url +from aiperf.workers.session_routing import ( + DynamoHeadersRouting, + DynamoNvextRouting, + SessionIdHeaderRouting, + SmgRoutingKeyRouting, +) + +_ROUTING_CLASSES = { + "dynamo_headers": DynamoHeadersRouting, + "dynamo_nvext": DynamoNvextRouting, + "smg_routing_key": SmgRoutingKeyRouting, + "session_id_header": SessionIdHeaderRouting, +} @pytest.fixture @@ -511,3 +525,339 @@ def mock_get_class(protocol, name): assert not pydantic_warnings, ( f"Unexpected Pydantic serialization warnings for {base_url!r}: {pydantic_warnings}" ) + + +class TestSessionRouting: + """Session-routing plugins wired through the InferenceClient chokepoint. + + The endpoint/transport plugins are mocked as before; the session_routing + protocol resolves to the real routing classes so the chokepoint exercises + genuine header/body transforms and the notify_session_end pass-through. + """ + + def _build_client( + self, + mock_http_transport_entry, + *, + session_routing: str | None, + session_routing_opts: dict | None = None, + ) -> InferenceClient: + model_endpoint = ModelEndpointInfo( + models=ModelListInfo( + models=[ModelInfo(name="test-model")], + model_selection_strategy=ModelSelectionStrategy.ROUND_ROBIN, + ), + endpoint=EndpointInfo( + type=EndpointType.CHAT, + base_url="http://localhost:8000/v1/test", + session_routing=session_routing, + session_routing_opts=session_routing_opts or {}, + ), + ) + mock_transport = MagicMock() + mock_endpoint = MagicMock() + mock_endpoint.get_endpoint_headers.return_value = {} + mock_endpoint.get_endpoint_params.return_value = {} + mock_endpoint.format_payload.return_value = { + "model": "test-model", + "messages": [{"role": "user", "content": "hello"}], + } + + def mock_get_class(protocol, name): + if protocol == "endpoint": + return lambda **kwargs: mock_endpoint + if protocol == "transport": + return lambda **kwargs: mock_transport + if protocol == "session_routing": + return _ROUTING_CLASSES[name] + raise ValueError(f"Unknown protocol: {protocol}") + + with ( + patch( + "aiperf.workers.inference_client.plugins.get_class", + side_effect=mock_get_class, + ), + patch( + "aiperf.workers.inference_client.plugins.list_entries", + return_value=[mock_http_transport_entry], + ), + ): + client = InferenceClient( + model_endpoint=model_endpoint, service_id="test-service-id" + ) + client.transport.send_request = AsyncMock(return_value=RequestRecord()) + return client + + def _request_info( + self, + client: InferenceClient, + *, + x_correlation_id: str = "corr-1", + parent_correlation_id: str | None = None, + is_final_turn: bool = False, + raw_payload: dict | None = None, + ) -> RequestInfo: + turn = ( + Turn(role="user", raw_payload=raw_payload) + if raw_payload is not None + else Turn(role="user", texts=[Text(contents=["hello"])]) + ) + return RequestInfo( + model_endpoint=client.model_endpoint, + turns=[turn], + turn_index=0, + credit_num=0, + credit_phase=CreditPhase.PROFILING, + x_request_id="req-1", + x_correlation_id=x_correlation_id, + parent_correlation_id=parent_correlation_id, + conversation_id="conv-template", + is_final_turn=is_final_turn, + ) + + def _sent_payload(self, client: InferenceClient) -> dict: + return client.transport.send_request.call_args.kwargs["payload"] + + @pytest.mark.asyncio + async def test_dynamo_headers_mode_emits_headers_and_leaves_body( + self, mock_http_transport_entry + ): + client = self._build_client( + mock_http_transport_entry, session_routing="dynamo_headers" + ) + request_info = self._request_info( + client, parent_correlation_id="parent-corr", is_final_turn=False + ) + + await client._send_request_to_transport(request_info) + + assert request_info.endpoint_headers["X-Dynamo-Session-ID"] == "corr-1" + assert ( + request_info.endpoint_headers["X-Dynamo-Parent-Session-ID"] == "parent-corr" + ) + assert "nvext" not in self._sent_payload(client) + + @pytest.mark.asyncio + async def test_dynamo_headers_root_omits_parent_header( + self, mock_http_transport_entry + ): + client = self._build_client( + mock_http_transport_entry, session_routing="dynamo_headers" + ) + request_info = self._request_info(client, parent_correlation_id=None) + + await client._send_request_to_transport(request_info) + + assert request_info.endpoint_headers["X-Dynamo-Session-ID"] == "corr-1" + assert "X-Dynamo-Parent-Session-ID" not in request_info.endpoint_headers + + @pytest.mark.asyncio + async def test_dynamo_nvext_mode_binds_then_closes(self, mock_http_transport_entry): + client = self._build_client( + mock_http_transport_entry, + session_routing="dynamo_nvext", + session_routing_opts={"timeout_seconds": "123"}, + ) + + non_final = self._request_info(client, is_final_turn=False) + await client._send_request_to_transport(non_final) + assert self._sent_payload(client)["nvext"]["session_control"] == { + "session_id": "corr-1", + "action": "bind", + "timeout": 123, + } + + final = self._request_info(client, is_final_turn=True) + await client._send_request_to_transport(final) + assert self._sent_payload(client)["nvext"]["session_control"] == { + "session_id": "corr-1", + "action": "close", + } + + @pytest.mark.asyncio + async def test_dynamo_nvext_on_raw_payload_turn_does_not_mutate_cache( + self, mock_http_transport_entry + ): + """The nvext transform runs on a cached raw_payload dict; the copy-on-write + contract must leave the shared Turn.raw_payload untouched.""" + client = self._build_client( + mock_http_transport_entry, session_routing="dynamo_nvext" + ) + raw_payload = {"model": "m", "messages": [{"role": "user", "content": "hi"}]} + request_info = self._request_info(client, raw_payload=raw_payload) + + await client._send_request_to_transport(request_info) + + # The wire payload carries the injected session_control ... + assert ( + self._sent_payload(client)["nvext"]["session_control"]["action"] == "bind" + ) + # ... but the cached raw_payload dict is unchanged (no nvext leaked in). + assert "nvext" not in raw_payload + assert request_info.turns[0].raw_payload == { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + } + + @pytest.mark.asyncio + async def test_smg_routing_key_mode_emits_header(self, mock_http_transport_entry): + client = self._build_client( + mock_http_transport_entry, session_routing="smg_routing_key" + ) + request_info = self._request_info(client) + + await client._send_request_to_transport(request_info) + + assert request_info.endpoint_headers["X-SMG-Routing-Key"] == "corr-1" + assert "nvext" not in self._sent_payload(client) + + @pytest.mark.asyncio + async def test_session_id_header_custom_name(self, mock_http_transport_entry): + client = self._build_client( + mock_http_transport_entry, + session_routing="session_id_header", + session_routing_opts={"header_name": "X-Affinity"}, + ) + request_info = self._request_info(client) + + await client._send_request_to_transport(request_info) + + assert request_info.endpoint_headers["X-Affinity"] == "corr-1" + assert "nvext" not in self._sent_payload(client) + + @pytest.mark.asyncio + async def test_routing_unset_no_headers_no_body_change( + self, mock_http_transport_entry + ): + client = self._build_client(mock_http_transport_entry, session_routing=None) + assert client._routing is None + request_info = self._request_info(client, parent_correlation_id="parent-corr") + + await client._send_request_to_transport(request_info) + + payload = self._sent_payload(client) + assert "nvext" not in payload + assert "X-Dynamo-Session-ID" not in request_info.endpoint_headers + assert "X-Dynamo-Parent-Session-ID" not in request_info.endpoint_headers + + @pytest.mark.asyncio + async def test_notify_session_end_reaches_plugin(self, mock_http_transport_entry): + client = self._build_client( + mock_http_transport_entry, session_routing="dynamo_headers" + ) + client._routing.on_session_end = MagicMock() + + # Pass-through must not dedupe: idempotency is the plugin's job. + client.notify_session_end("corr-1") + client.notify_session_end("corr-1") + + assert client._routing.on_session_end.call_count == 2 + client._routing.on_session_end.assert_called_with("corr-1") + + def test_notify_session_end_noop_when_routing_unset( + self, mock_http_transport_entry + ): + client = self._build_client(mock_http_transport_entry, session_routing=None) + # No routing plugin: the hook is a safe no-op (never raises). + client.notify_session_end("corr-1") + + def test_notify_session_end_swallows_plugin_error_and_warns( + self, mock_http_transport_entry + ): + """A raising on_session_end must NOT propagate (core eviction must + proceed); the failure is logged with the plugin + session named.""" + client = self._build_client( + mock_http_transport_entry, session_routing="dynamo_headers" + ) + client._routing.on_session_end = MagicMock(side_effect=RuntimeError("boom")) + + with patch.object(client, "warning") as warn: + # Must not raise. + client.notify_session_end("corr-err") + + client._routing.on_session_end.assert_called_once_with("corr-err") + warn.assert_called_once() + msg = warn.call_args.args[0] + assert "dynamo_headers" in msg and "corr-err" in msg + + @pytest.mark.asyncio + async def test_raising_headers_produces_plugin_attributed_error_record( + self, mock_http_transport_entry + ): + """A plugin exception in headers() surfaces as an error record whose + message names the routing plugin, not the inference server.""" + client = self._build_client( + mock_http_transport_entry, session_routing="dynamo_headers" + ) + client._routing.headers = MagicMock( + side_effect=RuntimeError("bad header build") + ) + request_info = self._request_info(client) + + record = await client._send_request_internal(request_info) + + assert record.error is not None + assert "dynamo_headers" in record.error.message + assert "headers()" in record.error.message + # The transport was never reached (the fault is pre-send). + client.transport.send_request.assert_not_called() + + @pytest.mark.asyncio + async def test_raising_transform_body_produces_plugin_attributed_error_record( + self, mock_http_transport_entry + ): + """A plugin exception in transform_body() is attributed to the plugin.""" + client = self._build_client( + mock_http_transport_entry, session_routing="dynamo_nvext" + ) + client._routing.transform_body = MagicMock( + side_effect=RuntimeError("bad body transform") + ) + request_info = self._request_info(client) + + record = await client._send_request_internal(request_info) + + assert record.error is not None + assert "dynamo_nvext" in record.error.message + assert "transform_body()" in record.error.message + + @pytest.mark.asyncio + async def test_bytes_path_stamps_headers_and_sends_bytes_verbatim( + self, mock_http_transport_entry + ): + """Graph-IR verbatim-bytes payloads still get routing HEADERS, and the + pre-serialized bytes reach the transport untouched (no re-serialize).""" + client = self._build_client( + mock_http_transport_entry, session_routing="dynamo_headers" + ) + body = orjson.dumps({"model": "recorded", "messages": []}) + request_info = self._request_info(client) + request_info.turns[-1].raw_payload_bytes = body + + await client._send_request_to_transport(request_info) + + assert request_info.endpoint_headers["X-Dynamo-Session-ID"] == "corr-1" + assert self._sent_payload(client) is body + assert request_info.payload_bytes is body + + @pytest.mark.asyncio + async def test_bytes_path_skips_body_mutating_mode_with_one_warning( + self, mock_http_transport_entry + ): + """A body-mutating mode (dynamo_nvext) cannot transform verbatim bytes: + the transform is skipped, the bytes ride untouched, and the client + warns exactly ONCE per worker (never silently).""" + client = self._build_client( + mock_http_transport_entry, session_routing="dynamo_nvext" + ) + client.warning = MagicMock() + body = orjson.dumps({"model": "recorded", "messages": []}) + for _ in range(2): + request_info = self._request_info(client) + request_info.turns[-1].raw_payload_bytes = body + await client._send_request_to_transport(request_info) + assert self._sent_payload(client) is body + assert b"session_control" not in request_info.payload_bytes + + client.warning.assert_called_once() + assert "dynamo_nvext" in client.warning.call_args[0][0] diff --git a/tests/unit/workers/test_session_fork_refcount.py b/tests/unit/workers/test_session_fork_refcount.py index d7a30773bd..d1e2a9b77c 100644 --- a/tests/unit/workers/test_session_fork_refcount.py +++ b/tests/unit/workers/test_session_fork_refcount.py @@ -46,6 +46,20 @@ def _conv_with_n_forks(n: int, session_id: str = "root") -> Conversation: ) +def _conv_with_packed_forks(n: int, session_id: str = "root") -> Conversation: + """One FORK branch entry carrying ``n`` children -- the dag_jsonl loader shape.""" + branch = ConversationBranchInfo( + branch_id=f"{session_id}:0", + mode=ConversationBranchMode.FORK, + child_conversation_ids=[f"child-{i}" for i in range(n)], + ) + return Conversation( + session_id=session_id, + turns=[Turn(branch_ids=[branch.branch_id])], + branches=[branch], + ) + + class TestForkRefcount: def test_default_refcount_zero(self, session_manager: UserSessionManager) -> None: session = session_manager.create_and_store( @@ -124,3 +138,76 @@ def test_evict_if_unpinned_unknown_session_is_noop( self, session_manager: UserSessionManager ) -> None: session_manager.evict_if_unpinned("does-not-exist") + + +class TestPendingForkEvictionLateSibling: + """A pending-eviction FORK parent must survive until every declared child pins. + + Regression for the fork-minimal byte-parity race: with two children + forking off the parent's terminal turn, an early child can complete + its WHOLE conversation (pin -> seed -> run -> release) before the + sibling's credit reaches the worker. The release drops the refcount + back to 0 with ``pending_fork_eviction`` set; evicting there yanks + the parent out from under the late sibling, which then dispatches + with NO inherited history (bare user message on the wire). + """ + + @pytest.mark.parametrize( + "conversation_builder", + [ + pytest.param(lambda: _conv_with_n_forks(2), id="one-entry-per-child"), + # The REAL dag_jsonl loader shape: _apply_forks packs every fork + # declared on one turn into a SINGLE ConversationBranchInfo whose + # child_conversation_ids lists all of them. Counting branch + # ENTRIES instead of CHILDREN here read expected=1 and evicted + # the parent after the first child's release. + pytest.param( + lambda: _conv_with_packed_forks(2), id="single-entry-two-children" + ), + ], + ) # fmt: skip + def test_parent_survives_release_until_all_expected_children_pin( + self, session_manager: UserSessionManager, conversation_builder + ) -> None: + session = session_manager.create_and_store( + x_correlation_id="corr-late-sibling", + conversation=conversation_builder(), + num_turns=1, + ) + assert session.fork_children_expected == 2 + + # Parent's terminal turn fired before any child arrived: eviction + # deferred, session stays resident. + session.pending_fork_eviction = True + session_manager.evict_if_unpinned("corr-late-sibling") + assert session_manager.get("corr-late-sibling") is not None + + # Child A pins, seeds, runs its whole conversation, and releases -- + # all before child B's credit reaches the worker. + session_manager.pin_for_fork_child("corr-late-sibling") + session_manager.release_fork_child("corr-late-sibling") + assert session_manager.get("corr-late-sibling") is not None, ( + "parent evicted after the FIRST child's release while the " + "second declared child had not pinned yet" + ) + + # Child B arrives late, pins + seeds, then releases: NOW every + # declared child has joined and the parent is collected. + session_manager.pin_for_fork_child("corr-late-sibling") + session_manager.release_fork_child("corr-late-sibling") + assert session_manager.get("corr-late-sibling") is None + + def test_parent_without_pending_eviction_survives_all_releases( + self, session_manager: UserSessionManager + ) -> None: + """Mid-conversation forks: children join while the parent is still + running (no pending eviction) -- releases never collect the session.""" + session_manager.create_and_store( + x_correlation_id="corr-mid-conv", + conversation=_conv_with_n_forks(2), + num_turns=1, + ) + for _ in range(2): + session_manager.pin_for_fork_child("corr-mid-conv") + session_manager.release_fork_child("corr-mid-conv") + assert session_manager.get("corr-mid-conv") is not None diff --git a/tests/unit/workers/test_session_routing.py b/tests/unit/workers/test_session_routing.py new file mode 100644 index 0000000000..be47095f1d --- /dev/null +++ b/tests/unit/workers/test_session_routing.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from pydantic import ValidationError + +from aiperf.workers.session_routing import ( + DynamoHeadersRouting, + DynamoNvextOptions, + DynamoNvextRouting, + RoutingContext, + SessionIdHeaderOptions, + SessionIdHeaderRouting, + SessionRoutingBase, + SmgRoutingKeyRouting, +) + + +def _ctx(**overrides) -> RoutingContext: + defaults = dict( + x_correlation_id="corr-1", + parent_correlation_id=None, + root_correlation_id="corr-1", + is_final_turn=False, + is_parent_final=None, + is_tree_final=False, + ) + defaults.update(overrides) + return RoutingContext(**defaults) + + +class TestDynamoHeadersRouting: + def test_root_emits_session_header_only(self): + plugin = DynamoHeadersRouting(DynamoHeadersRouting.Options()) + assert plugin.headers(_ctx()) == {"X-Dynamo-Session-ID": "corr-1"} + assert plugin.mutates_body is False + + def test_child_emits_parent_header(self): + plugin = DynamoHeadersRouting(DynamoHeadersRouting.Options()) + headers = plugin.headers(_ctx(parent_correlation_id="parent-1")) + assert headers == { + "X-Dynamo-Session-ID": "corr-1", + "X-Dynamo-Parent-Session-ID": "parent-1", + } + + def test_body_untouched(self): + plugin = DynamoHeadersRouting(DynamoHeadersRouting.Options()) + payload = {"messages": []} + assert plugin.transform_body(payload, _ctx()) is payload + + +class TestDynamoNvextRouting: + def test_non_final_turn_binds_with_timeout(self): + plugin = DynamoNvextRouting(DynamoNvextOptions(timeout_seconds=123)) + merged = plugin.transform_body({"messages": []}, _ctx()) + assert merged["nvext"]["session_control"] == { + "session_id": "corr-1", + "action": "bind", + "timeout": 123, + } + assert plugin.mutates_body is True + + def test_final_turn_closes_without_timeout(self): + plugin = DynamoNvextRouting(DynamoNvextOptions()) + merged = plugin.transform_body({}, _ctx(is_final_turn=True)) + assert merged["nvext"]["session_control"] == { + "session_id": "corr-1", + "action": "close", + } + + def test_never_mutates_input_payload(self): + nested_sc = {"existing": "keep"} + nvext = {"trace": "keep", "session_control": nested_sc} + payload = {"nvext": nvext} + plugin = DynamoNvextRouting(DynamoNvextOptions()) + merged = plugin.transform_body(payload, _ctx()) + assert payload == { + "nvext": {"trace": "keep", "session_control": {"existing": "keep"}} + } + assert nvext == {"trace": "keep", "session_control": {"existing": "keep"}} + assert merged is not payload + assert merged["nvext"]["session_control"]["existing"] == "keep" + + def test_options_default_and_bounds(self): + assert DynamoNvextOptions().timeout_seconds == 300 + with pytest.raises(ValidationError): + DynamoNvextOptions(timeout_seconds=0) + + def test_options_reject_unknown_keys(self): + with pytest.raises(ValidationError): + DynamoNvextOptions(timeout_secs=5) + + def test_typed_options_access(self): + plugin = DynamoNvextRouting(DynamoNvextOptions(timeout_seconds=42)) + assert plugin.options.timeout_seconds == 42 + + def test_legacy_open_contract_opens_once_then_bare_session_id(self): + """contract=open (released Dynamo v1.2.x, SessionAction open/close + only): 'open' (+timeout) on the FIRST request this worker sends for + the session -- open is NOT idempotent -- then bare session_id.""" + plugin = DynamoNvextRouting( + DynamoNvextOptions(contract="open", timeout_seconds=77) + ) + first = plugin.transform_body({}, _ctx()) + assert first["nvext"]["session_control"] == { + "session_id": "corr-1", + "action": "open", + "timeout": 77, + } + second = plugin.transform_body({}, _ctx()) + assert second["nvext"]["session_control"] == {"session_id": "corr-1"} + # A DIFFERENT session opens independently. + other = plugin.transform_body({}, _ctx(x_correlation_id="corr-2")) + assert other["nvext"]["session_control"]["action"] == "open" + + def test_legacy_open_contract_final_closes_and_releases_marker(self): + plugin = DynamoNvextRouting(DynamoNvextOptions(contract="open")) + plugin.transform_body({}, _ctx()) # opens + final = plugin.transform_body({}, _ctx(is_final_turn=True)) + assert final["nvext"]["session_control"] == { + "session_id": "corr-1", + "action": "close", + } + # Marker released: a recycled session with the same corr re-opens. + reopened = plugin.transform_body({}, _ctx()) + assert reopened["nvext"]["session_control"]["action"] == "open" + + def test_legacy_open_marker_released_by_on_session_end(self): + """Cancellation paths never send the final turn; on_session_end is + the terminal hook that must release the per-worker opened marker + (idempotent) so the set cannot grow unboundedly.""" + plugin = DynamoNvextRouting(DynamoNvextOptions(contract="open")) + plugin.transform_body({}, _ctx()) + plugin.on_session_end("corr-1") + plugin.on_session_end("corr-1") # idempotent + reopened = plugin.transform_body({}, _ctx()) + assert reopened["nvext"]["session_control"]["action"] == "open" + + def test_bind_contract_keeps_no_per_session_state(self): + plugin = DynamoNvextRouting(DynamoNvextOptions()) + plugin.transform_body({}, _ctx()) + assert plugin._opened == set() + + def test_options_reject_unknown_contract(self): + with pytest.raises(ValidationError): + DynamoNvextOptions(contract="handshake") + + +class TestSmgRoutingKeyRouting: + def test_emits_routing_key(self): + plugin = SmgRoutingKeyRouting(SmgRoutingKeyRouting.Options()) + assert plugin.headers(_ctx()) == {"X-SMG-Routing-Key": "corr-1"} + + def test_rejects_any_opt(self): + with pytest.raises(ValidationError): + SmgRoutingKeyRouting.Options(anything="x") + + +class TestSessionIdHeaderRouting: + def test_default_header_name(self): + plugin = SessionIdHeaderRouting(SessionIdHeaderOptions()) + assert plugin.headers(_ctx()) == {"X-Session-ID": "corr-1"} + + def test_custom_header_name(self): + plugin = SessionIdHeaderRouting( + SessionIdHeaderOptions(header_name="X-Affinity") + ) + assert plugin.headers(_ctx()) == {"X-Affinity": "corr-1"} + + +class TestBaseDefaults: + def test_on_session_end_default_noop_and_idempotent(self): + class Passthrough(SessionRoutingBase): + pass + + plugin = Passthrough(Passthrough.Options()) + plugin.on_session_end("corr-1") + plugin.on_session_end("corr-1") + + def test_stateful_open_once_lifecycle_expressible(self): + """The legacy-nvext shape: open-once instance state, bounded by on_session_end.""" + + class OpenOnce(SessionRoutingBase): + mutates_body = True + + def __init__(self, options): + super().__init__(options) + self.opened: set[str] = set() + + def transform_body(self, payload, ctx): + merged = dict(payload) + if ctx.x_correlation_id not in self.opened: + self.opened.add(ctx.x_correlation_id) + merged["action"] = "open" + return merged + + def on_session_end(self, x_correlation_id): + self.opened.discard(x_correlation_id) + + plugin = OpenOnce(OpenOnce.Options()) + assert plugin.transform_body({}, _ctx())["action"] == "open" + assert "action" not in plugin.transform_body({}, _ctx()) + plugin.on_session_end("corr-1") + plugin.on_session_end("corr-1") # idempotent + assert plugin.opened == set() diff --git a/tests/unit/workers/test_streamed_stamp.py b/tests/unit/workers/test_streamed_stamp.py new file mode 100644 index 0000000000..7d9190efcf --- /dev/null +++ b/tests/unit/workers/test_streamed_stamp.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from aiperf.common.enums import ModelSelectionStrategy +from aiperf.common.models.model_endpoint_info import ( + EndpointInfo, + ModelEndpointInfo, + ModelInfo, + ModelListInfo, +) +from aiperf.common.models.record_models import ErrorDetails, RequestRecord +from aiperf.plugin.enums import EndpointType, TransportType +from aiperf.workers.inference_client import InferenceClient + + +@pytest.fixture +def mock_http_transport_entry(): + """Create a mock transport entry with http/https url_schemes.""" + entry = MagicMock() + entry.name = TransportType.HTTP.value + entry.metadata = {"url_schemes": ["http", "https"]} + return entry + + +@pytest.fixture +def model_endpoint(): + """Create a test ModelEndpointInfo.""" + return ModelEndpointInfo( + models=ModelListInfo( + models=[ModelInfo(name="test-model")], + model_selection_strategy=ModelSelectionStrategy.ROUND_ROBIN, + ), + endpoint=EndpointInfo( + type=EndpointType.CHAT, + base_url="http://localhost:8000/v1/test", + ), + ) + + +@pytest.fixture +def inference_client(model_endpoint, mock_http_transport_entry): + """Create an InferenceClient with a mocked endpoint and transport.""" + mock_transport = MagicMock() + mock_endpoint = MagicMock() + mock_endpoint.get_endpoint_headers.return_value = {} + mock_endpoint.get_endpoint_params.return_value = {} + mock_endpoint.format_payload.return_value = {} + + def mock_get_class(protocol, name): + if protocol == "endpoint": + return lambda **kwargs: mock_endpoint + if protocol == "transport": + return lambda **kwargs: mock_transport + raise ValueError(f"Unknown protocol: {protocol}") + + with ( + patch( + "aiperf.workers.inference_client.plugins.get_class", + side_effect=mock_get_class, + ), + patch( + "aiperf.workers.inference_client.plugins.list_entries", + return_value=[mock_http_transport_entry], + ), + ): + return InferenceClient( + model_endpoint=model_endpoint, service_id="test-service-id" + ) + + +@pytest.mark.asyncio +async def test_streaming_request_record_stamped_true( + inference_client, sample_request_info +): + """A per-request streaming override of True stamps ``streamed`` True.""" + request_info = sample_request_info + request_info.stream_override = True + inference_client.transport.send_request = AsyncMock( + return_value=RequestRecord(request_info=request_info) + ) + + record = await inference_client.send_request(request_info) + + assert record.streamed is True + + +@pytest.mark.asyncio +async def test_non_streaming_request_record_stamped_false( + inference_client, sample_request_info +): + """A per-request override of False wins over a global streaming=True.""" + request_info = sample_request_info + request_info.model_endpoint.endpoint.streaming = True + request_info.stream_override = False + inference_client.transport.send_request = AsyncMock( + return_value=RequestRecord(request_info=request_info) + ) + + record = await inference_client.send_request(request_info) + + assert record.streamed is False + + +@pytest.mark.asyncio +async def test_error_record_still_stamped_with_effective_mode( + inference_client, sample_request_info +): + """An error record reflects the mode the request was SENT with. + + A mid-stream error is still a streamed send, so the ground-truth stamp + tracks the effective wire mode rather than whether a response arrived. + """ + request_info = sample_request_info + request_info.stream_override = True + error_record = RequestRecord( + request_info=request_info, + error=ErrorDetails(type="TestError", message="boom"), + ) + inference_client.transport.send_request = AsyncMock(return_value=error_record) + + record = await inference_client.send_request(request_info) + + assert record.has_error + assert record.streamed is True diff --git a/tests/unit/workers/test_worker.py b/tests/unit/workers/test_worker.py index aa06724920..4bbaae32be 100644 --- a/tests/unit/workers/test_worker.py +++ b/tests/unit/workers/test_worker.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import asyncio -from unittest.mock import AsyncMock, Mock +from unittest.mock import AsyncMock, MagicMock, Mock import pytest @@ -478,3 +478,184 @@ async def test_falls_back_to_dataset_manager_when_no_client_and_not_stopping( assert result == expected_conversation mock_fallback.assert_called_once_with("test-conv-123", sample_credit_context) + + +# --- Session-Routing Terminal Hook Tests --- + + +@pytest.mark.asyncio +class TestReleaseAndEvictForTerminal: + """The terminal-eviction path fires the routing plugin's post-session hook + (``InferenceClient.notify_session_end``) so stateful plugins release the + session on ANY terminal outcome, even one abandoned before its final turn + (e.g. cancellation).""" + + async def test_release_and_evict_for_terminal_notifies_session_end( + self, mock_worker, sample_credit_context + ): + credit = sample_credit_context.credit + mock_worker.inference_client.notify_session_end = Mock() + + mock_worker._release_and_evict_for_terminal(credit, credit.x_correlation_id) + + mock_worker.inference_client.notify_session_end.assert_called_once_with( + credit.x_correlation_id + ) + + async def test_raising_on_session_end_does_not_block_eviction( + self, mock_worker, sample_credit_context + ): + """A routing plugin whose on_session_end raises must not abort the + worker's eviction: the InferenceClient hook swallows + warns, so the + session still gets evicted.""" + from aiperf.workers.inference_client import InferenceClient + + credit = sample_credit_context.credit + fake_client = MagicMock() + fake_client._routing.on_session_end.side_effect = RuntimeError("plugin boom") + fake_client._routing_mode = "dynamo_headers" + # Drive the REAL hook logic (try/except-log) bound to the fake client. + fake_client.notify_session_end = InferenceClient.notify_session_end.__get__( + fake_client + ) + mock_worker.inference_client = fake_client + mock_worker.session_manager.get = Mock(return_value=None) + mock_worker.session_manager.evict = Mock() + + # Must not raise despite the plugin fault. + mock_worker._release_and_evict_for_terminal(credit, credit.x_correlation_id) + + mock_worker.session_manager.evict.assert_called_once_with( + credit.x_correlation_id + ) + fake_client._routing.on_session_end.assert_called_once_with( + credit.x_correlation_id + ) + fake_client.warning.assert_called_once() + + +class TestSessionRoutingTerminalHooks: + """Every terminal disposition this codebase has reaches + ``InferenceClient.notify_session_end`` so routing plugins get their + idempotent post-session hook: the final-turn / cancellation eviction path + and the done-callback cancel-before-start branch (which the finally block + never sees).""" + + def _done_callback(self, *, returned: bool) -> MagicMock: + """Drive the real done-callback on a mock worker with a not-yet-returned + (or already-returned) credit context.""" + worker = MagicMock() + worker.inference_client = MagicMock() + worker.service_id = "worker-1" + credit = MagicMock() + credit.id = 7 + credit.x_correlation_id = "conv-done" + ctx = MagicMock() + ctx.returned = returned + ctx.credit = credit + ctx.error = None + task = MagicMock() + task.cancelled.return_value = True + Worker._on_credit_drop_message_task_done.__get__(worker)(task, ctx) + return worker + + def test_done_callback_not_returned_branch_notifies_session_end(self) -> None: + """The cancel-before-start path (task done, credit never returned) is a + terminal disposition the finally block never sees, so the done-callback + must fire the hook too.""" + worker = self._done_callback(returned=False) + worker.inference_client.notify_session_end.assert_called_once_with("conv-done") + + def test_done_callback_already_returned_does_not_double_notify(self) -> None: + """When the credit already returned, the finally block already notified; + the done-callback short-circuits without a second hook call.""" + worker = self._done_callback(returned=True) + worker.inference_client.notify_session_end.assert_not_called() + + +@pytest.mark.asyncio +class TestCreateRequestInfoFinality: + """Touch #3: worker -> RequestInfo lineage-finality plumb. + + The Credit is a REAL struct (not a MagicMock, which would auto-create the + attributes and mask a missed plumb) carrying both finality facts plus a + tree root id. + """ + + async def test_create_request_info_plumbs_finality_from_credit(self, mock_worker): + from aiperf.common.models import Turn + + credit_context = CreditContext( + credit=Credit( + id=1, + phase=CreditPhase.PROFILING, + conversation_id="test-conv", + x_correlation_id="child-x", + turn_index=0, + num_turns=1, + issued_at_ns=0, + agent_depth=1, + parent_correlation_id="root-x", + root_correlation_id="root-x", + is_parent_final=True, + is_tree_final=True, + ), + drop_perf_ns=0, + ) + + # Session fields the builder reads, set to real values (the guard is on + # the Credit, so a lightweight session stand-in is fine here). + session = MagicMock() + session.x_correlation_id = "child-x" + session.conversation.session_id = "test-conv" + session.turn_index = 0 + session.turn_list = [Turn()] + session.url_index = None + + request_info = mock_worker._create_request_info( + x_request_id="request-id", + session=session, + credit_context=credit_context, + ) + + assert request_info.is_parent_final is True + assert request_info.is_tree_final is True + # effective_root_correlation_id resolves to the stamped root id. + assert request_info.root_correlation_id == "root-x" + + async def test_create_request_info_root_defaults_to_own_correlation( + self, mock_worker + ): + """A root credit (no root_correlation_id) surfaces its own + x_correlation_id as the record's tree root, with conservative finality. + """ + from aiperf.common.models import Turn + + credit_context = CreditContext( + credit=Credit( + id=2, + phase=CreditPhase.PROFILING, + conversation_id="test-conv", + x_correlation_id="root-x", + turn_index=0, + num_turns=1, + issued_at_ns=0, + ), + drop_perf_ns=0, + ) + session = MagicMock() + session.x_correlation_id = "root-x" + session.conversation.session_id = "test-conv" + session.turn_index = 0 + session.turn_list = [Turn()] + session.url_index = None + + request_info = mock_worker._create_request_info( + x_request_id="request-id", + session=session, + credit_context=credit_context, + ) + + assert request_info.root_correlation_id == "root-x" + assert request_info.is_parent_final is None + assert request_info.is_tree_final is False diff --git a/tests/unit/workers/test_worker_graph_predispatch_errors.py b/tests/unit/workers/test_worker_graph_predispatch_errors.py new file mode 100644 index 0000000000..4d7cade959 --- /dev/null +++ b/tests/unit/workers/test_worker_graph_predispatch_errors.py @@ -0,0 +1,269 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pre-dispatch graph credit failures must emit a synthetic error record (WK2). + +RecordsManager's completion barrier counts success+error RECORDS against the +credit-side ``final_requests_completed``; a graph credit that errors BEFORE +dispatch (missing store, missing envelope, missing pool entry) returns its +credit with an error but -- without the fix -- produced no ``RequestRecord``, +starving the barrier and hanging the run at "please wait for the results". +These tests drive the real ``Worker._process_graph_credit`` / +``Worker._send_graph_error_record`` bound onto a mock self and assert every +pre-dispatch error path pushes exactly one error ``InferenceResultsMessage`` +record while preserving the credit-return error attribution. +""" + +from __future__ import annotations + +import types +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +import orjson +import pytest + +from aiperf.common.enums import CreditPhase, ModelSelectionStrategy +from aiperf.common.models import ErrorDetails, RequestRecord +from aiperf.common.models.model_endpoint_info import ( + EndpointInfo, + ModelEndpointInfo, + ModelInfo, + ModelListInfo, +) +from aiperf.credit.structs import Credit, CreditContext +from aiperf.dataset.graph_segment_unified_store import ( + GraphSegmentUnifiedBackingStore, + GraphSegmentUnifiedClient, +) +from aiperf.graph.dynamic_pool import GraphDynamicPool, GraphPoolMissingError +from aiperf.plugin.enums import EndpointType +from aiperf.workers.worker import Worker + +TRACE = "t-1" +# Instance id is ``{template}::{nonce}``; the worker strips the ``::{nonce}`` +# back to the base template id for catalog / store lookups. +INSTANCE = "t-1::inst0" + + +def _model_endpoint() -> ModelEndpointInfo: + """A real (validating) ModelEndpointInfo for the synthetic RequestInfo.""" + return ModelEndpointInfo( + models=ModelListInfo( + models=[ModelInfo(name="test-model")], + model_selection_strategy=ModelSelectionStrategy.ROUND_ROBIN, + ), + endpoint=EndpointInfo( + type=EndpointType.CHAT, + base_url="http://localhost:8000/v1", + ), + ) + + +def _credit_context(node_ordinal: int = 0) -> CreditContext: + return CreditContext( + credit=Credit( + id=1, + phase=CreditPhase.PROFILING, + conversation_id=TRACE, + x_correlation_id="t-1::corr0", + turn_index=0, + num_turns=1, + issued_at_ns=0, + trace_id=INSTANCE, + node_ordinal=node_ordinal, + ), + drop_perf_ns=0, + ) + + +def _mock_worker(store_reader) -> MagicMock: + """Mock self carrying the REAL pre-dispatch error methods under test.""" + self = MagicMock() + self._graph_store_reader = store_reader + self._graph_dynamic_pool = GraphDynamicPool(max_bytes=1024 * 1024) + self.model_endpoint = _model_endpoint() + self._send_inference_result_message = AsyncMock() + self.inference_client.send_request = AsyncMock() + # Bind the real methods under test onto the mock self. + self._process_graph_credit = types.MethodType(Worker._process_graph_credit, self) + self._send_graph_error_record = types.MethodType( + Worker._send_graph_error_record, self + ) + self._set_graph_envelope_missing = types.MethodType( + Worker._set_graph_envelope_missing, self + ) + return self + + +def _sole_sent_record(self: MagicMock) -> RequestRecord: + assert self._send_inference_result_message.await_count == 1 + record = self._send_inference_result_message.await_args.args[0] + assert isinstance(record, RequestRecord) + return record + + +async def _client_with_node( + tmp_path: Path, envelope_extra: dict | None = None +) -> GraphSegmentUnifiedClient: + """A real finalized unified store carrying ONE node manifest at ordinal 0.""" + store = GraphSegmentUnifiedBackingStore(base_path=tmp_path, benchmark_id="bench") + handle = store.put_segment("seg0", "user", "hello") + store.add_node_manifest( + TRACE, + 0, + "profiling", + orjson.dumps({"handles": [handle], **(envelope_extra or {})}), + ) + await store.finalize() + return GraphSegmentUnifiedClient(tmp_path, "bench").open() + + +@pytest.mark.asyncio +async def test_store_missing_credit_emits_error_record() -> None: + """A missing graph store sends a synthetic error InferenceResults record.""" + store_error = ErrorDetails( + type="GraphStoreUnavailable", message="no store could be opened" + ) + + def _reader(credit_context: CreditContext) -> None: + # Mirrors the real _graph_store_reader failure contract: attribute the + # error on the context and return None. + credit_context.error = store_error + return None + + self = _mock_worker(MagicMock(side_effect=_reader)) + ctx = _credit_context() + await self._process_graph_credit(ctx, "x-req-1", None) + + record = _sole_sent_record(self) + assert record.error is store_error + assert record.valid is False + assert record.request_info.x_request_id == "x-req-1" + assert record.request_info.x_correlation_id == ctx.credit.x_correlation_id + # Credit-return semantics preserved: the context still carries the error. + assert ctx.error is store_error + # Nothing was dispatched to the inference server. + self.inference_client.send_request.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_envelope_missing_credit_emits_error_record(tmp_path: Path) -> None: + """An unaddressable node ordinal sends a GraphEnvelopeMissing error record.""" + client = await _client_with_node(tmp_path) + self = _mock_worker(MagicMock(return_value=client)) + ctx = _credit_context(node_ordinal=7) + await self._process_graph_credit(ctx, "x-req-2", None) + + record = _sole_sent_record(self) + assert record.error is not None + assert record.error.type == "GraphEnvelopeMissing" + assert record.valid is False + assert ctx.error is record.error + self.inference_client.send_request.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_pool_missing_credit_emits_error_record( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A missing dynamic-pool entry sends an error record carrying the sniffable prefix.""" + import aiperf.graph.worker_materialize as wm + + client = await _client_with_node(tmp_path, envelope_extra={"items": []}) + + def _raise_pool_missing(*args, **kwargs): + raise GraphPoolMissingError(3) + + monkeypatch.setattr(wm, "materialize_graph_request_unified", _raise_pool_missing) + self = _mock_worker(MagicMock(return_value=client)) + ctx = _credit_context() + await self._process_graph_credit(ctx, "x-req-3", None) + + record = _sole_sent_record(self) + assert record.error is not None + assert record.valid is False + assert "aiperf.graph.pool_missing:" in record.error.message + # The context keeps the raw prefixed string the dispatch adapter sniffs. + assert isinstance(ctx.error, str) + assert ctx.error.startswith("aiperf.graph.pool_missing:") + self.inference_client.send_request.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_successful_dispatch_sends_single_record(tmp_path: Path) -> None: + """Control: the happy path still sends exactly one (dispatch-built) record.""" + client = await _client_with_node(tmp_path) + self = _mock_worker(MagicMock(return_value=client)) + self._dispatch_graph_request = AsyncMock() + self._build_graph_request_info = MagicMock(return_value=MagicMock()) + ctx = _credit_context() + await self._process_graph_credit(ctx, "x-req-4", None) + + self._dispatch_graph_request.assert_awaited_once() + # No synthetic pre-dispatch record on the happy path. + self._send_inference_result_message.assert_not_awaited() + assert ctx.error is None + + +@pytest.mark.asyncio +async def test_unexpected_exception_emits_error_record_and_credit_attribution( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An UNANTICIPATED raiser inside ``_process_graph_credit`` (here a corrupt + envelope failing ``orjson.loads`` in ``read_node_envelope``) must still emit + exactly one error record AND attribute the error on the CreditReturn -- + without the catch-all the credit returns error=None (counted as a success) + while no record flows, starving the RecordsManager barrier.""" + import aiperf.graph.worker_materialize as wm + + def _raise_decode_error(*args, **kwargs): + raise orjson.JSONDecodeError("corrupt envelope bytes", "", 0) + + monkeypatch.setattr(wm, "read_node_envelope", _raise_decode_error) + self = _mock_worker(MagicMock(return_value=MagicMock())) + self._prefill_concurrency_enabled = False + self.credit_return_push_client.send = AsyncMock() + # Drive the full credit task so the CreditReturn built in its finally is + # observable alongside the record. + self._process_credit = types.MethodType(Worker._process_credit, self) + self._on_credit_drop_message_task = types.MethodType( + Worker._on_credit_drop_message_task, self + ) + ctx = _credit_context() + + await self._on_credit_drop_message_task(ctx) + + record = _sole_sent_record(self) + assert record.error is not None + assert record.valid is False + assert "corrupt envelope bytes" in record.error.message + # Credit-return attribution: the barrier-side counterpart must see an error. + assert ctx.error is not None + credit_return = self.credit_return_push_client.send.await_args.args[0] + assert credit_return.error is not None, ( + "an escaped exception must not be counted as a completed request" + ) + self.inference_client.send_request.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_failure_after_dispatch_record_does_not_double_emit( + tmp_path: Path, +) -> None: + """A raiser AFTER the dispatch path resolved (here the pool bracket close) + attributes the error but must NOT emit a second, synthetic record.""" + client = await _client_with_node(tmp_path) + self = _mock_worker(MagicMock(return_value=client)) + self._dispatch_graph_request = AsyncMock() + self._build_graph_request_info = MagicMock(return_value=MagicMock()) + self._graph_dynamic_pool.credit_finished = MagicMock( + side_effect=RuntimeError("post-record boom") + ) + ctx = _credit_context() + + await self._process_graph_credit(ctx, "x-req-5", None) + + self._dispatch_graph_request.assert_awaited_once() + # The dispatch path owns the record; the catch-all must not add another. + self._send_inference_result_message.assert_not_awaited() + assert ctx.error is not None diff --git a/tests/unit/workers/test_worker_graph_store_discovery.py b/tests/unit/workers/test_worker_graph_store_discovery.py new file mode 100644 index 0000000000..fc171823e5 --- /dev/null +++ b/tests/unit/workers/test_worker_graph_store_discovery.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Worker opens the graph unified store from the broadcast, not env conventions.""" + +from types import SimpleNamespace + +from aiperf.common.models.dataset_models import GraphSegmentClientMetadata +from aiperf.workers.worker import Worker + + +def _bare_worker(meta: GraphSegmentClientMetadata | None) -> Worker: + w = Worker.__new__(Worker) + w.run = SimpleNamespace(benchmark_id="b1") + w._graph_client_metadata = meta + w._graph_unified_open_attempted = False + w._graph_unified_client = None + w._graph_unified_open_failure = None + w.warning = lambda *a, **k: None + w.debug = lambda *a, **k: None + return w + + +def test_reader_without_graph_broadcast_records_failure(tmp_path): + w = _bare_worker(None) + assert w._graph_unified_reader() is None + # Absent metadata records an actionable failure (feeds GraphStoreUnavailable), + # never an env fallback. + assert "GraphSegmentClientMetadata" in (w._graph_unified_open_failure or "") + + +def test_reader_opens_from_broadcast_paths(tmp_path, monkeypatch): + meta = GraphSegmentClientMetadata( + store_base_path=tmp_path, + benchmark_id="b1", + sidecar_path=tmp_path / "aiperf_graph_meta_b1" / "graph_meta.msgpack", + ) + w = _bare_worker(meta) + captured = {} + + class _FakeClient: + def __init__(self, *, base_path, benchmark_id): + captured["base_path"] = base_path + captured["benchmark_id"] = benchmark_id + + def open(self): + return self + + # worker.py binds ``GraphSegmentUnifiedClient`` at import time, so patch the + # name in the worker module, not its defining module. + monkeypatch.setattr( + "aiperf.workers.worker.GraphSegmentUnifiedClient", + _FakeClient, + ) + assert w._graph_unified_reader() is not None + assert captured == {"base_path": tmp_path, "benchmark_id": "b1"} diff --git a/tools/dag_jsonl_fidelity.py b/tools/dag_jsonl_fidelity.py new file mode 100644 index 0000000000..f724c6b698 --- /dev/null +++ b/tools/dag_jsonl_fidelity.py @@ -0,0 +1,497 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Empirical-proof tool: legacy ``dag_jsonl`` vs graph ``dag_jsonl`` raw-export parity. + +This is the Task-7 EXTERNAL acceptance gate for the dag_jsonl graph adapter. It +does NOT touch the live run path; it reads the deterministic offline artifacts +two profiling runs produced (each a ``profile_export_raw.jsonl``) -- one from the +LEGACY custom-dataset path (``--custom-dataset-type dag_jsonl``) and one from the +GRAPH adapter path (``--graph-format dag_jsonl``) -- and proves, per matching +request keyed by ``(session_id, turn_index)``, three independent fidelity +criteria plus two coverage invariants: + +1. :attr:`ParityReport.messages` (criterion 1) -- the exported ``payload.messages`` + are byte-identical (``orjson.dumps`` equal) between the two planes. +2. :attr:`ParityReport.payload` (criterion 2) -- the full parsed request payloads + are equal (order-insensitive dict equality: same keys, same values). +3. :attr:`ParityReport.body` (criterion 3) -- the CANONICAL re-serializations are + byte-identical: ``orjson.dumps(payload, option=orjson.OPT_SORT_KEYS)`` on each + side. Key ORDER is deliberately NOT compared (the graph plane authors the + stream flag / token cap / model through native node fields, so their wire + positions differ from legacy); canonical bytes stay strict on VALUES AND + TYPES -- e.g. ``1`` vs ``1.0`` or ``True`` vs ``1`` serialize differently + even though Python ``==`` (criterion 2) treats them as equal. + +Plus COUNT equality (both planes dispatched the same number of PROFILING requests) +and NO-UNMATCHED-KEYS (every ``(session, turn)`` present on one plane is present +on the other, with the SAME multiplicity on both). A proof that compared ZERO +requests FAILS -- an empty overlap must never pass vacuously. + +Node-identity recovery (the linchpin): the two planes carry request identity +differently in the raw export. + +* LEGACY rows carry identity directly: ``metadata.conversation_id`` is the dag + ``session_id`` and ``metadata.turn_index`` is the turn within that session. +* GRAPH rows carry the trajectory TEMPLATE id in ``metadata.conversation_id`` + (shared across every node of a trajectory and duplicated across recycles), so + the per-node identity is recovered from ``metadata.x_request_id``, which the + worker's ``_mint_x_request_id`` folds as ``{node_id}::{nonce}`` (the nonce is a + ``uuid4().hex`` keeping the id fresh per dispatch, and it never contains ``::``). + The node id is ``"[#n]:"`` (the ``#n`` instance suffix appears + only for repeated SPAWN instances); we split off the trailing ``::`` from + the RIGHT, strip the ``#n`` suffix, and split the trailing ``:`` back + out. Rows whose ``x_request_id`` lacks the ``::`` nonce delimiter (or a numeric + trailing turn) are non-graph/malformed and keep a loud sentinel key so coverage + failures never pass vacuously. + +REPEATED SPAWN TEMPLATES -- multiset keying: stripping the ``#n`` suffix collapses +every instance of one SPAWN template onto the same ``(session, turn)`` key, so a +template instantiated N times per tree contributes multiplicity N to that key. The +proof keys each plane into ``dict[key, list[body]]`` and, per key, sorts both +lists by CANONICAL serialized body and compares them ELEMENT-WISE: parity therefore +means the two planes emitted the SAME multiset of bodies for that key, and the +NO-UNMATCHED-KEYS invariant additionally requires equal per-key MULTIPLICITY (a +key fired twice on one plane and once on the other fails). This works without +instance-level pairing because a repeated template's instances carry identical +authored bytes; a corrupted instance breaks the sorted element-wise equality and +surfaces with an instance-indexed pointed diff. + +Runtime multiplicity assumption (verified by the ``spawn-repeat`` gate case in +``tests/component_integration/graph/test_dag_jsonl_byte_parity.py``): the LEGACY +plane keys each spawn-child dispatch by its TEMPLATE ``conversation_id`` (the dag +``session_id``), not an instance-unique id, so a template spawned N times emits N +rows under the same ``(session, turn)`` legacy key -- matching the graph plane's N +``#n`` node instances. If a future runtime change makes legacy conversation ids +instance-unique, :func:`_legacy_key` must strip that instance discriminator to +keep the multiplicities aligned. + +Run offline against any pair of artifacts; the live runs that produce them are +driven by ``tests/component_integration/graph/test_dag_jsonl_fidelity.py``. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import orjson + +_PROFILING = "profiling" + +Key = tuple[str, int] + + +# --- report --------------------------------------------------------------- + + +@dataclass +class Mismatch: + """One failing comparison: which request key and why (pointed, not a blob). + + ``instance`` is the position within a multiplicity>1 key's sorted body list + (``None`` for a single-instance key), so a repeated SPAWN template's failing + instance is pinpointed rather than blurred across its siblings. + """ + + where: str + detail: str + instance: int | None = None + + +@dataclass +class Criterion: + """Structured pass/fail result of one parity criterion. + + ``checked`` counts every ``(session, turn)`` compared; ``passes`` counts the + subset that matched. ``passed`` requires at least one comparison AND zero + mismatches -- a criterion that checked nothing is not a pass. + """ + + name: str + checked: int = 0 + passes: int = 0 + mismatches: list[Mismatch] = field(default_factory=list) + + @property + def passed(self) -> bool: + return not self.mismatches and self.checked > 0 + + def fail(self, where: str, detail: str, *, instance: int | None = None) -> None: + self.mismatches.append(Mismatch(where=where, detail=detail, instance=instance)) + + def render(self) -> str: + head = ( + f"{self.name}: {'PASS' if self.passed else 'FAIL'} " + f"(checked={self.checked}, passed={self.passes}, " + f"mismatches={len(self.mismatches)})" + ) + lines = [head] + for m in self.mismatches: + inst = f" instance#{m.instance}" if m.instance is not None else "" + lines.append(f" MISMATCH @ {m.where}{inst}: {m.detail}") + return "\n".join(lines) + + +@dataclass +class ParityReport: + """Full legacy-vs-graph parity result: three criteria + coverage invariants. + + ``passed`` is the conjunction of: matching PROFILING request COUNTS (and at + least one request), NO unmatched keys AND equal per-key multiplicity, and + every criterion passing. Any single failure fails the whole proof; the CLI + exit code is ``0`` iff ``passed``. + """ + + legacy_profiling: int + graph_profiling: int + only_legacy: list[Key] = field(default_factory=list) + only_graph: list[Key] = field(default_factory=list) + multiplicity_mismatch: list[tuple[Key, int, int]] = field(default_factory=list) + """``(key, legacy_count, graph_count)`` for each shared key whose per-key + multiplicity differs between the planes (a repeated SPAWN template fired a + different number of times on each side).""" + messages: Criterion = field( + default_factory=lambda: Criterion("messages_byte_equal") + ) + payload: Criterion = field( + default_factory=lambda: Criterion("payload_parsed_equal") + ) + body: Criterion = field(default_factory=lambda: Criterion("body_canonical_equal")) + + @property + def counts_match(self) -> bool: + """Both planes dispatched the same NONZERO number of PROFILING requests.""" + return ( + self.legacy_profiling == self.graph_profiling and self.legacy_profiling > 0 + ) + + @property + def keys_match(self) -> bool: + """Same key SET on both planes AND the same per-key multiplicity.""" + return not (self.only_legacy or self.only_graph or self.multiplicity_mismatch) + + @property + def passed(self) -> bool: + return ( + self.counts_match + and self.keys_match + and self.messages.passed + and self.payload.passed + and self.body.passed + ) + + def render(self) -> str: + head = ( + f"dag_jsonl parity: {'PASS' if self.passed else 'FAIL'} " + f"(legacy_profiling={self.legacy_profiling}, " + f"graph_profiling={self.graph_profiling})" + ) + lines = [head] + if not self.counts_match: + lines.append( + f" COUNT MISMATCH: legacy={self.legacy_profiling} " + f"graph={self.graph_profiling}" + ) + for key in self.only_legacy: + lines.append(f" ONLY-LEGACY key {key} (missing from graph export)") + for key in self.only_graph: + lines.append(f" ONLY-GRAPH key {key} (missing from legacy export)") + for key, legacy_n, graph_n in self.multiplicity_mismatch: + lines.append( + f" MULTIPLICITY MISMATCH key {key}: " + f"legacy fired {legacy_n}x, graph fired {graph_n}x" + ) + lines.append(self.messages.render()) + lines.append(self.payload.render()) + lines.append(self.body.render()) + return "\n".join(lines) + + +# --- raw export record ---------------------------------------------------- + + +@dataclass +class _ExportRecord: + """The parity-relevant projection of one raw-export JSONL line. + + ``payload`` is the FULL wire request dict (criteria 2 and 3 need every key, + not just ``messages``). ``conversation_id`` / ``turn_index`` carry the legacy + identity directly; ``x_request_id`` carries the graph node identity folded by + the worker's ``_mint_x_request_id`` as ``{node_id}::{nonce}``. + """ + + conversation_id: str + turn_index: int | None + x_request_id: str | None + phase: str + payload: dict[str, Any] + + @property + def messages(self) -> list[dict[str, Any]]: + return self.payload.get("messages", []) or [] + + +def load_raw_export(raw_jsonl: Path) -> list[_ExportRecord]: + """Parse a ``profile_export_raw.jsonl`` into parity projections. + + Each line is one dispatched request. ``metadata`` carries ``conversation_id`` + / ``turn_index`` / ``x_request_id`` / ``benchmark_phase``; ``payload`` is the + full wire request dict. Blank lines are skipped. + """ + records: list[_ExportRecord] = [] + for line in Path(raw_jsonl).read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + obj = orjson.loads(line) + meta = obj.get("metadata", {}) or {} + payload = obj.get("payload", {}) or {} + records.append( + _ExportRecord( + conversation_id=str(meta.get("conversation_id") or ""), + turn_index=meta.get("turn_index"), + x_request_id=meta.get("x_request_id"), + phase=str(meta.get("benchmark_phase") or ""), + payload=payload, + ) + ) + return records + + +def _legacy_key(rec: _ExportRecord) -> Key: + """Legacy identity straight off the record: ``(conversation_id, turn_index)``.""" + return (rec.conversation_id, int(rec.turn_index or 0)) + + +def _graph_key(rec: _ExportRecord) -> Key | None: + """Recover ``(session_id, turn_index)`` from the graph ``x_request_id``. + + Per-dispatch node identity rides ``x_request_id`` as ``{node_id}::{nonce}`` + (the worker's ``_mint_x_request_id``); the node id is the dag coordinate + ``"[#n]:"`` and the nonce is a ``uuid4().hex`` with no ``::``. + We split off the trailing ``::`` from the RIGHT, then split the node + id's trailing ``:`` and strip any ``#n`` SPAWN-instance suffix from the + session. Returns ``None`` for any x_request_id lacking the ``::`` nonce + delimiter or a numeric trailing turn (non-graph/malformed rows). + """ + xreq = rec.x_request_id + if not xreq or "::" not in xreq: + return None + node_id = xreq.rsplit("::", 1)[0] + session_part, sep, turn_str = node_id.rpartition(":") + if not sep or not turn_str.isdigit(): + return None + session = session_part.split("#", 1)[0] + return (session, int(turn_str)) + + +def _index_legacy( + records: list[_ExportRecord], +) -> dict[Key, list[dict[str, Any]]]: + """Index legacy PROFILING payloads by ``(session, turn)`` into a multiset. + + A template spawned N times keys N payloads under the same ``(session, turn)`` + (legacy carries the TEMPLATE ``conversation_id``, not an instance-unique id), + so the value is a LIST -- one entry per dispatch under that key. + """ + keyed: dict[Key, list[dict[str, Any]]] = {} + for rec in records: + keyed.setdefault(_legacy_key(rec), []).append(rec.payload) + return keyed + + +def _index_graph( + records: list[_ExportRecord], +) -> dict[Key, list[dict[str, Any]]]: + """Index graph PROFILING payloads by recovered ``(session, turn)`` into a multiset. + + Repeated SPAWN instances (distinct ``#n`` node ids) strip to the same key and + accumulate in the key's payload list. A record whose ``x_request_id`` does not + parse is keyed under a UNIQUE sentinel so it surfaces as an ONLY-GRAPH + unmatched key rather than being silently dropped (an identity-scheme drift + must not pass vacuously). + """ + keyed: dict[Key, list[dict[str, Any]]] = {} + for i, rec in enumerate(records): + key = _graph_key(rec) + if key is None: + key = (f"", -1) + keyed.setdefault(key, []).append(rec.payload) + return keyed + + +# --- pointed diffs -------------------------------------------------------- + + +def _snip(value: Any, limit: int = 80) -> str: + """A short repr of a value for a pointed (non-blob) mismatch detail.""" + text = repr(value) + return text if len(text) <= limit else text[: limit - 3] + "..." + + +def _messages_diff(legacy: list[dict[str, Any]], graph: list[dict[str, Any]]) -> str: + """A compact description of the first differing message.""" + if len(legacy) != len(graph): + return f"message count graph={len(graph)} != legacy={len(legacy)}" + for idx, (le, gr) in enumerate(zip(legacy, graph, strict=False)): + if le != gr: + if le.get("role") != gr.get("role"): + return ( + f"msg[{idx}] role graph={gr.get('role')!r} != " + f"legacy={le.get('role')!r}" + ) + return ( + f"msg[{idx}] role={gr.get('role')!r} content " + f"graph={_snip(gr.get('content'))} != " + f"legacy={_snip(le.get('content'))}" + ) + return "messages differ under byte compare but no positional field diff found" + + +def _payload_diff(legacy: dict[str, Any], graph: dict[str, Any]) -> str: + """A compact description of the first payload key that differs in value/presence.""" + only_legacy = sorted(set(legacy) - set(graph)) + only_graph = sorted(set(graph) - set(legacy)) + if only_legacy: + return f"key {only_legacy[0]!r} present in legacy, absent in graph" + if only_graph: + return f"key {only_graph[0]!r} present in graph, absent in legacy" + for key in legacy: + if legacy[key] != graph[key]: + if key == "messages": + return f"messages differ: {_messages_diff(legacy[key], graph[key])}" + return ( + f"key {key!r}: graph={_snip(graph[key])} != legacy={_snip(legacy[key])}" + ) + return "payloads differ under equality but no key diff found" + + +def _canonical(payload: dict[str, Any]) -> bytes: + """Canonical order-insensitive serialization: sorted keys at every level.""" + return orjson.dumps(payload, option=orjson.OPT_SORT_KEYS) + + +def _body_diff(legacy: dict[str, Any], graph: dict[str, Any]) -> str: + """Describe a parsed-equal but canonically byte-different body (value/type + drift Python ``==`` conflates, e.g. 1 vs 1.0 or True vs 1).""" + return ( + f"parsed-equal but CANONICALLY byte-different (value/type drift): " + f"graph={_canonical(graph)!r} != legacy={_canonical(legacy)!r}" + ) + + +# --- the proof ------------------------------------------------------------ + + +def prove_parity(legacy_export: Path, graph_export: Path) -> ParityReport: + """Prove legacy-vs-graph dag_jsonl raw-export parity, keyed ``(session, turn)``. + + Loads both exports, keeps only PROFILING records (the comparable subset -- + warmup dispatch counts may legitimately differ), keys each plane into a + ``dict[key, list[body]]`` multiset (:func:`_index_legacy` / :func:`_index_graph`), + then for every common key sorts both bodies lists by serialized bytes and + checks the three criteria ELEMENT-WISE across the aligned pairs. Coverage + invariants (count equality, no unmatched keys, equal per-key multiplicity) are + recorded on the returned :class:`ParityReport`. + + Repeated SPAWN templates are supported: their ``#n`` instances strip to one + ``(session, turn)`` key and are compared as a multiset. A per-key multiplicity + difference (fired N times on one plane, M on the other) is a + MULTIPLICITY-MISMATCH failure; the criteria are compared only for keys whose + multiplicities agree, so a corrupted instance under an agreeing key still fails + loudly with an instance-indexed pointed diff. + """ + legacy_recs = [r for r in load_raw_export(legacy_export) if r.phase == _PROFILING] + graph_recs = [r for r in load_raw_export(graph_export) if r.phase == _PROFILING] + + legacy_by_key = _index_legacy(legacy_recs) + graph_by_key = _index_graph(graph_recs) + + only_legacy = sorted(set(legacy_by_key) - set(graph_by_key)) + only_graph = sorted(set(graph_by_key) - set(legacy_by_key)) + + report = ParityReport( + legacy_profiling=len(legacy_recs), + graph_profiling=len(graph_recs), + only_legacy=only_legacy, + only_graph=only_graph, + ) + + for key in sorted(set(legacy_by_key) & set(graph_by_key)): + legacy_bodies = sorted(legacy_by_key[key], key=_canonical) + graph_bodies = sorted(graph_by_key[key], key=_canonical) + if len(legacy_bodies) != len(graph_bodies): + report.multiplicity_mismatch.append( + (key, len(legacy_bodies), len(graph_bodies)) + ) + continue + + multi = len(legacy_bodies) > 1 + where = f"session={key[0]!r} turn={key[1]}" + for idx, (legacy_payload, graph_payload) in enumerate( + zip(legacy_bodies, graph_bodies, strict=True) + ): + instance = idx if multi else None + + # Criterion 1: messages byte-equal. + report.messages.checked += 1 + legacy_msgs = legacy_payload.get("messages", []) or [] + graph_msgs = graph_payload.get("messages", []) or [] + if orjson.dumps(legacy_msgs) != orjson.dumps(graph_msgs): + report.messages.fail( + where, _messages_diff(legacy_msgs, graph_msgs), instance=instance + ) + else: + report.messages.passes += 1 + + # Criterion 2: full payload parsed-equal. + report.payload.checked += 1 + if legacy_payload != graph_payload: + report.payload.fail( + where, + _payload_diff(legacy_payload, graph_payload), + instance=instance, + ) + else: + report.payload.passes += 1 + + # Criterion 3: canonical (sorted-keys) body byte-equal. + report.body.checked += 1 + if _canonical(legacy_payload) != _canonical(graph_payload): + report.body.fail( + where, _body_diff(legacy_payload, graph_payload), instance=instance + ) + else: + report.body.passes += 1 + + return report + + +def _main(argv: list[str] | None = None) -> int: + """CLI: ``python tools/dag_jsonl_fidelity.py ``.""" + import argparse + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "legacy", type=Path, help="legacy (--custom-dataset-type dag_jsonl) raw export" + ) + parser.add_argument( + "graph", type=Path, help="graph (--graph-format dag_jsonl) raw export" + ) + args = parser.parse_args(argv) + report = prove_parity(args.legacy, args.graph) + print(report.render()) + return 0 if report.passed else 1 + + +__all__ = [ + "Criterion", + "Mismatch", + "ParityReport", + "load_raw_export", + "prove_parity", +] + + +if __name__ == "__main__": + raise SystemExit(_main()) diff --git a/tools/ergonomics_baseline.json b/tools/ergonomics_baseline.json index 3585618a31..e7940384bd 100644 --- a/tools/ergonomics_baseline.json +++ b/tools/ergonomics_baseline.json @@ -46,6 +46,51 @@ "src/aiperf/dataset/agentic_code_gen/reporting/cache_explorer.py", "_classify_turn_blocks" ], + [ + "keyword-only-args", + "src/aiperf/dataset/graph/native_lowering.py", + "_assemble_prompt" + ], + [ + "keyword-only-args", + "src/aiperf/dataset/graph/native_lowering.py", + "_content_parts" + ], + [ + "keyword-only-args", + "src/aiperf/dataset/graph/native_lowering.py", + "_delta_messages" + ], + [ + "keyword-only-args", + "src/aiperf/dataset/graph/native_lowering.py", + "_gate_array_splice_shape" + ], + [ + "keyword-only-args", + "src/aiperf/dataset/graph/native_lowering.py", + "_gate_init_bearing_root_reads" + ], + [ + "keyword-only-args", + "src/aiperf/dataset/graph/native_lowering.py", + "_ordered_writer_chain" + ], + [ + "keyword-only-args", + "src/aiperf/dataset/graph/native_lowering.py", + "_stamp_node" + ], + [ + "keyword-only-args", + "src/aiperf/dataset/graph/segment_ir/trie_content.py", + "_assemble_messages_from" + ], + [ + "keyword-only-args", + "src/aiperf/dataset/graph/segment_ir/trie_content.py", + "assemble_messages" + ], [ "keyword-only-args", "src/aiperf/exporters/metrics_csv_exporter.py", @@ -231,6 +276,11 @@ "src/aiperf/dataset/agentic_code_gen/reporting/simulation_engine.py", "simulate" ], + [ + "nesting-depth", + "src/aiperf/graph/worker_materialize.py", + "_assemble_items" + ], [ "nesting-depth", "src/aiperf/plot/config.py", @@ -305,6 +355,11 @@ "pydantic-fields", "src/aiperf/common/models/export_models.py", "JsonExportData" + ], + [ + "pydantic-fields", + "src/aiperf/timing/config.py", + "CreditPhaseConfig" ] ] } diff --git a/tools/generate_cli_docs.py b/tools/generate_cli_docs.py index 8f108497e9..a07dc3a966 100755 --- a/tools/generate_cli_docs.py +++ b/tools/generate_cli_docs.py @@ -418,10 +418,14 @@ def generate_markdown(app: Any, data: dict[str, dict[str, list[Param]]]) -> str: ] ) groups = data[name] - if len(groups) > 1 or list(groups.keys())[0] not in ( - "Parameters", - "Options", - "General", + if groups and ( + len(groups) > 1 + or list(groups.keys())[0] + not in ( + "Parameters", + "Options", + "General", + ) ): links = [ f"[{g}](#{g.lower().replace(' ', '-').replace('(', '').replace(')', '')})" diff --git a/tools/ruff_baseline.json b/tools/ruff_baseline.json index e0d93179fb..68b84d482f 100644 --- a/tools/ruff_baseline.json +++ b/tools/ruff_baseline.json @@ -244,6 +244,11 @@ "src/aiperf/common/tokenizer_validator.py", "preload_tokenizers" ], + [ + "C901", + "src/aiperf/config/dataset/resolver.py", + "DatasetResolver._resolve_one" + ], [ "C901", "src/aiperf/config/distributions.py", @@ -279,6 +284,31 @@ "src/aiperf/dataset/generator/video.py", "VideoGenerator._get_ffmpeg_install_instructions" ], + [ + "C901", + "src/aiperf/dataset/graph/adapters/dynamo/trie_lowering.py", + "dynamo_trie_nodes" + ], + [ + "C901", + "src/aiperf/dataset/graph/native_lowering.py", + "_analyze_dynamic_refs" + ], + [ + "C901", + "src/aiperf/dataset/graph/native_lowering.py", + "_assemble_prompt" + ], + [ + "C901", + "src/aiperf/dataset/graph/parser.py", + "_auto_inject_start_end" + ], + [ + "C901", + "src/aiperf/dataset/graph/parser.py", + "_expand_single_doc" + ], [ "C901", "src/aiperf/dataset/synthesis/synthesizer.py", @@ -294,6 +324,11 @@ "src/aiperf/gpu_telemetry/dcgm_collector.py", "DCGMTelemetryCollector._parse_metrics_to_records" ], + [ + "C901", + "src/aiperf/graph/worker_materialize.py", + "_assemble_items" + ], [ "C901", "src/aiperf/orchestrator/jsonl_loader.py", @@ -499,6 +534,11 @@ "src/aiperf/server_metrics/json_exporter.py", "ServerMetricsJsonExporter._build_hybrid_metrics" ], + [ + "C901", + "src/aiperf/timing/strategies/graph_ir_replay.py", + "GraphIRReplayStrategy._run_lanes" + ], [ "C901", "src/aiperf/timing/strategies/request_rate.py", diff --git a/tools/weka_trace_fidelity.py b/tools/weka_trace_fidelity.py new file mode 100644 index 0000000000..0d53f80257 --- /dev/null +++ b/tools/weka_trace_fidelity.py @@ -0,0 +1,1270 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Empirical-proof tool: validate a trie ``--export-level raw`` run vs the REAL trace. + +This is the Task-7 acceptance gate for the weka segment-trie IR rewrite. It does +NOT touch the live run path; it reads the deterministic offline artifacts a +profiling run produced (``profile_export_raw.jsonl``) plus the ORIGINAL recorded +Weka trace file, rebuilds the trie graph + segment pool from that trace exactly +as the build plane does, and proves three independent fidelity criteria: + +1. :func:`content_byte_exact_vs_v04` (criterion 1) -- per matching trace, the + ours-vs-v0.4 raw payloads are byte-identical after stripping the + ``[rid:<12hex>]`` first-user cache-bust marker. +2. :func:`content_vs_real_trace` (criterion 2) -- each exported profiling record's + reconstructed prompt ``messages`` equal what the recorded trace prescribes for + that request (the trie node's root->tip segment path, recomputed from the trace + via ``build_trie_graph`` + ``SegmentPool.materialize``). +3. :func:`causality_timing_vs_real_trace` (criterion 3 -- THE REQUIRED PROOF) -- + from the raw export ALONE, reconstruct each dispatched request's causal + predecessor (the trie ``StaticEdge`` dependency, recovered from the per-record + node id + the rebuilt trie graph) and its relative dispatch offset, and assert + (a) the causal ORDER matches the recorded trace's causal order (a request never + dispatched before its recorded predecessor) and (b) the relative inter-request + offsets match the recorded ``t`` / ``api_time`` / ``think_time`` structure + within a tolerance. + +Node-identity recovery (the linchpin): the raw export carries no ``node_id`` +field, but the worker folds the node id into ``x_request_id`` as the +4th ``|``-delimited field -- +``{conversation_id}|{instance_hash}|{trace_id}|{node_id}|{phase_variant}`` (see +``aiperf.graph.credit_dispatch_adapter._mint``). We split it back out; the +``conversation_id`` base (the part before ``#``) selects the trace. + +Run offline against any pair of artifacts; the live run that produces them is +driven by ``tests/component_integration/graph/test_weka_trace_fidelity.py``. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import orjson + +# --- report --------------------------------------------------------------- + + +@dataclass +class Mismatch: + """One failing comparison: which record/node and why.""" + + where: str + detail: str + + +@dataclass +class Report: + """Structured pass/fail result of one fidelity check. + + ``checked`` counts every comparison ATTEMPTED (passes and failures alike); + ``passes`` counts the subset that produced zero mismatches. A single checked + record can append more than one mismatch (e.g. a causal-order violation AND a + relative-timing violation), so pass counts must be read from ``passes``, never + derived as ``checked - len(mismatches)``. + + The timing criterion additionally splits its two distinct assertion kinds: + ``order_checks`` / ``order_failures`` count per-edge causal-ORDER comparisons + (a node never dispatched before a dispatched recorded predecessor), while + ``timing_checks`` / ``timing_failures`` count relative-TIMING gate comparisons + (observed dispatch vs the warped expected gate). ``exact_edges`` / + ``idle_capped_edges`` further break the passing timing gates into those whose + warped expected delay equals the raw recorded gap (``exact`` -- the export + reproduces the REAL recorded think-time exactly) and those whose recorded gap + exceeded the idle-gap cap and was warped down to it (``idle_capped`` -- + bounded by the documented faithful cap). Their sum need not equal ``checked``: + anchors and predecessor-less roots are checked but carry no inter-request + edge to classify, and a fan-in gate won by a delay-0.0 AND-join edge carries + no think-time to classify. + """ + + name: str + checked: int = 0 + passes: int = 0 + order_checks: int = 0 + order_failures: int = 0 + timing_checks: int = 0 + timing_failures: int = 0 + exact_edges: int = 0 + idle_capped_edges: int = 0 + mismatches: list[Mismatch] = field(default_factory=list) + notes: list[str] = field(default_factory=list) + + @property + def passed(self) -> bool: + return not self.mismatches and self.checked > 0 + + def fail(self, where: str, detail: str) -> None: + self.mismatches.append(Mismatch(where=where, detail=detail)) + + def render(self) -> str: + head = ( + f"{self.name}: {'PASS' if self.passed else 'FAIL'} " + f"(checked={self.checked}, passed={self.passes}, " + f"mismatches={len(self.mismatches)})" + ) + lines = [head] + if self.order_checks or self.timing_checks: + lines.append( + f" causal-order: {self.order_checks - self.order_failures}/" + f"{self.order_checks} relative-timing: " + f"{self.timing_checks - self.timing_failures}/{self.timing_checks}" + ) + if self.exact_edges or self.idle_capped_edges: + lines.append( + f" edges: exact={self.exact_edges} " + f"idle_capped={self.idle_capped_edges}" + ) + lines.extend(f" note: {n}" for n in self.notes) + for m in self.mismatches: + lines.append(f" MISMATCH @ {m.where}: {m.detail}") + return "\n".join(lines) + + +# --- raw export record ---------------------------------------------------- + +# ``[rid:<12hex>]\n\n`` per ``timing/strategies/cache_bust.py`` (FIRST_TURN_PREFIX). +_RID_MARKER_RE = re.compile(r"^\[rid:[0-9a-f]{12}\]\n\n") + +_PROFILING = "profiling" + + +@dataclass +class _ExportRecord: + """The fidelity-relevant projection of one raw-export JSONL line. + + ``start_perf_ns`` (``perf_counter`` request start) and + ``first_response_perf_ns`` (the first response's ``perf_ns``, same clock) + recover a streaming request's OBSERVED time-to-first-token as their DIFFERENCE + -- a duration valid to add to that record's wall-clock ``request_start_ns``. A + NON-streaming record still carries one ``TextResponse`` whose ``perf_ns`` is + the COMPLETION time, so its ``observed_ttft_ns`` is the FULL request duration, + NOT ``None``; the pair is ``None`` only when a response is truly absent (an + errored request that streamed nothing, or a zero-latency-style export with no + responses). That duration is CONSULTED only for post-TTFT anchor parents -- + which stream by construction -- so a non-streaming record's full-duration + value never actually anchors; the timing proof falls back to the dispatch + anchor whenever the first token is unrecoverable. + """ + + conversation_id: str + node_id: str | None + phase: str + request_start_ns: int | None + credit_issued_ns: int | None + messages: list[dict[str, str]] + start_perf_ns: int | None = None + first_response_perf_ns: int | None = None + + @property + def trace_base(self) -> str: + """The recorded trace id: the conversation_id stripped of its ``#inst`` tail.""" + return self.conversation_id.split("#", 1)[0] + + @property + def observed_ttft_ns(self) -> int | None: + """Observed request-start-to-first-token DURATION ns, or ``None``. + + ``first_response_perf_ns - start_perf_ns`` (both ``perf_counter``): a + duration valid to add to the wall-clock ``request_start_ns``. ``None`` when + either perf timestamp is missing or the difference is negative (a clock + anomaly we refuse to trust rather than silently mis-anchor). + """ + if self.start_perf_ns is None or self.first_response_perf_ns is None: + return None + dur = self.first_response_perf_ns - self.start_perf_ns + return dur if dur >= 0 else None + + +def _node_id_from_request_id(x_request_id: str | None) -> str | None: + """Recover the dispatched node id from the export's ``x_request_id``. + + Worker-minted graph shape: ``{node_id}::{uuid4().hex}`` where the node id + is the legacy-shaped ``{scope}:{turn}`` coordinate. Returns ``None`` for + any id without the ``::`` nonce separator (non-graph records). + """ + if not x_request_id or "::" not in x_request_id: + return None + return x_request_id.rsplit("::", 1)[0] + + +def _norm_messages(messages: list[dict[str, Any]]) -> list[dict[str, str]]: + """Project wire messages to the comparable ``{role, content}`` shape.""" + return [{"role": m["role"], "content": m["content"]} for m in messages] + + +def load_raw_export(raw_jsonl: Path) -> list[_ExportRecord]: + """Parse a ``profile_export_raw.jsonl`` into fidelity projections. + + Each line is one dispatched request. ``metadata`` carries + ``conversation_id`` / ``x_request_id`` / ``benchmark_phase`` / + ``request_start_ns`` / ``credit_issued_ns``; top-level ``start_perf_ns`` + + ``responses[*].perf_ns`` recover the observed first token (see + :attr:`_ExportRecord.observed_ttft_ns`); ``payload.messages`` is the wire + prompt. Blank lines are skipped. + """ + records: list[_ExportRecord] = [] + for line in Path(raw_jsonl).read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + obj = orjson.loads(line) + meta = obj.get("metadata", {}) or {} + payload = obj.get("payload", {}) or {} + responses = obj.get("responses") or [] + first_response_perf_ns = ( + responses[0].get("perf_ns") + if responses and isinstance(responses[0], dict) + else None + ) + records.append( + _ExportRecord( + conversation_id=str(meta.get("conversation_id") or ""), + node_id=_node_id_from_request_id(meta.get("x_request_id")), + phase=str(meta.get("benchmark_phase") or ""), + request_start_ns=meta.get("request_start_ns"), + credit_issued_ns=meta.get("credit_issued_ns"), + messages=_norm_messages(payload.get("messages", []) or []), + start_perf_ns=obj.get("start_perf_ns"), + first_response_perf_ns=first_response_perf_ns, + ) + ) + return records + + +# --- recorded-trace model (rebuilt trie + recorded timing) ---------------- + + +@dataclass +class _RecordedNode: + """The recorded facts for one trie node: prompt + timing + causal predecessor. + + Timing carries BOTH clocks: ``start_s`` / ``end_s`` are on the (possibly + idle-gap-warped) clock the rebuilt trie placed the node on -- the SAME clock + the capped RUN dispatched against -- while ``raw_start_s`` / ``raw_end_s`` are + the unwarped recorded ``request.t`` / ``t + api_time``. ``pred_delay_us`` maps + each recorded predecessor (trie ``StaticEdge`` source) to that edge's warped + ``delay_after_predecessor_us`` -- the EXPECTED end-to-start delay the proof + compares against. Comparing the warped expected against the raw recorded gap + classifies each edge as exact (gap <= cap) or idle-capped (gap > cap). + """ + + node_id: str + messages: list[dict[str, str]] + start_s: float # warped request start (== raw request.t when uncapped) + end_s: float # warped completion (warped start + raw api_time) + raw_start_s: float # unwarped recorded request.t + raw_end_s: float # unwarped recorded t + api_time + predecessors: list[str] # trie StaticEdge sources (excluding START) + pred_delay_us: dict[str, float] # warped StaticEdge delay per predecessor + rooted_at_start: bool # True iff its only recorded edge is from START + # The ``StaticEdge(source="START").min_start_delay_us`` the trie builder + # stamped on this node -- its ABSOLUTE warped arrival offset from the instance + # run-origin. The executor fires a START-rooted node at + # ``anchor_wall + min_start_delay_us`` (``_compute_firing_gate_us`` with + # ``absolute_start_offsets=True``), so the proof anchors START-roots to this + # offset rather than to the anchor's recorded END. ``None`` for a node with a + # non-START predecessor edge. + min_start_delay_us: float | None = None + # (parent node id, warped start-to-start delay us D, first-token delay us D' | + # None) when this node's ONLY recorded edge is start-anchored (mid-flight spawn + # / chain overlap). The timing proof compares OBSERVED child_dispatch against + # the parent's dispatch + D, and the order check requires parent dispatch + # BEFORE child. When D' is not None the node was recorded POST-TTFT: the proof + # prefers ``parent_first_token + D'`` whenever the parent's observed first token + # is recoverable from the export, falling back to dispatch + D (LOUDLY) when it + # is not. D' is None for a pre-TTFT child or a non-streaming parent. + start_anchor: tuple[str, float, float | None] | None = None + + +@dataclass +class _RecordedTrace: + """Every trie node of one recorded trace, keyed by node id.""" + + trace_id: str + nodes: dict[str, _RecordedNode] + + +def build_recorded_trace( + trace_file: Path, + idle_gap_cap_seconds: float | None = 60.0, + *, + tokenizer_name: str = "builtin", + prompt_corpus: str = "coding", + root_seed: int | None = None, +) -> _RecordedTrace: + """Rebuild the trie graph + segment pool from a recorded Weka trace file. + + Reproduces the build-plane realization (``build_trie_graph``) so each node's + EXPECTED prompt is ``pool.materialize(prompt_segment_ids)`` and its recorded + causal predecessors are the trie ``StaticEdge`` sources into it. + + ``tokenizer_name`` / ``prompt_corpus`` / ``root_seed`` are the content knobs + the run under proof was built with; the defaults mirror a bare live run + (builtin tokenizer, ``"coding"`` corpus, no ``--random-seed``). A run built + with different knobs (e.g. ``--tokenizer gpt2`` or an explicit seed) must + pass the same values here or every content comparison spuriously fails. + + ``idle_gap_cap_seconds`` is passed straight through to ``build_trie_graph`` so + the rebuilt trie's ``StaticEdge.delay_after_predecessor_us`` and the node + arrival offsets sit on the SAME idle-gap-warped clock a capped RUN dispatched + against (60.0s default == the ``inferencex-agentx-mvp`` scenario cap). The + expected end-to-start delay the timing proof compares is that WARPED edge + delay, captured per predecessor in ``_RecordedNode.pred_delay_us``. The RAW + recorded ``t`` / ``api_time`` is kept alongside (``raw_start_s`` / + ``raw_end_s``) so the proof can classify each edge exact-vs-idle-capped and + report it transparently. Pass ``None`` to disable the warp (raw timeline). + + Imports are local so the (heavy) aiperf graph stack stays off the import path + for callers that only need the pure :class:`Report` helpers. + """ + from aiperf.dataset.graph.adapters.weka.trace_models import WekaTrace + from aiperf.dataset.graph.adapters.weka.trie_build import ( + _flatten_requests, + build_trie_graph, + ) + from aiperf.dataset.graph.models import LlmNode, StaticEdge + + raw = orjson.loads(Path(trace_file).read_bytes()) + trace = WekaTrace.model_validate(raw) + parsed, pool = build_trie_graph( + trace, + tokenizer_name=tokenizer_name, + prompt_corpus=prompt_corpus, + root_seed=root_seed, + idle_gap_cap_seconds=idle_gap_cap_seconds, + ) + graph = parsed.graph + + # RAW (unwarped) per-node timing straight from the recorded request.t / + # api_time. ``_flatten_requests`` here does NOT apply the warp (warped_start + # stays at its 0.0 default), so we read the raw request fields directly rather + # than the warped ``.start`` / ``.end`` properties. + flat = _flatten_requests(trace.requests, root_scope=trace.id) + raw_timing: dict[str, tuple[float, float]] = { + n.node_id: (n.request.t, n.request.t + (n.request.api_time or 0.0)) + for n in flat + } + + # WARPED per-node start from the rebuilt trie's arrival offset (the clock the + # capped run used). Warped end = warped start + RAW api_time (api_time is the + # node's own processing, not an inter-request idle gap, so it is not warped -- + # mirrors ``_Node.end``). + raw_api_time: dict[str, float] = { + n.node_id: (n.request.api_time or 0.0) for n in flat + } + + # Causal predecessors + the WARPED edge delay per predecessor edge. A node + # whose only edge is from START is a recorded root (its recorded arrival is + # its own warped t). + preds: dict[str, list[str]] = {nid: [] for nid in graph.nodes} + pred_delay_us: dict[str, dict[str, float]] = {nid: {} for nid in graph.nodes} + start_rooted: dict[str, bool] = {nid: False for nid in graph.nodes} + start_min_delay_us: dict[str, float] = {} + start_anchor: dict[str, tuple[str, float, float | None]] = {} + for edge in graph.edges: + if not isinstance(edge, StaticEdge): + continue + tgt = edge.target + if tgt not in preds: + continue + if edge.source == "START": + start_rooted[tgt] = True + start_min_delay_us[tgt] = edge.min_start_delay_us or 0.0 + elif edge.delay_after_predecessor_start_us is not None: + # Start-anchored: gates off the predecessor's DISPATCH (dispatch-to- + # dispatch), NOT its completion, so it does NOT model an end-to-start + # wait. Keep it out of preds / pred_delay_us (those drive the + # end-to-start gate) and record the (parent, D, D' | None) separately. + start_anchor[tgt] = ( + edge.source, + edge.delay_after_predecessor_start_us, + edge.delay_after_predecessor_first_token_us, + ) + else: + preds[tgt].append(edge.source) + pred_delay_us[tgt][edge.source] = edge.delay_after_predecessor_us or 0.0 + + nodes: dict[str, _RecordedNode] = {} + for nid, node in graph.nodes.items(): + if not isinstance(node, LlmNode): + continue + path = node.metadata["trie"]["prompt_segment_ids"] + warped_start_s = (node.arrival_offset_us or 0) / 1e6 + warped_end_s = warped_start_s + raw_api_time.get(nid, 0.0) + raw_start_s, raw_end_s = raw_timing.get(nid, (0.0, 0.0)) + nodes[nid] = _RecordedNode( + node_id=nid, + messages=pool.materialize(path), + start_s=warped_start_s, + end_s=warped_end_s, + raw_start_s=raw_start_s, + raw_end_s=raw_end_s, + predecessors=preds.get(nid, []), + pred_delay_us=pred_delay_us.get(nid, {}), + rooted_at_start=start_rooted.get(nid, False) and not preds.get(nid), + min_start_delay_us=( + start_min_delay_us.get(nid) + if start_rooted.get(nid, False) and not preds.get(nid) + else None + ), + start_anchor=start_anchor.get(nid), + ) + return _RecordedTrace(trace_id=trace.id, nodes=nodes) + + +# --- criterion 2: content vs real trace ----------------------------------- + + +def content_vs_real_trace( + raw_jsonl: Path | None, + trace_file: Path | None, + *, + recorded: _RecordedTrace | None = None, + records: list[_ExportRecord] | None = None, + tokenizer_name: str = "builtin", + prompt_corpus: str = "coding", + root_seed: int | None = None, +) -> Report: + """Each profiling export record's prompt == the recorded trace's prescribed content. + + Maps every profiling record to its recorded trie node (by the node id folded + into ``x_request_id``) and asserts the exported ``payload.messages`` equal + ``pool.materialize(node.prompt_segment_ids)`` recomputed from the trace. + + The exported first-user message may carry a ``[rid:...]`` cache-bust prefix + (when the run used ``--cache-bust first_turn_prefix``); it is stripped before + comparison so the check proves CONTENT fidelity independent of the bust marker. + + The corpus driver supplies a pre-built ``recorded`` / ``records`` so each + recorded trace is built once and both criteria run against it; the single-file + entrypoint builds them from ``trace_file`` / ``raw_jsonl``. Content fidelity is + cap-INDEPENDENT (prompt content doesn't depend on timing). + """ + report = Report(name="content_vs_real_trace") + if recorded is None: + assert trace_file is not None + recorded = build_recorded_trace( + trace_file, + tokenizer_name=tokenizer_name, + prompt_corpus=prompt_corpus, + root_seed=root_seed, + ) + if records is None: + assert raw_jsonl is not None + records = load_raw_export(raw_jsonl) + + profiling = [ + r + for r in records + if r.phase == _PROFILING and r.trace_base == recorded.trace_id + ] + if not profiling: + report.fail("export", "no profiling records found in raw export") + return report + + for i, rec in enumerate(profiling): + where = f"profiling[{i}] node={rec.node_id} conv={rec.conversation_id}" + report.checked += 1 + node = recorded.nodes.get(rec.node_id or "") + if node is None: + report.fail(where, f"node id {rec.node_id!r} not in rebuilt trie graph") + continue + got = _strip_rid_marker(rec.messages) + if got != node.messages: + report.fail(where, _first_message_diff(node.messages, got)) + continue + report.passes += 1 + return report + + +def _strip_rid_marker(messages: list[dict[str, str]]) -> list[dict[str, str]]: + """Return a copy with the ``[rid:...]`` prefix stripped from the first user msg.""" + out = [dict(m) for m in messages] + for m in out: + if m.get("role") == "user": + m["content"] = _RID_MARKER_RE.sub("", m["content"], count=1) + break + return out + + +def _first_message_diff( + expected: list[dict[str, str]], got: list[dict[str, str]] +) -> str: + """A compact description of the first differing message (for failure detail).""" + if len(expected) != len(got): + return f"message count {len(got)} != expected {len(expected)}" + for idx, (e, g) in enumerate(zip(expected, got, strict=False)): + if e != g: + return ( + f"msg[{idx}] role {g.get('role')!r}/{e.get('role')!r}; " + f"content {g.get('content', '')[:60]!r} != " + f"{e.get('content', '')[:60]!r}" + ) + return "messages differ (no positional diff found)" + + +# --- criterion 3: causality + timing vs real trace ------------------------ + +# Timing tolerance for the recorded-relative-offset comparison. Wall-clock is +# absolute-different (a fresh run-origin), but the RECORDED RELATIVE gap between a +# node and its causal predecessor must survive the replay. We allow the larger of +# an absolute floor (scheduling jitter, ZMQ transit, the gap-warp rounding) and a +# relative fraction of the recorded gap (longer recorded gaps accrue +# proportionally more replay slack). +_TIMING_ABS_TOLERANCE_S = 0.75 +_TIMING_REL_TOLERANCE = 0.15 + +# Sub-second slack distinguishing a warped (idle-capped) edge from an exact one: +# float us->s round-trip on the warped delay introduces tiny noise, so treat a +# warped expected within this of the raw recorded gap as "exact" (gap <= cap). +_CLASSIFY_EPSILON_S = 1e-3 + + +def causality_timing_vs_real_trace( + raw_jsonl: Path | None, + trace_file: Path | None, + *, + recorded: _RecordedTrace | None = None, + records: list[_ExportRecord] | None = None, + idle_gap_cap_seconds: float | None = 60.0, + tokenizer_name: str = "builtin", + prompt_corpus: str = "coding", + root_seed: int | None = None, +) -> Report: + """Reconstruct dispatch causality + relative timing from the export ALONE. + + For every profiling record we recover its node id (from ``x_request_id``) + and look up the recorded trie causal predecessors (``StaticEdge`` sources). We + then assert, against the OBSERVED ``request_start_ns`` of each dispatched + record (warmup + profiling, since a profiling node's predecessor may have been + chopped into warmup by t*): + + (a) Causal ORDER: for each recorded predecessor that was ALSO dispatched, the + record's observed ``request_start_ns`` is >= the predecessor's observed + ``request_start_ns`` -- a request never dispatched before its recorded + predecessor (the closed-loop waits-for edge is honored on the wire). + + (b) Relative TIMING: the WARPED end-to-start gap between the node and its + (latest-completing) recorded predecessor -- the rebuilt trie's + ``StaticEdge.delay_after_predecessor_us`` on the idle-gap-CAPPED clock, + captured in ``_RecordedNode.pred_delay_us`` -- matches the OBSERVED + start-to-start gap within :data:`_TIMING_ABS_TOLERANCE_S` (abs) or + :data:`_TIMING_REL_TOLERANCE` (relative to the warped gap). + + Cap-awareness (principled, NOT a fudge): a faithful replay applies the + agentx idle-gap cap, so for any recorded end-to-start gap > cap the RUN + dispatched after only ``cap`` seconds, not the raw recorded think-time. + The EXPECTED delay must therefore be the WARPED edge delay, not the raw + recorded gap. Each checked edge is CLASSIFIED by comparing the warped + expected against the raw recorded gap: + + * "exact" when warped == raw (gap <= cap) -- the export reproduces the + REAL recorded think-time exactly. + * "idle-capped" when warped < raw (gap > cap) -- bounded by the documented + faithful cap (warped expected == cap). + + Tolerances are NOT weakened and no edge is skipped to force a pass. + + Why END-to-start, not start-to-start: the executor parks each node for its + ``delay_after_predecessor_us`` AFTER the predecessor RETURNS, then + dispatches. The proof run drives the in-repo mock at ``--ttft 0 --itl 0``, + so the predecessor returns ~instantly after IT dispatched; the observed + dispatch-clock gap therefore collapses onto the (warped) end-to-start edge + delay (the recorded ``api_time`` is the predecessor's own processing, which + the zero-latency mock does not reproduce). + + For a recorded node with NO dispatched predecessor (a t*-chop survivor + re-rooted at START, or a true root), the reference gap is measured + (warped) start-to-start from the trace's earliest dispatched node (the run + anchor) and compared to the observed gap from that anchor's dispatch. + + START-ANCHORED node (a mid-flight spawn / chain overlap whose only edge is + ``StaticEdge.delay_after_predecessor_start_us`` D): the runtime schedules it + at its parent's DISPATCH and gates it at dispatch + D, so the reference is + DISPATCH-to-dispatch -- observed ``request_start_ns(child) - + request_start_ns(parent)`` must equal the warped start-anchor delay (same + tolerance), and the parent must dispatch before the child. These edges are + counted in ``exact_edges`` (they reproduce the recorded start-to-start gap + exactly; the idle-gap cap warps the delay itself, not the anchor kind). + + POST-TTFT start-anchored node (its edge also carries + ``delay_after_predecessor_first_token_us`` D'): the runtime re-anchors it + onto the parent's OBSERVED first token, so the proof compares the child + against ``parent_first_token + D'`` whenever that first token is recoverable + from the export (``responses[0].perf_ns - start_perf_ns`` added to the + parent's dispatch wall clock; post-TTFT anchor parents STREAM by + construction). The runtime falls back to dispatch + D -- and the proof does + too, emitting a LOUD ``FALLBACK`` note (the edge is still CHECKED, never + silently skipped) -- only when that first token is truly unrecoverable: an + errored parent that streamed nothing, or a zero-latency export with no + responses. A non-streaming record is not a fallback cause here -- its lone + ``TextResponse`` still yields a duration, and it never sources a post-TTFT + anchor edge anyway. + + Absolute wall-clock differs run-to-run; only the recorded RELATIVE timing + + causal order are asserted. The corpus driver supplies a pre-built ``recorded`` + / ``records`` (built ONCE with the cap); the single-file entrypoint builds + them from ``trace_file`` / ``raw_jsonl`` with ``idle_gap_cap_seconds``. + """ + report = Report(name="causality_timing_vs_real_trace") + if recorded is None: + assert trace_file is not None + recorded = build_recorded_trace( + trace_file, + idle_gap_cap_seconds=idle_gap_cap_seconds, + tokenizer_name=tokenizer_name, + prompt_corpus=prompt_corpus, + root_seed=root_seed, + ) + if records is None: + assert raw_jsonl is not None + records = load_raw_export(raw_jsonl) + + def _earliest_by_node(phase: str | None) -> dict[str, int]: + """Earliest observed ``request_start_ns`` per node id, optionally one phase.""" + out: dict[str, int] = {} + for rec in records: + if rec.trace_base != recorded.trace_id or rec.node_id is None: + continue + if rec.request_start_ns is None: + continue + if phase is not None and rec.phase != phase: + continue + prev = out.get(rec.node_id) + if prev is None or rec.request_start_ns < prev: + out[rec.node_id] = rec.request_start_ns + return out + + # CAUSAL-ORDER uses the cross-phase earliest dispatch (a profiling node's + # recorded predecessor may have been chopped into warmup by t*). RELATIVE + # TIMING must stay phase-consistent: a profiling node's recorded predecessor, + # when itself dispatched in profiling, replays the SAME re-rooted t*-relative + # schedule, so we compare profiling-vs-profiling timestamps (warmup primes a + # node ~immediately and would alias the edge delay). + any_phase_ns = _earliest_by_node(None) + profiling_ns = _earliest_by_node(_PROFILING) + + def _first_token_duration_by_node() -> dict[str, int]: + """Observed first-token DURATION ns for each node's EARLIEST profiling record. + + Pairs the same earliest-dispatch record ``profiling_ns`` picks, so a node's + recovered first token is ``profiling_ns[node] + this[node]`` on the wall + clock. Absent only for a record with NO responses (an errored request that + streamed nothing, or a zero-latency export); a non-streaming record's lone + ``TextResponse`` still yields its full duration here, but this map is + consulted only for post-TTFT anchor parents (which stream by construction), + so that value never anchors -- when absent the timing proof falls back to + dispatch + D. + """ + best_start: dict[str, int] = {} + out: dict[str, int] = {} + for rec in records: + if rec.trace_base != recorded.trace_id or rec.node_id is None: + continue + if rec.phase != _PROFILING or rec.request_start_ns is None: + continue + prev = best_start.get(rec.node_id) + if prev is not None and rec.request_start_ns >= prev: + continue + best_start[rec.node_id] = rec.request_start_ns + dur = rec.observed_ttft_ns + if dur is None: + out.pop(rec.node_id, None) + else: + out[rec.node_id] = dur + return out + + profiling_ft_ns = _first_token_duration_by_node() + + if not any_phase_ns: + report.fail("export", "no dispatched trie nodes recovered from raw export") + return report + + # Instance run-origin for START-rooted timing. The executor fires a + # START-rooted node at ``anchor_wall + min_start_delay_us`` + # (``_compute_firing_gate_us``, ``absolute_start_offsets=True``), where + # ``anchor_wall`` is the SHARED instance run-start pinned once on the top + # executor's run. We recover that origin from the profiled START-root with the + # SMALLEST ``min_start_delay_us`` (typically the recorded root, offset 0): + # ``origin_ns = obs_ns(root) - root.min_start_delay_us``. Every other + # START-root must then dispatch at ``origin_ns + its own min_start_delay_us``, + # NOT relative to the root's END (independent roots share one t* origin and do + # NOT chain off the root's processing time). + if not profiling_ns: + report.fail("export", "no profiling records for this trace in raw export") + return report + start_roots = [ + nid + for nid in profiling_ns + if nid in recorded.nodes and recorded.nodes[nid].min_start_delay_us is not None + ] + if start_roots: + # Prefer the profiled START-root with the SMALLEST min_start_delay (the + # recorded root, offset 0) to recover the shared instance run-origin. + anchor_id = min( + start_roots, + key=lambda nid: recorded.nodes[nid].min_start_delay_us or 0.0, + ) + origin_ns: float | None = ( + profiling_ns[anchor_id] + - (recorded.nodes[anchor_id].min_start_delay_us or 0.0) * 1e3 + ) + # Independent START-roots must all imply the SAME origin within tolerance; + # a disagreement means a root fired off-schedule (a real fidelity signal), + # so surface it as a note (each root is also asserted per-node below). + for nid in start_roots: + implied = ( + profiling_ns[nid] + - (recorded.nodes[nid].min_start_delay_us or 0.0) * 1e3 + ) + drift_s = abs(implied - origin_ns) / 1e9 + if drift_s > _TIMING_ABS_TOLERANCE_S: + report.notes.append( + f"START-root {nid} implies origin drift {drift_s:.3f}s vs " + f"anchor {anchor_id} (checked per-node below)" + ) + else: + # No profiled START-root (the profiled subset is a mid-chain slice whose + # roots were never dispatched in profiling): there is no run-origin to + # anchor START-roots against. Fall back to the earliest-RECORDED profiled + # node as a bare relative-timing anchor that carries NO constraint (it is + # the origin of the comparable slice). START-root expecteds are then + # unreachable, so every profiled node is timed off its profiled preds. + anchor_id = min( + profiling_ns, + key=lambda nid: ( + recorded.nodes[nid].start_s if nid in recorded.nodes else float("inf") + ), + ) + origin_ns = None + report.notes.append( + f"trace={recorded.trace_id} dispatched_nodes={sorted(any_phase_ns)} " + f"profiled={sorted(profiling_ns)} anchor={anchor_id}" + ) + + profiling = [ + r + for r in records + if r.phase == _PROFILING + and r.node_id is not None + and r.trace_base == recorded.trace_id + ] + + def _check_record(rec: _ExportRecord, where: str) -> None: + """Run every order/timing assertion for one profiling record. + + Appends mismatches (possibly several -- causal-order and relative-timing + are DISTINCT assertion kinds, tallied separately on the report) and + returns; the caller owns the record-level ``checked``/``passes`` tally. + """ + nid = rec.node_id + node = recorded.nodes.get(nid or "") + if node is None: + report.fail(where, f"node id {nid!r} not in rebuilt trie graph") + return + if rec.request_start_ns is None: + report.fail(where, "record has no request_start_ns") + return + obs_s = rec.request_start_ns / 1e9 + + # Start-anchored node: its ONE recorded edge gates off the parent's + # DISPATCH (dispatch-to-dispatch), not its completion. Compare the + # OBSERVED start-to-start gap to the warped start-anchor delay, and require + # the parent to have dispatched before this child. (Start-anchored edges + # are kept out of predecessors/pred_delay_us, so this branch is the only + # place they are timed.) + if node.start_anchor is not None: + parent, delay_us, first_token_delay_us = node.start_anchor + # (a) Causal order: never dispatched before the recorded parent (same + # tolerance the end-to-start order check uses). + parent_any_ns = any_phase_ns.get(parent) + if parent_any_ns is not None: + report.order_checks += 1 + if obs_s + _TIMING_ABS_TOLERANCE_S < parent_any_ns / 1e9: + report.order_failures += 1 + report.fail( + where, + f"dispatched {parent_any_ns / 1e9 - obs_s:.3f}s BEFORE " + f"recorded start-anchor parent {parent} " + f"(causal-order violation)", + ) + # (b) Relative timing. If the parent was never dispatched in profiling + # there is no gate to compare against; the order check above still + # guards it. + parent_prof_ns = profiling_ns.get(parent) + if parent_prof_ns is None: + return + # POST-TTFT node (D' present): the runtime re-anchors it onto the + # parent's OBSERVED first token, so the proof must too -- expected == + # parent_first_token + D' whenever that first token is recoverable from + # the export. Post-TTFT anchor parents STREAM by construction; the first + # token is unrecoverable only when the parent has NO responses (an + # errored parent that streamed nothing, or a zero-latency export), in + # which case the runtime falls back to dispatch + D and so does the + # proof, LOUDLY: the edge is still CHECKED, never silently skipped. A + # pre-TTFT node (D' is None) always gates dispatch-to-dispatch on D. + parent_ft_dur_ns = ( + profiling_ft_ns.get(parent) + if first_token_delay_us is not None + else None + ) + if parent_ft_dur_ns is not None: + expected_gap_s = first_token_delay_us / 1e6 + expected_s = (parent_prof_ns + parent_ft_dur_ns) / 1e9 + expected_gap_s + ref = f"first_token({parent})+D'" + else: + expected_gap_s = delay_us / 1e6 + expected_s = parent_prof_ns / 1e9 + expected_gap_s + if first_token_delay_us is not None: + report.notes.append( + f"FALLBACK first-token edge {parent}->{nid}: parent observed " + f"TTFT unrecoverable from raw export; compared child dispatch " + f"vs parent dispatch + D ({expected_gap_s:.3f}s)" + ) + ref = f"dispatch({parent})+D [first-token fallback]" + else: + ref = f"dispatch({parent})+D" + report.timing_checks += 1 + tol = max( + _TIMING_ABS_TOLERANCE_S, _TIMING_REL_TOLERANCE * abs(expected_gap_s) + ) + if abs(obs_s - expected_s) > tol: + report.timing_failures += 1 + report.fail( + where, + f"start-anchor offset vs {ref}: observed dispatch " + f"{obs_s - expected_s:+.3f}s off expected delay " + f"{expected_gap_s:.3f}s (tol {tol:.3f}s)", + ) + return + report.exact_edges += 1 + return + + # Causal order considers any-phase dispatch; relative timing stays within + # the profiling phase (phase-consistent re-rooted schedule). + order_preds = [p for p in node.predecessors if p in any_phase_ns] + timing_preds = [p for p in node.predecessors if p in profiling_ns] + + # (a) Causal order: never dispatched before any dispatched recorded pred. + for pred in order_preds: + report.order_checks += 1 + pred_obs_s = any_phase_ns[pred] / 1e9 + if obs_s + _TIMING_ABS_TOLERANCE_S < pred_obs_s: + report.order_failures += 1 + report.fail( + where, + f"dispatched {pred_obs_s - obs_s:.3f}s BEFORE recorded " + f"predecessor {pred} (causal-order violation)", + ) + + # (b) Relative timing: compute the EXPECTED observed dispatch the way the + # executor computes its firing gate -- the AND-fan-in MAX over incoming + # edges -- using OBSERVED predecessor dispatch times: + # + # expected_ns = max over profiled preds p of + # ( profiling_ns[p] + warped_delay_us[p] * 1e3 ) + # + # and for a START-rooted node (no profiled predecessor): + # + # expected_ns = origin_ns + min_start_delay_us * 1e3 + # + # The BINDING predecessor is the argmax (the gate the executor actually + # waited on), which may NOT be the recorded latest-completing pred when a + # different edge's warped delay (e.g. a 60s-capped one) pushes its gate + # later. ``expected_gap_s`` (the binding edge's warped delay, or the + # START-root offset) sets the relative-tolerance magnitude exactly as + # before. ``raw_gap_s`` is the unwarped recorded gap of the binding edge, + # kept only to CLASSIFY exact-vs-idle-capped (NOT the comparison target). + classify = False + if timing_preds: + gates_ns = { + p: profiling_ns[p] + node.pred_delay_us.get(p, 0.0) * 1e3 + for p in timing_preds + } + cause = max(gates_ns, key=lambda p: gates_ns[p]) + expected_ns = gates_ns[cause] + expected_gap_s = node.pred_delay_us.get(cause, 0.0) / 1e6 + raw_gap_s = node.raw_start_s - recorded.nodes[cause].raw_end_s + ref = f"pred={cause}" + # Classify exact-vs-idle-capped only when the argmax edge carries a + # nonzero warped delay: a NON-BINDING AND-join edge has delay 0.0 by + # construction (build_interval_edges), so when it wins the argmax it + # carries no think-time to compare -- its positive raw end-to-start + # gap would misread as "idle-capped" although no gap exceeded the + # cap. (A capped binding edge always has warped delay == cap > 0, so + # no idle-capped edge is ever skipped by this guard.) + classify = expected_gap_s > 0.0 + elif node.min_start_delay_us is not None and origin_ns is not None: + expected_ns = origin_ns + node.min_start_delay_us * 1e3 + expected_gap_s = node.min_start_delay_us / 1e6 + raw_gap_s = expected_gap_s + ref = f"origin={anchor_id}" + else: + # A node with recorded non-START preds but NONE dispatched in profiling + # (its preds were chopped into warmup by t*): it has no profiling-phase + # gate to compare against. The causal-order check above still guards it; + # it stays checked without a relative-timing assertion. + return + + report.timing_checks += 1 + expected_s = expected_ns / 1e9 + tol = max(_TIMING_ABS_TOLERANCE_S, _TIMING_REL_TOLERANCE * abs(expected_gap_s)) + if abs(obs_s - expected_s) > tol: + report.timing_failures += 1 + report.fail( + where, + f"relative offset vs {ref}: observed dispatch " + f"{obs_s - expected_s:+.3f}s off warped-expected (binding warped " + f"delay {expected_gap_s:.3f}s, raw {raw_gap_s:.3f}s, tol {tol:.3f}s)", + ) + return + if classify: + # warped < raw => the recorded gap exceeded the cap and was compressed + # to it (idle-capped); warped == raw => sub-cap gap reproduced exactly. + if expected_gap_s + _CLASSIFY_EPSILON_S < raw_gap_s: + report.idle_capped_edges += 1 + else: + report.exact_edges += 1 + + for i, rec in enumerate(profiling): + report.checked += 1 + mismatches_before = len(report.mismatches) + _check_record(rec, f"profiling[{i}] node={rec.node_id}") + if len(report.mismatches) == mismatches_before: + report.passes += 1 + return report + + +# --- criterion 1: byte-exact vs v0.4 -------------------------------------- + + +def content_byte_exact_vs_v04(ours_raw: Path, v04_raw: Path) -> Report: + """Per matching trace, ours-vs-v0.4 raw payloads are byte-identical (rid-stripped). + + Matches records across the two raw exports by ``(conversation_id base, + node_id)`` -- the stable per-trace-per-node identity -- and asserts the + rid-stripped ``payload.messages`` are byte-identical. Only profiling records + are compared (warmup dispatch counts can legitimately differ between the two + builds). A trace/node present in one export but not the other is reported as a + coverage mismatch, so a vacuous (empty-overlap) pass cannot slip through. + """ + report = Report(name="content_byte_exact_vs_v04") + ours = _index_by_trace_node(load_raw_export(ours_raw)) + v04 = _index_by_trace_node(load_raw_export(v04_raw)) + + common = sorted(set(ours) & set(v04)) + only_ours = sorted(set(ours) - set(v04)) + only_v04 = sorted(set(v04) - set(ours)) + for key in only_ours: + report.fail(f"{key}", "present in OURS export but missing from v0.4 export") + for key in only_v04: + report.fail(f"{key}", "present in v0.4 export but missing from OURS export") + + for key in common: + report.checked += 1 + ours_msgs = _strip_rid_marker(ours[key]) + v04_msgs = _strip_rid_marker(v04[key]) + if ours_msgs != v04_msgs: + report.fail(str(key), _first_message_diff(v04_msgs, ours_msgs)) + continue + report.passes += 1 + return report + + +def _index_by_trace_node( + records: list[_ExportRecord], +) -> dict[tuple[str, str | None], list[dict[str, str]]]: + """Index profiling records by ``(trace_base, node_id)`` -> messages. + + Last write wins on a duplicate key (a recycled node re-fire is byte-identical + on content, so the choice is immaterial for the byte-exact comparison). + """ + out: dict[tuple[str, str | None], list[dict[str, str]]] = {} + for rec in records: + if rec.phase != _PROFILING: + continue + out[(rec.trace_base, rec.node_id)] = rec.messages + return out + + +# --- corpus-scale driver -------------------------------------------------- + + +@dataclass +class _TraceResult: + """Per-trace fidelity outcome the corpus driver aggregates + prints. + + ``dispatched_records`` counts the raw-export PROFILING records mapped to + this trace -- the records both criteria actually assert on. Zero means the + bounded run never profiled the trace (either never dispatched it at all, or + only warmup-primed it; ``warmup_records`` distinguishes the two) and is + reported as coverage (``SKIP``), never as a failure -- while a nonzero count + with a failing report is a REAL fidelity failure the exit code must surface. + """ + + trace_id: str + content: Report + timing: Report + dispatched_nodes: int + total_nodes: int + dispatched_records: int = 0 + warmup_records: int = 0 + + @property + def coverage(self) -> float: + """Fraction of recorded trie nodes this run actually profiled. + + A bounded run (``--request-count`` / a short window) may never reach the + deepest recorded turns, and graph auto-warmup may prime a node without + ever profiling it; those nodes are reported COVERAGE, not failures (the + proof only asserts the profiled subset). + """ + return self.dispatched_nodes / self.total_nodes if self.total_nodes else 0.0 + + @property + def passed(self) -> bool: + return self.content.passed and self.timing.passed + + +def prove_corpus( + raw_jsonl: Path, + trace_dir: Path, + idle_gap_cap_seconds: float | None = 60.0, + *, + tokenizer_name: str = "builtin", + prompt_corpus: str = "coding", + root_seed: int | None = None, +) -> int: + """Prove a whole corpus: one raw export vs every ``*.json`` trace in a dir. + + The raw export interleaves ALL traces; we map each record to its trace by + ``_ExportRecord.trace_base == recorded.trace_id`` and run BOTH + :func:`content_vs_real_trace` and :func:`causality_timing_vs_real_trace` + restricted to that trace's records. Each recorded trace is built ONCE (with + the cap + content knobs -- see :func:`build_recorded_trace`) and both + criteria share it. + + Prints a per-trace + TOTAL summary (records dispatched, content/causal-order/ + relative-timing pass counts, exact-vs-idle-capped edge counts, node coverage). + + Exit contract: returns nonzero iff ANY trace with PROFILING records fails + either criterion (including unresolvable node ids and records missing + ``request_start_ns`` -- a printed MISMATCH is never a passing exit), OR the + proof was VACUOUS (no trace files, or no export record checked against any + trace): a proof that checked nothing must not pass. Traces without profiling + records -- never dispatched at all, or only warmup-primed (graph auto-warmup + bursts priming credits before profiling) -- are reported as coverage + (``SKIP``), NOT failures; an ALL-warmup export still fails loudly through + the VACUOUS gate. + """ + records = load_raw_export(raw_jsonl) + by_trace: dict[str, list[_ExportRecord]] = {} + for rec in records: + by_trace.setdefault(rec.trace_base, []).append(rec) + + trace_files = sorted(Path(trace_dir).glob("*.json")) + results: list[_TraceResult] = [] + for trace_file in trace_files: + recorded = build_recorded_trace( + trace_file, + idle_gap_cap_seconds=idle_gap_cap_seconds, + tokenizer_name=tokenizer_name, + prompt_corpus=prompt_corpus, + root_seed=root_seed, + ) + trace_records = by_trace.get(recorded.trace_id, []) + profiling_records = [r for r in trace_records if r.phase == _PROFILING] + if not profiling_records: + # No PROFILING records for this trace: either the bounded run never + # dispatched it at all, or auto-warmup only primed it. Both criteria + # hard-fail on zero profiling records, so running them here would + # mislabel coverage as a fidelity failure -- report SKIP instead + # (the corpus-level VACUOUS gate still fails an all-warmup export). + results.append( + _TraceResult( + trace_id=recorded.trace_id, + content=Report(name="content_vs_real_trace"), + timing=Report(name="causality_timing_vs_real_trace"), + dispatched_nodes=0, + total_nodes=len(recorded.nodes), + dispatched_records=0, + warmup_records=len(trace_records), + ) + ) + continue + content = content_vs_real_trace( + None, None, recorded=recorded, records=trace_records + ) + timing = causality_timing_vs_real_trace( + None, None, recorded=recorded, records=trace_records + ) + dispatched = { + r.node_id + for r in profiling_records + if r.node_id is not None and r.node_id in recorded.nodes + } + results.append( + _TraceResult( + trace_id=recorded.trace_id, + content=content, + timing=timing, + dispatched_nodes=len(dispatched), + total_nodes=len(recorded.nodes), + dispatched_records=len(profiling_records), + warmup_records=len(trace_records) - len(profiling_records), + ) + ) + + _print_corpus_summary(results, idle_gap_cap_seconds) + + total_checked = sum(r.content.checked + r.timing.checked for r in results) + if not results or total_checked == 0: + print( + "VACUOUS: nothing checked (no recorded traces, or no raw-export " + "record maps to any trace) -- a proof that checked nothing FAILS" + ) + return 1 + failed = [r.trace_id for r in results if r.dispatched_records > 0 and not r.passed] + if failed: + print(f"corpus proof: FAIL ({len(failed)} trace(s): {', '.join(failed)})") + return 1 + print("corpus proof: PASS") + return 0 + + +def _print_corpus_summary( + results: list[_TraceResult], idle_gap_cap_seconds: float | None +) -> None: + """Print the per-trace + TOTAL corpus fidelity table. + + ``content`` counts profiling records whose prompt matched; ``order`` counts + per-edge causal-order comparisons; ``timing`` counts relative-timing gate + comparisons -- two DIFFERENT assertion kinds, so they get separate columns + rather than one mixed pass/checked pair. All pass counts come from the + reports' dedicated counters (never derived by subtracting mismatch counts, + which can exceed one per checked record). + """ + cap = "off" if idle_gap_cap_seconds is None else f"{idle_gap_cap_seconds:g}s" + print(f"weka trace-fidelity corpus proof (idle-gap cap = {cap})") + print( + f"{'trace':<28} {'disp/tot':>9} {'cov':>5} " + f"{'content':>14} {'order':>12} {'timing':>12} {'edges(ex/cap)':>14}" + ) + tot_disp = tot_total = 0 + tot_c_chk = tot_c_pass = 0 + tot_o_chk = tot_o_pass = tot_t_chk = tot_t_pass = 0 + tot_exact = tot_capped = 0 + for r in results: + o_pass = r.timing.order_checks - r.timing.order_failures + t_pass = r.timing.timing_checks - r.timing.timing_failures + if r.dispatched_records == 0: + status = ( + "SKIP (warmup-primed, not profiled)" + if r.warmup_records + else "SKIP (not dispatched)" + ) + elif r.passed: + status = "OK" + else: + status = "FAIL" + print( + f"{r.trace_id:<28} " + f"{r.dispatched_nodes:>4}/{r.total_nodes:<4} " + f"{r.coverage * 100:>4.0f}% " + f"{r.content.passes:>6}/{r.content.checked:<6} " + f"{o_pass:>5}/{r.timing.order_checks:<5} " + f"{t_pass:>5}/{r.timing.timing_checks:<5} " + f"{r.timing.exact_edges:>6}/{r.timing.idle_capped_edges:<6} " + f"{status}" + ) + for m in (*r.content.mismatches, *r.timing.mismatches): + print(f" MISMATCH @ {m.where}: {m.detail}") + for n in r.timing.notes: + if n.startswith("FALLBACK"): + print(f" {n}") + tot_disp += r.dispatched_nodes + tot_total += r.total_nodes + tot_c_chk += r.content.checked + tot_c_pass += r.content.passes + tot_o_chk += r.timing.order_checks + tot_o_pass += o_pass + tot_t_chk += r.timing.timing_checks + tot_t_pass += t_pass + tot_exact += r.timing.exact_edges + tot_capped += r.timing.idle_capped_edges + cov = (tot_disp / tot_total * 100) if tot_total else 0.0 + print( + f"{'TOTAL':<28} {tot_disp:>4}/{tot_total:<4} {cov:>4.0f}% " + f"{tot_c_pass:>6}/{tot_c_chk:<6} {tot_o_pass:>5}/{tot_o_chk:<5} " + f"{tot_t_pass:>5}/{tot_t_chk:<5} {tot_exact:>6}/{tot_capped:<6}" + ) + print( + f"edges: {tot_exact} exact (reproduce real recorded think-time), " + f"{tot_capped} idle-capped (bounded by the {cap} cap)" + ) + + +def _main(argv: list[str] | None = None) -> int: + """CLI: ``--raw --trace-dir [--idle-gap-cap-seconds 60] + [--tokenizer builtin] [--corpus coding] [--seed N]``.""" + import argparse + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--raw", type=Path, required=True, help="raw export JSONL") + parser.add_argument( + "--trace-dir", type=Path, required=True, help="dir of *.json recorded traces" + ) + parser.add_argument( + "--idle-gap-cap-seconds", + type=float, + default=60.0, + help="agentx idle-gap cap applied to the rebuilt trie (default 60; " + "pass a negative value to disable the warp)", + ) + parser.add_argument( + "--tokenizer", + default="builtin", + help="content tokenizer the proved run was built with " + "(default 'builtin', the bare live-run default)", + ) + parser.add_argument( + "--corpus", + default="coding", + help="prompt corpus the proved run was built with " + "(default 'coding', the live-run default)", + ) + parser.add_argument( + "--seed", + type=int, + default=None, + help="content root seed the proved run was built with " + "(the run's --random-seed; default none, the live-run default)", + ) + args = parser.parse_args(argv) + cap = args.idle_gap_cap_seconds if args.idle_gap_cap_seconds >= 0 else None + return prove_corpus( + args.raw, + args.trace_dir, + idle_gap_cap_seconds=cap, + tokenizer_name=args.tokenizer, + prompt_corpus=args.corpus, + root_seed=args.seed, + ) + + +__all__ = [ + "Mismatch", + "Report", + "build_recorded_trace", + "causality_timing_vs_real_trace", + "content_byte_exact_vs_v04", + "content_vs_real_trace", + "load_raw_export", + "prove_corpus", +] + + +if __name__ == "__main__": + raise SystemExit(_main()) diff --git a/tools/weka_trie_timing_sim.py b/tools/weka_trie_timing_sim.py new file mode 100755 index 0000000000..95516df470 --- /dev/null +++ b/tools/weka_trie_timing_sim.py @@ -0,0 +1,304 @@ +#!/usr/bin/env python3 + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Discrete-event timing simulator for the weka segment-trie IR. + +Replays a trie ``ParsedGraph`` with each request's RECORDED ``api_time`` as its +processing duration and checks that the reconstructed per-node start times land +byte-exact on the ORIGINAL recorded timeline (idle-gap-warped). This is a pure +simulation -- no inference server, no ``aiperf`` run -- so it validates the +trie's causal + timing MODEL directly against the recorded ground truth, rather +than (as ``weka_trace_fidelity.py`` does) against a run that itself replays the +same builder. + +The check is mathematically tight. The builder fires each node at +``max over incoming edges of (predecessor_completion + edge_delay)`` (START-roots +at ``min_start_delay``; a START-ANCHORED edge -- +``delay_after_predecessor_start_us``, a mid-flight spawn / overlap -- fires at +``predecessor_DISPATCH + start_delay`` instead of its completion). If, for every +node, the BINDING cause is the +latest-completing one and its ``delay_after_predecessor_us`` is the warped +end-to-start gap ``warped_start(node) - warped_end(binding)``, then feeding the +recorded ``api_time`` back in reconstructs ``warped_start(node)`` EXACTLY: + + sim_end(binding) = warped_start(binding) + api(binding) = warped_end(binding) + sim_start(node) = warped_end(binding) + (warped_start(node) - warped_end(binding)) + = warped_start(node) # exact + +and every wait-only edge (delay 0) contributes ``warped_end(pred) <= +warped_end(binding) <= warped_start(node)``, so the ``max`` is the binding's +term. Any divergence therefore flags a real model error: a wrong binding, a bad +delay, a mis-attributed AND-join, or a warp that did not preserve temporal shape. + +The recorded warped timeline used as the reference is computed by an INDEPENDENT +idle warp here (over the union of recorded active intervals), so the simulator +does not borrow the builder's own warp for its ground truth. + +Usage: + python tools/weka_trie_timing_sim.py [--cap 60] [--tol 1e-3] +""" + +from __future__ import annotations + +import argparse +import heapq +import sys +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +import orjson + +_DEFAULT_CAP = 60.0 +_DEFAULT_TOL = 1e-3 # seconds; float round-trip slack on us<->s conversions + + +class _IdleWarp: + """Independent idle warp over the UNION of active intervals (rule 1). + + A true IDLE gap is dead air -- ``next_start`` minus the running max end + (nothing active in between). Idle > cap collapses to cap; everything after + shifts left. Active stretches are never cut, so temporal shape is preserved. + Deliberately re-implemented here (not imported) so the reference timeline is + independent of the builder. + """ + + def __init__(self, intervals: list[tuple[float, float]], cap: float) -> None: + self._cuts: list[tuple[float, float]] = [] + if not intervals: + return + ordered = sorted(intervals) + running_end = ordered[0][1] + cumulative = 0.0 + for start, end in ordered[1:]: + if start > running_end and (start - running_end) > cap: + cumulative += (start - running_end) - cap + self._cuts.append((start, cumulative)) + if end > running_end: + running_end = end + + def map(self, t: float) -> float: + shift = 0.0 + for next_start, cumulative in self._cuts: + if t < next_start: + break + shift = cumulative + return t - shift + + +def _walk_ids(reqs: list, scope: str) -> Iterator[tuple[str, dict]]: + """Yield ``(node_id, leaf_request)`` in recorded order, keyed by the same + ``{scope}:{turn}`` id scheme the trie builder uses. + + ``scope`` is the trajectory scope -- the trace id for the top-level chain, + the recorded ``agent_id`` for each subagent marker (nested markers use their + own ``agent_id``). ``turn`` is the 0-based leaf index WITHIN that scope, so a + subagent marker never consumes a top-level turn index. + """ + turn = 0 + for req in reqs: + if isinstance(req, dict) and "requests" in req: + yield from _walk_ids(req["requests"], req["agent_id"]) + else: + yield f"{scope}:{turn}", req + turn += 1 + + +def _flatten_api(trace_dict: dict) -> dict[str, float]: + """Recorded ``api_time`` per node id, keyed by the same ``{scope}:{turn}`` id + scheme the trie builder uses.""" + return { + nid: float(req.get("api_time") or 0.0) + for nid, req in _walk_ids(trace_dict["requests"], trace_dict["id"]) + } + + +def _raw_starts(trace_dict: dict) -> dict[str, float]: + return { + nid: float(req["t"]) + for nid, req in _walk_ids(trace_dict["requests"], trace_dict["id"]) + } + + +def _dependency_order( + trace_file: Path, + llm: dict[str, Any], + incoming: dict[str, list[Any]], + recorded_start: dict[str, float], +) -> list[str]: + """Topological order over the static edges among LLM nodes. + + Guarantees every predecessor (end- OR start-anchored) is simulated before + its dependents. A plain ``(recorded_start, node_id)`` sort breaks when a + start-anchored child ties its parent's recorded start and its node-id string + sorts first (e.g. ``agent:0`` < ``trace:0``), which would leave the parent + unsimulated at the child's gate. Ties inside the ready set break by recorded + start then node id so output stays deterministic. + """ + indegree = {nid: 0 for nid in llm} + dependents: dict[str, list[str]] = {nid: [] for nid in llm} + for nid in llm: + for e in incoming[nid]: + if e.source in llm: + indegree[nid] += 1 + dependents[e.source].append(nid) + ready = [ + (recorded_start.get(nid, 0.0), nid) for nid, deg in indegree.items() if deg == 0 + ] + heapq.heapify(ready) + order: list[str] = [] + while ready: + _, nid = heapq.heappop(ready) + order.append(nid) + for succ in dependents[nid]: + indegree[succ] -= 1 + if indegree[succ] == 0: + heapq.heappush(ready, (recorded_start.get(succ, 0.0), succ)) + if len(order) != len(llm): + unresolved = sorted(set(llm) - set(order)) + raise RuntimeError( + f"{trace_file}: dependency cycle among trie LLM nodes; cannot " + f"simulate (unresolved: {unresolved[:5]})" + ) + return order + + +def simulate_trace( + trace_file: Path, cap: float, tol: float +) -> tuple[int, int, list[str], int]: + """Return (n_checked, n_exact, divergence_lines, n_first_token_edges) for one trace. + + ``n_first_token_edges`` counts start-anchored edges that ALSO carry + ``delay_after_predecessor_first_token_us`` (a POST-TTFT overlap): a distinct + kind label. Their gate is unchanged -- a pure replay observes ttft == recorded + ttft, so ``first_token + D' == dispatch + D`` and the start-anchored branch's + ``sim_start(source) + start_delay`` already reconstructs them exactly. + """ + from aiperf.dataset.graph.adapters.weka.trace_models import WekaTrace + from aiperf.dataset.graph.adapters.weka.trie_build import build_trie_graph + from aiperf.dataset.graph.models import LlmNode, StaticEdge + + raw = orjson.loads(Path(trace_file).read_bytes()) + trace = WekaTrace.model_validate(raw) + graph = build_trie_graph( + trace, + tokenizer_name="builtin", + prompt_corpus="coding", + root_seed=None, + idle_gap_cap_seconds=cap, + )[0].graph + + api = _flatten_api(raw) + starts = _raw_starts(raw) + # Independent recorded warped timeline (ground truth) from raw starts + api. + warp = _IdleWarp([(starts[nid], starts[nid] + api[nid]) for nid in starts], cap) + recorded_start = {nid: warp.map(t) for nid, t in starts.items()} + + # Incoming static edges per node. + incoming: dict[str, list[StaticEdge]] = {nid: [] for nid in graph.nodes} + for e in graph.edges: + if isinstance(e, StaticEdge) and e.target in incoming: + incoming[e.target].append(e) + + llm = {nid: n for nid, n in graph.nodes.items() if isinstance(n, LlmNode)} + order = _dependency_order(trace_file, llm, incoming, recorded_start) + + sim_start: dict[str, float] = {} + sim_end: dict[str, float] = {} + checked = exact = first_token_edges = 0 + diverged: list[str] = [] + for nid in order: + gate = 0.0 + for e in incoming[nid]: + if e.source == "START": + gate = max(gate, (e.min_start_delay_us or 0.0) / 1e6) + continue + if e.source not in sim_start: + # Dependency order guarantees every LLM predecessor is simulated + # first; an unknown source here is a real model error (a dangling + # or non-LLM endpoint) that zero-defaulting would silently absorb. + raise RuntimeError( + f"{trace_file}: edge {e.source} -> {nid} references a " + f"predecessor that was never simulated (unknown or non-LLM " + f"source); refusing to zero-default its gate" + ) + if e.delay_after_predecessor_start_us is not None: + # Start-anchored: gate off the predecessor's DISPATCH (its own + # sim_start), NOT its completion -- a mid-flight spawn / overlap + # fires ``delay`` after the parent dispatched, not after it ended. + # A POST-TTFT edge (first-token delay also present) shares this + # gate: in a pure replay first_token + D' == dispatch + D. + if e.delay_after_predecessor_first_token_us is not None: + first_token_edges += 1 + d = e.delay_after_predecessor_start_us / 1e6 + gate = max(gate, sim_start[e.source] + d) + else: + d = (e.delay_after_predecessor_us or 0.0) / 1e6 + gate = max(gate, sim_end[e.source] + d) + sim_start[nid] = gate + sim_end[nid] = gate + api.get(nid, 0.0) + + ref = recorded_start.get(nid) + if ref is None: + continue + checked += 1 + err = abs(gate - ref) + if err <= tol: + exact += 1 + else: + diverged.append( + f" {nid}: sim_start={gate:.3f}s != recorded_warped={ref:.3f}s " + f"(|d|={err:.3f}s)" + ) + return checked, exact, diverged, first_token_edges + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("path", type=Path, help="trace dir or single .json trace file") + ap.add_argument("--cap", type=float, default=_DEFAULT_CAP) + ap.add_argument("--tol", type=float, default=_DEFAULT_TOL) + ap.add_argument( + "--show", type=int, default=5, help="max divergences to print per trace" + ) + args = ap.parse_args(argv) + + files = sorted(args.path.glob("*.json")) if args.path.is_dir() else [args.path] + print(f"weka trie timing simulator (idle cap = {args.cap}s, tol = {args.tol}s)") + if not files: + print("VACUOUS: nothing checked (no *.json trace files found)") + return 1 + print(f"{'trace':36s} {'checked':>8s} {'exact':>8s} {'ft_edges':>9s} status") + grand_checked = grand_exact = grand_first_token = 0 + failed = False + for tf in files: + checked, exact, diverged, first_token_edges = simulate_trace( + tf, args.cap, args.tol + ) + ok = exact == checked + failed = failed or not ok + print( + f"{tf.stem[:36]:36s} {checked:8d} {exact:8d} {first_token_edges:9d} " + f"{'OK' if ok else 'DIVERGE'}" + ) + for line in diverged[: args.show]: + print(line) + grand_checked += checked + grand_exact += exact + grand_first_token += first_token_edges + print( + f"\nTOTAL: {grand_exact}/{grand_checked} nodes reconstruct the recorded warped timeline" + ) + print( + f"first-token edges (post-TTFT overlap, gated at dispatch + D): " + f"{grand_first_token}" + ) + if grand_checked == 0: + print("VACUOUS: nothing checked (zero LLM nodes simulated)") + return 1 + return 1 if failed else 0 + + +if __name__ == "__main__": + sys.exit(main())